#!/usr/bin/env python3
"""
任務規劃視覺化驗證工具（含動畫模擬）

使用方式:
    1. 由 GUI 自動觸發:  python3 verify_waypoints.py --file /tmp/mission_plan.json
    2. 獨立手動執行:      python3 verify_waypoints.py

操作:
    - 啟動後先顯示靜態航點圖
    - 點擊圖下方的按鈕控制動畫
"""
import sys, os
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..'))

import json
import argparse
import math
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.widgets import Button
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
from mission_planner import FormationPlanner, MissionType, CoordinateConverter


# ================================================================================
# 色彩定義
# ================================================================================
COLORS = ['#378ADD', '#1D9E75', '#BA7517', '#D85A30', '#7F77DD', '#D4537E',
          '#E24B4A', '#639922', '#00BFFF', '#FF69B4']

# 動畫參數
FRAMES_PER_SEGMENT = 40    # 每段航點之間的畫面數
TRAIL_LENGTH = 60          # 拖尾長度（畫面數）
INTERVAL_MS = 50           # 每幀間隔 (ms)


# ================================================================================
# 靜態繪圖
# ================================================================================

def plot_grid_sweep(ax, data, conv):
    """畫 Grid Sweep 俯視圖"""
    if 'rect_corners' in data:
        rect_local = [conv.gps_to_local(c[0], c[1], 0)[:2] for c in data['rect_corners']]
        xs = [p[0] for p in rect_local] + [rect_local[0][0]]
        ys = [p[1] for p in rect_local] + [rect_local[0][1]]
        ax.plot(xs, ys, 'k--', linewidth=1.5, label='Task area')
        ax.fill(xs, ys, alpha=0.05, color='gray')

    _draw_waypoint_paths(ax, data, conv, show_sweep_labels=True)

    total_wps = sum(len(wps) for wps in data['waypoints'])
    ax.set_title(f'Grid Sweep - {len(data["drone_ids"])} drones, {total_wps} waypoints')
    ax.set_xlabel('East (m)')
    ax.set_ylabel('North (m)')
    ax.set_aspect('equal')
    ax.grid(True, alpha=0.3)


def plot_formation(ax, data, conv):
    """畫 M-Formation / Circle / Leader-Follower 俯視圖"""
    _draw_waypoint_paths(ax, data, conv, show_sweep_labels=False)

    if 'target' in data:
        t = data['target']
        tx, ty, _ = conv.gps_to_local(t[0], t[1], t[2] if len(t) > 2 else 0)
        ax.plot(tx, ty, 'r*', markersize=18, zorder=5)
        ax.annotate('Target', (tx, ty + 1), fontsize=9, color='red', ha='center', va='bottom')

    if 'route_waypoints' in data:
        route_local = [conv.gps_to_local(wp[0], wp[1], 0)[:2] for wp in data['route_waypoints']]
        rxs = [p[0] for p in route_local]
        rys = [p[1] for p in route_local]
        ax.plot(rxs, rys, 'r--', linewidth=1.5, alpha=0.5, label='Route center')
        for i, (rx, ry) in enumerate(route_local):
            ax.plot(rx, ry, 'ro', markersize=6, alpha=0.5)
            ax.annotate(f'R{i+1}', (rx, ry + 0.8), fontsize=7, color='red',
                       ha='center', va='bottom', alpha=0.7)

    mission_type = data.get('mission_type', 'formation')
    total_wps = sum(len(wps) for wps in data['waypoints'])
    ax.set_title(f'{mission_type} - {len(data["drone_ids"])} drones, {total_wps} waypoints')
    ax.set_xlabel('East (m)')
    ax.set_ylabel('North (m)')
    ax.set_aspect('equal')
    ax.grid(True, alpha=0.3)


def _draw_waypoint_paths(ax, data, conv, show_sweep_labels=False):
    """共用的航點路徑繪圖"""
    drone_ids = data['drone_ids']
    waypoints = data['waypoints']
    drones_gps = data.get('drones_gps', [])

    for i, pos in enumerate(drones_gps):
        x, y, _ = conv.gps_to_local(pos[0], pos[1], pos[2] if len(pos) > 2 else 0)
        c = COLORS[i % len(COLORS)]
        ax.plot(x, y, 'o', color=c, markersize=10, zorder=5)
        ax.annotate(f'{drone_ids[i]}', (x, y + 1), fontsize=8, fontweight='bold',
                    ha='center', va='bottom', color=c)

    for i, wps in enumerate(waypoints):
        c = COLORS[i % len(COLORS)]
        local_wps = [conv.gps_to_local(wp[0], wp[1], wp[2]) for wp in wps]
        xs = [p[0] for p in local_wps]
        ys = [p[1] for p in local_wps]
        ax.plot(xs, ys, '-', color=c, linewidth=1.5, alpha=0.7)

        for j, (x, y, z) in enumerate(local_wps):
            if show_sweep_labels:
                if j == 0:
                    ax.plot(x, y, 's', color=c, markersize=8, zorder=4)
                    ax.annotate('gather', (x, y), fontsize=6, ha='right', va='top')
                elif j == 1:
                    ax.plot(x, y, '^', color=c, markersize=8, zorder=4)
                    ax.annotate('entry', (x, y), fontsize=6, ha='right', va='top')
                else:
                    ax.plot(x, y, '.', color=c, markersize=4)
            else:
                marker = 's' if j == 0 else '*'
                ax.plot(x, y, marker, color=c, markersize=10, zorder=4)
                ax.annotate(f'WP{j}\n({z:.0f}m)', (x, y), fontsize=6, ha='center', va='bottom')

        if local_wps:
            lx, ly, _ = local_wps[-1]
            ax.plot(lx, ly, 'X', color=c, markersize=10, markeredgewidth=2, zorder=4)


# ================================================================================
# 動畫模擬
# ================================================================================

class MissionAnimator:
    """任務動畫控制器"""

    def __init__(self, fig, ax, data, conv):
        self.fig = fig
        self.ax = ax
        self.data = data
        self.conv = conv
        self.is_playing = False
        self.anim = None
        self.current_frame = 0

        drone_ids = data['drone_ids']
        waypoints = data['waypoints']
        drones_gps = data.get('drones_gps', [])
        self.num_drones = len(drone_ids)

        # 建立完整航點序列：初始位置 → WP0 → WP1 → ...
        self.all_local_wps = []
        for i, wps in enumerate(waypoints):
            local_wps = [conv.gps_to_local(wp[0], wp[1], wp[2]) for wp in wps]

            # 把初始位置插入最前面
            if i < len(drones_gps):
                pos = drones_gps[i]
                start = conv.gps_to_local(pos[0], pos[1], pos[2] if len(pos) > 2 else 0)
                local_wps.insert(0, start)

            self.all_local_wps.append(local_wps)

        # 計算最大航點段數
        self.max_segments = max(len(wps) - 1 for wps in self.all_local_wps) if self.all_local_wps else 0
        self.total_frames = self.max_segments * FRAMES_PER_SEGMENT

        # 預計算位置
        self.positions = self._precompute_positions()

        # 動畫元素
        self.drone_dots = []
        self.trail_lines = []
        self.trail_data = [[] for _ in range(self.num_drones)]
        self.status_text = None

    def _precompute_positions(self):
        """預計算所有幀的位置 — 等時間步進"""
        positions = []

        for frame in range(self.total_frames + 1):
            seg_idx = frame // FRAMES_PER_SEGMENT
            seg_progress = (frame % FRAMES_PER_SEGMENT) / FRAMES_PER_SEGMENT

            frame_positions = []
            for drone_idx in range(self.num_drones):
                wps = self.all_local_wps[drone_idx]
                num_segs = len(wps) - 1

                if seg_idx >= num_segs:
                    frame_positions.append((wps[-1][0], wps[-1][1]))
                else:
                    x0, y0, _ = wps[seg_idx]
                    x1, y1, _ = wps[seg_idx + 1]
                    x = x0 + (x1 - x0) * seg_progress
                    y = y0 + (y1 - y0) * seg_progress
                    frame_positions.append((x, y))

            positions.append(frame_positions)

        return positions

    def setup(self):
        """建立動畫元素和按鈕"""
        for i in range(self.num_drones):
            c = COLORS[i % len(COLORS)]
            dot, = self.ax.plot([], [], 'o', color=c, markersize=12,
                               markeredgecolor='white', markeredgewidth=1.5, zorder=10)
            self.drone_dots.append(dot)
            trail, = self.ax.plot([], [], '-', color=c, linewidth=2.5, alpha=0.4, zorder=9)
            self.trail_lines.append(trail)

        self.status_text = self.ax.text(
            0.02, 0.98, 'Ready',
            transform=self.ax.transAxes, fontsize=10,
            verticalalignment='top',
            bbox=dict(boxstyle='round,pad=0.3', facecolor='white', alpha=0.8)
        )

        # 按鈕放在右上角圖內，避免擋到軸標籤
        ax_play = self.fig.add_axes([0.72, 0.91, 0.08, 0.035])
        self.btn_play = Button(ax_play, 'Play', color='#4CAF50', hovercolor='#66BB6A')
        self.btn_play.label.set_color('white')
        self.btn_play.label.set_fontweight('bold')
        self.btn_play.on_clicked(self._on_play)

        ax_pause = self.fig.add_axes([0.81, 0.91, 0.08, 0.035])
        self.btn_pause = Button(ax_pause, 'Pause', color='#FF9800', hovercolor='#FFB74D')
        self.btn_pause.label.set_color('white')
        self.btn_pause.label.set_fontweight('bold')
        self.btn_pause.on_clicked(self._on_pause)

        ax_reset = self.fig.add_axes([0.90, 0.91, 0.08, 0.035])
        self.btn_reset = Button(ax_reset, 'Reset', color='#F44336', hovercolor='#EF5350')
        self.btn_reset.label.set_color('white')
        self.btn_reset.label.set_fontweight('bold')
        self.btn_reset.on_clicked(self._on_reset)

    def _on_play(self, event):
        if self.is_playing:
            return
        if self.anim is None:
            self.anim = animation.FuncAnimation(
                self.fig, self._update_frame,
                frames=range(self.current_frame, self.total_frames + 1),
                interval=INTERVAL_MS,
                blit=False,
                repeat=False
            )
        else:
            self.anim.resume()
        self.is_playing = True
        self.fig.canvas.draw_idle()

    def _on_pause(self, event):
        if not self.is_playing or self.anim is None:
            return
        self.anim.pause()
        self.is_playing = False
        self.status_text.set_text(f'Paused (frame {self.current_frame}/{self.total_frames})')
        self.fig.canvas.draw_idle()

    def _on_reset(self, event):
        if self.anim is not None:
            self.anim.event_source.stop()
            self.anim = None
        self.is_playing = False
        self.current_frame = 0
        self.trail_data = [[] for _ in range(self.num_drones)]
        for dot in self.drone_dots:
            dot.set_data([], [])
        for trail in self.trail_lines:
            trail.set_data([], [])
        self.status_text.set_text('Ready')
        self.fig.canvas.draw_idle()

    def _update_frame(self, frame):
        self.current_frame = frame

        if frame >= len(self.positions):
            self.is_playing = False
            self.status_text.set_text('Done')
            return self.drone_dots + self.trail_lines + [self.status_text]

        # seg_idx - 1 是因為第 0 段是 start→WP0
        seg_idx = frame // FRAMES_PER_SEGMENT
        progress = (frame % FRAMES_PER_SEGMENT) / FRAMES_PER_SEGMENT

        # 顯示時把第 0 段標為 "Start -> WP0"
        if seg_idx == 0:
            label = f'Start -> WP0  Progress {progress:.0%}'
        else:
            label = f'WP{seg_idx-1} -> WP{seg_idx}  Progress {progress:.0%}'
        self.status_text.set_text(
            f'{label}  Frame {frame}/{self.total_frames}'
        )

        for i in range(self.num_drones):
            x, y = self.positions[frame][i]
            self.drone_dots[i].set_data([x], [y])

            self.trail_data[i].append((x, y))
            if len(self.trail_data[i]) > TRAIL_LENGTH:
                self.trail_data[i] = self.trail_data[i][-TRAIL_LENGTH:]
            trail_x = [p[0] for p in self.trail_data[i]]
            trail_y = [p[1] for p in self.trail_data[i]]
            self.trail_lines[i].set_data(trail_x, trail_y)

        return self.drone_dots + self.trail_lines + [self.status_text]


# ================================================================================
# 主流程
# ================================================================================

def visualize_from_file(filepath):
    """從 JSON 檔案讀取並視覺化"""
    with open(filepath, 'r') as f:
        data = json.load(f)

    origin = data['origin']
    conv = CoordinateConverter(origin[0], origin[1], 0)
    mission_type = data.get('mission_type', 'formation')
    is_sweep = mission_type == 'grid_sweep'

    fig, ax = plt.subplots(1, 1, figsize=(10, 8))
    fig.suptitle(f'Mission Verification - {mission_type}', fontsize=13, fontweight='bold')

    if is_sweep:
        plot_grid_sweep(ax, data, conv)
    else:
        plot_formation(ax, data, conv)

    _print_summary(data)

    animator = MissionAnimator(fig, ax, data, conv)
    animator.setup()

    plt.tight_layout(rect=[0, 0, 1, 0.95])
    plt.show()


def visualize_standalone():
    """獨立執行：使用內建測試資料"""
    drones = [
        (24.123450, 120.567800, 0.0),
        (24.123470, 120.567820, 0.0),
        (24.123440, 120.567810, 0.0),
        (24.123460, 120.567830, 0.0),
    ]
    rect_corners = [
        (24.12380, 120.56775),
        (24.12380, 120.56810),
        (24.12420, 120.56810),
        (24.12420, 120.56775),
    ]
    target = (24.12400, 120.56795, 10.0)

    planner = FormationPlanner(spacing=5.0, base_altitude=10.0, altitude_diff=2.0)

    fig = plt.figure(figsize=(16, 10))
    fig.suptitle('Mission Planner Verification (standalone)', fontsize=14, fontweight='bold')

    # M-Formation
    wps_m, origin_m, _ = planner.plan_formation_mission(drones, target, MissionType.M_FORMATION)
    conv_m = CoordinateConverter(origin_m[0], origin_m[1], 0)
    data_m = {
        'drone_ids': [f's0_{i + 1}' for i in range(len(drones))],
        'waypoints': wps_m,
        'drones_gps': drones,
        'target': target,
        'mission_type': 'M_FORMATION'
    }
    ax1 = fig.add_subplot(2, 2, 1)
    plot_formation(ax1, data_m, conv_m)

    # Grid Sweep 5m
    target_gs = (sum(c[0] for c in rect_corners) / 4,
                 sum(c[1] for c in rect_corners) / 4, 10.0)
    wps_g, origin_g, _ = planner.plan_formation_mission(
        drones, target_gs, MissionType.GRID_SWEEP,
        params={'rect_corners': rect_corners, 'line_spacing': 5.0, 'altitude': 10.0}
    )
    conv_g = CoordinateConverter(origin_g[0], origin_g[1], 0)
    data_g = {
        'drone_ids': [f's0_{i + 1}' for i in range(len(drones))],
        'waypoints': wps_g,
        'drones_gps': drones,
        'rect_corners': rect_corners,
        'mission_type': 'grid_sweep'
    }
    ax2 = fig.add_subplot(2, 2, 2)
    plot_grid_sweep(ax2, data_g, conv_g)

    # Leader-Follower
    route = [
        (24.12360, 120.56780),
        (24.12380, 120.56800),
        (24.12400, 120.56820),
        (24.12410, 120.56800),
        (24.12420, 120.56790),
    ]
    wps_lf, origin_lf, _ = planner.plan_formation_mission(
        drones, target, MissionType.LEADER_FOLLOWER,
        params={'route_waypoints': route, 'lateral_offset': 3.0,
                'longitudinal_spacing': 5.0, 'altitude': 10.0}
    )
    conv_lf = CoordinateConverter(origin_lf[0], origin_lf[1], 0)
    data_lf = {
        'drone_ids': [f's0_{i + 1}' for i in range(len(drones))],
        'waypoints': wps_lf,
        'drones_gps': drones,
        'route_waypoints': route,
        'target': target,
        'mission_type': 'LEADER_FOLLOWER'
    }
    ax3 = fig.add_subplot(2, 2, 3)
    plot_formation(ax3, data_lf, conv_lf)

    # 3D
    ax4 = fig.add_subplot(2, 2, 4, projection='3d')
    _plot_3d(ax4, data_g, conv_g)

    plt.tight_layout()
    plt.show()


def _plot_3d(ax, data, conv):
    """3D 視角"""
    if 'rect_corners' in data:
        rect_local = [conv.gps_to_local(c[0], c[1], 0)[:2] for c in data['rect_corners']]
        xs = [p[0] for p in rect_local] + [rect_local[0][0]]
        ys = [p[1] for p in rect_local] + [rect_local[0][1]]
        ax.plot(xs, ys, [0] * len(xs), 'k--', linewidth=1, alpha=0.4)

    for i, wps in enumerate(data['waypoints']):
        c = COLORS[i % len(COLORS)]
        local_wps = [conv.gps_to_local(wp[0], wp[1], wp[2]) for wp in wps]
        xs = [p[0] for p in local_wps]
        ys = [p[1] for p in local_wps]
        zs = [p[2] for p in local_wps]
        ax.plot(xs, ys, zs, '-', color=c, linewidth=1.5)
        if local_wps:
            ax.scatter(xs[0], ys[0], zs[0], color=c, s=50, marker='s')
            ax.scatter(xs[-1], ys[-1], zs[-1], color=c, s=50, marker='X')

    ax.set_title('3D view')
    ax.set_xlabel('East (m)')
    ax.set_ylabel('North (m)')
    ax.set_zlabel('Alt (m)')


def _print_summary(data):
    """終端印出摘要"""
    drone_ids = data['drone_ids']
    waypoints = data['waypoints']
    mission_type = data.get('mission_type', 'unknown')
    print(f"\n{'=' * 50}")
    print(f"Mission: {mission_type}")
    print(f"Drones: {len(drone_ids)}")
    print(f"{'=' * 50}")
    for i, did in enumerate(drone_ids):
        wps = waypoints[i]
        print(f"  {did}: {len(wps)} waypoints")
        for j, wp in enumerate(wps):
            print(f"    WP{j}: ({wp[0]:.6f}, {wp[1]:.6f}, {wp[2]:.1f}m)")
    print(f"{'=' * 50}\n")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Mission waypoint visualizer')
    parser.add_argument('--file', '-f', type=str, default=None,
                        help='JSON file from GUI (auto-generated)')
    args = parser.parse_args()

    if args.file:
        visualize_from_file(args.file)
    else:
        visualize_standalone()