import sys
import os
import re
import cv2
import numpy as np
import pandas as pd
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, 
                             QLabel, QLineEdit, QPushButton, QFileDialog, QTabWidget, 
                             QSlider, QDoubleSpinBox, QProgressBar, QComboBox, QGroupBox, QFormLayout,
                             QSpinBox)
from PyQt5.QtCore import Qt, QThread, pyqtSignal, QTimer
from PyQt5.QtGui import QImage, QPixmap
from scipy.spatial.transform import Rotation as R
from pyproj import Transformer

# =================================================================
# [NEXT GEN] advanced_tuner_v2.py: Streamlined GIS GUI Tuner
# 설계 원칙:
# 1. MP4 + SRT + CSV (WGS84) 워크플로우 통합
# 2. 실시간 좌표 변환 (Lat/Lon -> EPSG:5186)
# 3. 사용자 친화적 튜닝 인터페이스 제공
# =================================================================

class SRTParser:
    @staticmethod
    def parse(srt_path):
        data_dict = {}
        if not os.path.exists(srt_path): return data_dict
        try:
            with open(srt_path, 'r', encoding='utf-8', errors='ignore') as f:
                content = f.read()
            if "latitude" not in content.lower():
                with open(srt_path, 'r', encoding='cp949', errors='ignore') as f:
                    content = f.read()
            
            blocks = content.split('\n\n')
            for block in blocks:
                f_match = re.search(r'FrameCnt:\s*(\d+)', block)
                if not f_match: continue
                idx = int(f_match.group(1))
                
                lat = float(re.search(r'latitude:\s*([\d\.-]+)', block).group(1))
                lon = float(re.search(r'longitude:\s*([\d\.-]+)', block).group(1))
                alt = float(re.search(r'abs_alt:\s*([\d\.-]+)', block).group(1))
                yaw = float(re.search(r'gb_yaw:\s*([\d\.-]+)', block).group(1))
                pitch = float(re.search(r'gb_pitch:\s*([\d\.-]+)', block).group(1))
                roll = float(re.search(r'gb_roll:\s*([\d\.-]+)', block).group(1))
                focal = float(re.search(r'focal_len:\s*([\d\.-]+)', block).group(1))
                
                data_dict[idx] = {
                    'lat': lat, 'lon': lon, 'alt': alt,
                    'yaw': yaw, 'pitch': pitch, 'roll': roll,
                    'focal': focal
                }
            # Post-process: Filling gaps and Smoothing
            if not data_dict: return {}
            
            indices = sorted(data_dict.keys())
            min_idx, max_idx = indices[0], indices[-1]
            
            # 1. Interpolate Gaps
            all_indices = np.arange(min_idx, max_idx + 1)
            new_data = {}
            
            fields = ['lat', 'lon', 'alt', 'yaw', 'pitch', 'roll', 'focal']
            temp_arrays = {f: [] for f in fields}
            for idx in indices:
                for f in fields:
                    temp_arrays[f].append(data_dict[idx][f])
            
            interp_arrays = {}
            for f in fields:
                # Handle Yaw wrapping for smoother rotation (Optional but safer)
                if f == 'yaw':
                    y_arr = np.array(temp_arrays[f])
                    # Ensure continuous angles
                    y_arr = np.unwrap(np.radians(y_arr))
                    interp_y = np.interp(all_indices, indices, y_arr)
                    # 2. Smooth (Moving Average)
                    smooth_y = np.convolve(interp_y, np.ones(5)/5, mode='same')
                    interp_arrays[f] = np.degrees(smooth_y)
                else:
                    interp_f = np.interp(all_indices, indices, temp_arrays[f])
                    # 2. Smooth (Moving Average)
                    interp_arrays[f] = np.convolve(interp_f, np.ones(5)/5, mode='same')

            for i, idx in enumerate(all_indices):
                new_data[int(idx)] = {f: interp_arrays[f][i] for f in fields}
                
            return new_data
        except Exception as e:
            print(f"SRT Error: {e}")
            return {}

class RenderWorker(QThread):
    progress = pyqtSignal(int)
    finished = pyqtSignal(str)

    def __init__(self, params):
        super().__init__()
        self.p = params
        self.is_running = True

    def stop(self):
        self.is_running = False

    def run(self):
        try:
            cap = cv2.VideoCapture(self.p['video_path'])
            original_fps = cap.get(cv2.CAP_PROP_FPS)
            cap.set(cv2.CAP_PROP_POS_FRAMES, self.p['range'][0])
            
            target_w, target_h = self.p['resolution']
            target_fps = self.p['fps'] if self.p['fps'] > 0 else original_fps
            start_f, end_f = self.p['range']
            total_work = end_f - start_f + 1

            fourcc = cv2.VideoWriter_fourcc(*self.p['codec'])
            out = cv2.VideoWriter(self.p['output_path'], fourcc, target_fps, (target_w, target_h))

            transformer = Transformer.from_crs("epsg:4326", "epsg:5186")
            
            # Pre-swap line points if needed
            eff_line_pts = self.p['line_pts']
            if self.p['swap_xy']:
                eff_line_pts = eff_line_pts.copy()
                eff_line_pts[:, [0, 1]] = eff_line_pts[:, [1, 0]]

            for f_idx in range(start_f, end_f + 1):
                if not self.is_running: break
                ret, frame = cap.read()
                if not ret: break
                
                # Check for Sync
                if f_idx in self.p['srt_data']:
                    meta = self.p['srt_data'][f_idx]
                    dx, dy = transformer.transform(meta['lat'], meta['lon'])
                    
                    if self.p['swap_xy']:
                        dx, dy = dy, dx

                    # Resize frame if needed
                    original_h, original_w = frame.shape[:2]
                    if (target_w, target_h) != (original_w, original_h):
                        draw_frame = cv2.resize(frame, (target_w, target_h))
                    else:
                        draw_frame = frame.copy()

                    drone_pos = np.array([dx + self.p['off_x'], dy + self.p['off_y'], meta['alt'] + self.p['off_z']])
                    rvec, tvec, K = self.get_proj(meta, drone_pos, target_w, target_h)
                    R_w2c, _ = cv2.Rodrigues(rvec)
                    pts_cam = (eff_line_pts @ R_w2c.T) + tvec.T
                    
                    for i in range(0, len(pts_cam), 2):
                        p1c, p2c = pts_cam[i], pts_cam[i+1]
                        if p1c[2] < 0.1 and p2c[2] < 0.1: continue
                        if p1c[2] < 0.1 or p2c[2] < 0.1:
                            t = (0.1 - p1c[2]) / (p2c[2] - p1c[2])
                            p_mid = p1c + t * (p2c - p1c)
                            if p1c[2] < 0.1: p1c = p_mid
                            else: p2c = p_mid
                        
                        u1 = int(K[0, 0] * (p1c[0] / p1c[2]) + K[0, 2])
                        v1 = int(K[1, 1] * (p1c[1] / p1c[2]) + K[1, 2])
                        u2 = int(K[0, 0] * (p2c[0] / p2c[2]) + K[0, 2])
                        v2 = int(K[1, 1] * (p2c[1] / p2c[2]) + K[1, 2])
                        
                        ret_cli, pt1, pt2 = cv2.clipLine((0, 0, target_w, target_h), (u1, v1), (u2, v2))
                        if ret_cli:
                            cv2.line(draw_frame, pt1, pt2, (0, 0, 255), 10)
                    out.write(draw_frame)
                else:
                    if (target_w, target_h) != (frame.shape[1], frame.shape[0]):
                        out.write(cv2.resize(frame, (target_w, target_h)))
                    else:
                        out.write(frame)

                self.progress.emit(int(((f_idx - start_f + 1) / total_work) * 100))


            cap.release()
            out.release()
            self.finished.emit(self.p['output_path'])
        except Exception as e:
            self.finished.emit(f"Error: {e}")

    def get_proj(self, meta, drone_pos, w, h):
        yaw = np.radians(meta['yaw'] + self.p['off_yaw'])
        pitch = np.radians(meta['pitch'] + self.p['off_pitch'])
        roll = np.radians(meta['roll'] + self.p['off_roll'])
        
        R_b2w = (R.from_euler('z', -yaw) * R.from_euler('x', pitch) * R.from_euler('y', roll)).as_matrix()
        R_w2c = np.array([[1, 0, 0], [0, 0, -1], [0, 1, 0]]) @ R_b2w.T
        rvec, _ = cv2.Rodrigues(R_w2c)
        tvec = -R_w2c @ drone_pos
        
        f_px = (self.p['focal'] / self.p['sensor_w']) * w
        K = np.array([[f_px, 0, w/2], [0, f_px, h/2], [0, 0, 1]])
        return rvec, tvec, K

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Advanced GIS Tuner v2 (Streamlined)")
        self.resize(1280, 850)
        
        self.srt_data = {}
        self.line_points = np.array([])
        self.transformer = Transformer.from_crs("epsg:4326", "epsg:5186")
        
        self.init_ui()

    def init_ui(self):
        tabs = QTabWidget()
        self.setCentralWidget(tabs)
        
        # --- Tab 1: Data Input ---
        tab1 = QWidget()
        layout1 = QVBoxLayout()
        form1 = QFormLayout()
        
        self.btn_video = QPushButton("Find MP4")
        self.txt_video = QLineEdit("하행)회덕-대전조차장.MP4")
        self.btn_video.clicked.connect(lambda: self.find_file(self.txt_video, "Video (*.mp4)"))
        form1.addRow(self.btn_video, self.txt_video)
        
        self.btn_srt = QPushButton("Find SRT")
        self.txt_srt = QLineEdit("하행)회덕-대전조차장.srt")
        self.btn_srt.clicked.connect(lambda: self.find_file(self.txt_srt, "SRT (*.srt)"))
        form1.addRow(self.btn_srt, self.txt_srt)
        
        self.btn_csv = QPushButton("Find CSV (Lat/Lon/H)")
        self.txt_csv = QLineEdit("center.csv")
        self.btn_csv.clicked.connect(lambda: self.find_file(self.txt_csv, "CSV (*.csv)"))
        form1.addRow(self.btn_csv, self.txt_csv)
        
        layout1.addLayout(form1)
        self.btn_load = QPushButton("LOAD ALL DATA")
        self.btn_load.setStyleSheet("background-color: #4CAF50; color: white; height: 50px; font-weight: bold;")
        self.btn_load.clicked.connect(self.load_all_data)
        layout1.addWidget(self.btn_load)
        layout1.addStretch()
        tab1.setLayout(layout1)
        tabs.addTab(tab1, "1. 자료입력")

        # --- Tab 2: Tuner ---
        tab2 = QWidget()
        hbox2 = QHBoxLayout()
        
        # Left: Preview
        vbox_prev = QVBoxLayout()
        self.lbl_preview = QLabel("Video Preview (Load Data First)")
        self.lbl_preview.setAlignment(Qt.AlignCenter)
        self.lbl_preview.setStyleSheet("border: 2px solid gray; background-color: black;")
        self.lbl_preview.setMinimumSize(800, 450)
        vbox_prev.addWidget(self.lbl_preview)
        
        self.sld_frame = QSlider(Qt.Horizontal)
        self.sld_frame.valueChanged.connect(self.update_sync_frame)
        self.spn_frame = QDoubleSpinBox()
        self.spn_frame.setDecimals(0)
        self.spn_frame.setSingleStep(1)
        self.spn_frame.valueChanged.connect(self.update_sync_frame)
        
        self.btn_play = QPushButton("▶ PLAY")
        self.btn_play.setCheckable(True)
        self.btn_play.clicked.connect(self.toggle_play)
        
        self.play_timer = QTimer()
        self.play_timer.timeout.connect(self.next_frame)
        
        frame_hbox = QHBoxLayout()
        frame_hbox.addWidget(self.btn_play)
        frame_hbox.addWidget(QLabel("Frame:"))
        frame_hbox.addWidget(self.spn_frame)
        frame_hbox.addWidget(self.sld_frame, 10)
        vbox_prev.addLayout(frame_hbox)
        
        hbox2.addLayout(vbox_prev, 7)
        
        # Right: Controls
        vbox_ctrl = QVBoxLayout()
        group_offsets = QGroupBox("Offsets (Yaw/Pitch/Roll/XYZ)")
        form2 = QFormLayout()
        
        self.chk_swap_xy = QPushButton("Swap XY Coordinates: OFF")
        self.chk_swap_xy.setCheckable(True)
        self.chk_swap_xy.clicked.connect(self.toggle_swap_xy)
        form2.addRow("Axis:", self.chk_swap_xy)
        
        self.spn_yaw = self.create_spinbox(-180, 180, 0, form2, "Yaw Off")
        self.spn_pitch = self.create_spinbox(-180, 180, 0, form2, "Pitch Off")
        self.spn_roll = self.create_spinbox(-180, 180, 0, form2, "Roll Off")
        form2.addRow(QLabel("<hr>"))
        self.spn_off_x = self.create_spinbox(-1000, 1000, 0, form2, "Pos X Off (m)")
        self.spn_off_y = self.create_spinbox(-1000, 1000, 0, form2, "Pos Y Off (m)")
        self.spn_off_z = self.create_spinbox(-1000, 1000, 0, form2, "Pos Z Off (m)")
        form2.addRow(QLabel("<hr>"))
        self.spn_focal = self.create_spinbox(1, 200, 24, form2, "Focal (mm)")
        self.spn_sensor = self.create_spinbox(1, 100, 36, form2, "Sensor W (mm)")
        
        group_offsets.setLayout(form2)
        vbox_ctrl.addWidget(group_offsets)
        vbox_ctrl.addStretch()
        hbox2.addLayout(vbox_ctrl, 3)
        
        tab2.setLayout(hbox2)
        tabs.addTab(tab2, "2. 튜너")

        # --- Tab 3: Export ---
        tab3 = QWidget()
        vbox3 = QVBoxLayout()
        form3 = QFormLayout()
        
        self.cmb_res = QComboBox()
        self.cmb_res.addItems(["Original", "4K (3840x2160)", "2K (2560x1440)", "1080p (1920x1080)"])
        form3.addRow("Resolution", self.cmb_res)
        
        self.cmb_fps = QComboBox()
        self.cmb_fps.addItems(["Original", "60", "30"])
        form3.addRow("FPS", self.cmb_fps)
        
        self.cmb_codec = QComboBox()
        self.cmb_codec.addItems(["H.264 (avc1)", "MPEG-4 (mp4v)"])
        form3.addRow("Codec", self.cmb_codec)
        
        form3.addRow(QLabel("<hr>"))
        self.spn_start = QSpinBox()
        self.spn_start.setRange(0, 999999)
        form3.addRow("Start Frame", self.spn_start)
        
        self.spn_end = QSpinBox()
        self.spn_end.setRange(0, 999999)
        form3.addRow("End Frame", self.spn_end)
        
        vbox3.addLayout(form3)
        self.btn_export = QPushButton("START EXPORT")
        self.btn_export.setStyleSheet("height: 60px; font-weight: bold; background-color: #f44336; color: white;")
        self.btn_export.clicked.connect(self.start_export)
        vbox3.addWidget(self.btn_export)
        
        self.pbar = QProgressBar()
        vbox3.addWidget(self.pbar)
        vbox3.addStretch()
        tab3.setLayout(vbox3)
        tabs.addTab(tab3, "3. 출력")

    def create_spinbox(self, min_v, max_v, def_v, layout, label):
        sb = QDoubleSpinBox()
        sb.setRange(min_v, max_v)
        sb.setValue(def_v)
        sb.setDecimals(2)
        sb.setSingleStep(0.1)
        sb.valueChanged.connect(self.update_preview)
        layout.addRow(label, sb)
        return sb

    def toggle_play(self):
        if self.btn_play.isChecked():
            self.btn_play.setText("■ PAUSE")
            # Calculate interval based on FPS (or default 33ms for 30fps)
            cap = cv2.VideoCapture(self.txt_video.text())
            fps = cap.get(cv2.CAP_PROP_FPS)
            cap.release()
            interval = int(1000 / fps) if fps > 0 else 33
            self.play_timer.start(interval)
        else:
            self.btn_play.setText("▶ PLAY")
            self.play_timer.stop()

    def next_frame(self):
        curr = self.sld_frame.value()
        if curr < self.total_frames - 1:
            self.sld_frame.setValue(curr + 1)
        else:
            self.btn_play.setChecked(False)
            self.toggle_play()

    def toggle_swap_xy(self):
        if self.chk_swap_xy.isChecked():
            self.chk_swap_xy.setText("Swap XY Coordinates: ON")
        else:
            self.chk_swap_xy.setText("Swap XY Coordinates: OFF")
        self.update_preview()

    def update_sync_frame(self, val):
        # 도우미: 슬라이더와 스핀박스 동기화
        sender = self.sender()
        if sender == self.sld_frame:
            self.spn_frame.blockSignals(True)
            self.spn_frame.setValue(val)
            self.spn_frame.blockSignals(False)
        else:
            self.sld_frame.blockSignals(True)
            self.sld_frame.setValue(int(val))
            self.sld_frame.blockSignals(False)
        self.update_preview()

    def find_file(self, line_edit, filt):
        path, _ = QFileDialog.getOpenFileName(self, "Select File", "", filt)
        if path: line_edit.setText(path)

    def load_all_data(self):
        print("Loading...")
        self.srt_data = SRTParser.parse(self.txt_srt.text())
        
        # Load CSV
        try:
            df = pd.read_csv(self.txt_csv.text(), encoding='cp949')
            # lat, lon, 타원체고(h)
            raw_pts = []
            for _, row in df.iterrows():
                x, y = self.transformer.transform(row['lat'], row['lon'])
                raw_pts.append([x, y, row['타원체고(h)']])
            
            # Reconstruct segments: 0-1, 1-2, 2-3...
            seg_pts = []
            for i in range(len(raw_pts)-1):
                seg_pts.extend([raw_pts[i], raw_pts[i+1]])
                
            self.line_points = np.array(seg_pts, dtype=np.float64)
            print(f"Loaded {len(self.line_points)//2} segments.")
        except Exception as e:
            print(f"CSV Load Error: {e}")

        # Video Info
        cap = cv2.VideoCapture(self.txt_video.text())
        self.total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        self.video_fps = cap.get(cv2.CAP_PROP_FPS)
        
        self.sld_frame.setRange(0, self.total_frames - 1)
        self.spn_frame.setRange(0, self.total_frames - 1)
        self.spn_start.setRange(0, self.total_frames - 1)
        self.spn_end.setRange(0, self.total_frames - 1)
        self.spn_end.setValue(self.total_frames - 1)
        cap.release()
        
        self.update_preview()

    def update_preview(self):
        if not self.srt_data or self.line_points.size == 0: return
        
        f_idx = self.sld_frame.value()
        cap = cv2.VideoCapture(self.txt_video.text())
        cap.set(cv2.CAP_PROP_POS_FRAMES, f_idx)
        ret, frame = cap.read()
        cap.release()
        
        if not ret: return
        
        h, w = frame.shape[:2]
        if f_idx in self.srt_data:
            meta = self.srt_data[f_idx]
            
            dx, dy = self.transformer.transform(meta['lat'], meta['lon'])
            
            # Swap Logic
            if self.chk_swap_xy.isChecked():
                dx, dy = dy, dx
                line_pts_eff = self.line_points.copy()
                line_pts_eff[:, [0, 1]] = line_pts_eff[:, [1, 0]]
            else:
                line_pts_eff = self.line_points

            drone_pos = np.array([dx + self.spn_off_x.value(), 
                                  dy + self.spn_off_y.value(), 
                                  meta['alt'] + self.spn_off_z.value()])
            
            # 실시간 디버그 로그 (콘솔 출력)
            print(f"[Frame {f_idx}] XY_Order: {'Y,X' if self.chk_swap_xy.isChecked() else 'X,Y'} | Drone X={drone_pos[0]:.2f}, Y={drone_pos[1]:.2f}")
            
            # Projection
            yaw = np.radians(meta['yaw'] + self.spn_yaw.value())
            pitch = np.radians(meta['pitch'] + self.spn_pitch.value())
            roll = np.radians(meta['roll'] + self.spn_roll.value())
            
            r_yaw = R.from_euler('z', -yaw)
            r_pitch = R.from_euler('x', pitch)
            r_roll = R.from_euler('y', roll)
            
            R_b2w = (r_yaw * r_pitch * r_roll).as_matrix()
            R_align = np.array([[1, 0, 0], [0, 0, -1], [0, 1, 0]])
            R_w2c = R_align @ R_b2w.T
            
            rvec, _ = cv2.Rodrigues(R_w2c)
            tvec = -R_w2c @ drone_pos
            
            f_px = (self.spn_focal.value() / self.spn_sensor.value()) * w
            K = np.array([[f_px, 0, w/2], [0, f_px, h/2], [0, 0, 1]])
            
            pts_cam = (line_pts_eff @ R_w2c.T) + tvec.T
            for i in range(0, len(pts_cam), 2):
                p1c, p2c = pts_cam[i], pts_cam[i+1]
                if p1c[2] < 0.1 and p2c[2] < 0.1: continue
                if p1c[2] < 0.1 or p2c[2] < 0.1:
                    t = (0.1 - p1c[2]) / (p2c[2] - p1c[2])
                    p_mid = p1c + t * (p2c - p1c)
                    if p1c[2] < 0.1: p1c = p_mid
                    else: p2c = p_mid
                
                u1 = int(K[0, 0] * (p1c[0] / p1c[2]) + K[0, 2])
                v1 = int(K[1, 1] * (p1c[1] / p1c[2]) + K[1, 2])
                u2 = int(K[0, 0] * (p2c[0] / p2c[2]) + K[0, 2])
                v2 = int(K[1, 1] * (p2c[1] / p2c[2]) + K[1, 2])
                
                ret_cli, pt1, pt2 = cv2.clipLine((0, 0, w, h), (u1, v1), (u2, v2))
                if ret_cli:
                    cv2.line(frame, pt1, pt2, (0, 0, 255), 10)

        # Convert to QPixmap
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        img = QImage(frame.data, w, h, w*3, QImage.Format_RGB888)
        pixmap = QPixmap.fromImage(img).scaled(self.lbl_preview.width(), self.lbl_preview.height(), Qt.KeepAspectRatio)
        self.lbl_preview.setPixmap(pixmap)

    def start_export(self):
        res_map = {"Original": (0, 0), "4K (3840x2160)": (3840, 2160), "2K (2560x1440)": (2560, 1440), "1080p (1920x1080)": (1920, 1080)}
        codec_map = {"H.264 (avc1)": "avc1", "MPEG-4 (mp4v)": "mp4v"}
        
        target_res = res_map[self.cmb_res.currentText()]
        if target_res == (0, 0):
            cap = cv2.VideoCapture(self.txt_video.text())
            target_res = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
            cap.release()
            
        params = {
            'video_path': self.txt_video.text(),
            'srt_data': self.srt_data,
            'line_pts': self.line_points,
            'off_yaw': self.spn_yaw.value(),
            'off_pitch': self.spn_pitch.value(),
            'off_roll': self.spn_roll.value(),
            'off_x': self.spn_off_x.value(),
            'off_y': self.spn_off_y.value(),
            'off_z': self.spn_off_z.value(),
            'focal': self.spn_focal.value(),
            'sensor_w': self.spn_sensor.value(),
            'resolution': target_res,
            'fps': int(self.cmb_fps.currentText()) if self.cmb_fps.currentText() != "Original" else 0,
            'range': (self.spn_start.value(), self.spn_end.value()),
            'swap_xy': self.chk_swap_xy.isChecked(),
            'codec': codec_map[self.cmb_codec.currentText()],
            'output_path': "output_tuner_v2.mp4"
        }
        
        self.btn_export.setEnabled(False)
        self.worker = RenderWorker(params)
        self.worker.progress.connect(self.pbar.setValue)
        self.worker.finished.connect(self.export_done)
        self.worker.start()

    def export_done(self, msg):
        self.btn_export.setEnabled(True)
        print(f"Export Result: {msg}")

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())