#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
AGV Path Tracking Visualization Script
Reads control sequence and trajectory data from CSV files and visualizes them
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.patches import Rectangle
from matplotlib.widgets import Button
import sys
import os

class AGVVisualizer:
    def __init__(self, trajectory_file=None, control_file=None):
        """
        Initialize AGV Visualizer

        Args:
            trajectory_file: Path to trajectory CSV file
            control_file: Path to control sequence CSV file
        """
        self.trajectory_file = trajectory_file
        self.control_file = control_file
        self.trajectory_data = None
        self.control_data = None

    def load_data(self):
        """Load data from CSV files"""
        if self.trajectory_file and os.path.exists(self.trajectory_file):
            print(f"Loading trajectory from: {self.trajectory_file}")
            self.trajectory_data = np.loadtxt(
                self.trajectory_file,
                delimiter=',',
                skiprows=2,  # Skip comment lines
                usecols=(0, 1, 2)  # x, y, theta(rad)
            )
            print(f"Loaded {len(self.trajectory_data)} trajectory points")
        else:
            print(f"Warning: Trajectory file not found: {self.trajectory_file}")

        if self.control_file and os.path.exists(self.control_file):
            print(f"Loading control sequence from: {self.control_file}")
            self.control_data = np.loadtxt(
                self.control_file,
                delimiter=',',
                skiprows=2,  # Skip comment lines
                usecols=(0, 1, 2, 3)  # time, velocity, steering(rad), steering(deg)
            )
            print(f"Loaded {len(self.control_data)} control points")
        else:
            print(f"Warning: Control file not found: {self.control_file}")

    def plot_trajectory(self):
        """Plot trajectory and AGV path"""
        if self.trajectory_data is None:
            print("Error: No trajectory data loaded!")
            return

        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 7))

        # Plot 1: XY Trajectory
        ax1.set_title('AGV Trajectory (XY Plane)', fontsize=14, fontweight='bold')
        ax1.set_xlabel('X (m)', fontsize=12)
        ax1.set_ylabel('Y (m)', fontsize=12)
        ax1.grid(True, alpha=0.3)
        ax1.axis('equal')

        # Extract data
        x = self.trajectory_data[:, 0]
        y = self.trajectory_data[:, 1]
        theta = self.trajectory_data[:, 2]

        # Plot trajectory line
        ax1.plot(x, y, 'b-', linewidth=2, label='Trajectory', alpha=0.7)

        # Plot start and end points
        ax1.plot(x[0], y[0], 'go', markersize=15, label='Start', markeredgecolor='black', markeredgewidth=2)
        ax1.plot(x[-1], y[-1], 'ro', markersize=15, label='End', markeredgecolor='black', markeredgewidth=2)

        # Plot heading arrows at intervals
        step = max(1, len(x) // 20)  # Show about 20 arrows
        for i in range(0, len(x), step):
            arrow_length = 0.3
            dx = arrow_length * np.cos(theta[i])
            dy = arrow_length * np.sin(theta[i])
            ax1.arrow(x[i], y[i], dx, dy,
                     head_width=0.15, head_length=0.1,
                     fc='red', ec='red', alpha=0.6)

        ax1.legend(fontsize=10, loc='upper right')

        # Plot 2: Heading angle over time
        ax2.set_title('AGV Heading Angle', fontsize=14, fontweight='bold')
        ax2.set_xlabel('Step', fontsize=12)
        ax2.set_ylabel('Heading Angle (degrees)', fontsize=12)
        ax2.grid(True, alpha=0.3)

        theta_deg = np.degrees(theta)
        ax2.plot(theta_deg, 'b-', linewidth=2)
        ax2.axhline(y=0, color='k', linestyle='--', alpha=0.3)

        plt.tight_layout()
        plt.savefig('trajectory_plot.png', dpi=300, bbox_inches='tight')
        print("Trajectory plot saved as 'trajectory_plot.png'")
        plt.show()

    def plot_controls(self):
        """Plot control sequence (velocity and steering)"""
        if self.control_data is None:
            print("Error: No control data loaded!")
            return

        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 8))

        time = self.control_data[:, 0]
        velocity = self.control_data[:, 1]
        steering_deg = self.control_data[:, 3]

        # Plot 1: Velocity
        ax1.set_title('AGV Velocity Profile', fontsize=14, fontweight='bold')
        ax1.set_xlabel('Time (s)', fontsize=12)
        ax1.set_ylabel('Velocity (m/s)', fontsize=12)
        ax1.grid(True, alpha=0.3)
        ax1.plot(time, velocity, 'b-', linewidth=2)
        ax1.fill_between(time, 0, velocity, alpha=0.3)

        # Add statistics
        avg_vel = np.mean(velocity)
        max_vel = np.max(velocity)
        min_vel = np.min(velocity)
        ax1.axhline(y=avg_vel, color='r', linestyle='--', linewidth=1.5,
                   label=f'Average: {avg_vel:.3f} m/s')
        ax1.legend(fontsize=10)

        # Plot 2: Steering Angle
        ax2.set_title('AGV Steering Angle', fontsize=14, fontweight='bold')
        ax2.set_xlabel('Time (s)', fontsize=12)
        ax2.set_ylabel('Steering Angle (degrees)', fontsize=12)
        ax2.grid(True, alpha=0.3)
        ax2.plot(time, steering_deg, 'g-', linewidth=2)
        ax2.fill_between(time, 0, steering_deg, alpha=0.3,
                        where=(steering_deg >= 0), color='green', interpolate=True)
        ax2.fill_between(time, 0, steering_deg, alpha=0.3,
                        where=(steering_deg < 0), color='red', interpolate=True)
        ax2.axhline(y=0, color='k', linestyle='-', alpha=0.3)

        # Add statistics
        avg_steer = np.mean(steering_deg)
        ax2.axhline(y=avg_steer, color='r', linestyle='--', linewidth=1.5,
                   label=f'Average: {avg_steer:.3f}°')
        ax2.legend(fontsize=10)

        plt.tight_layout()
        plt.savefig('control_plot.png', dpi=300, bbox_inches='tight')
        print("Control plot saved as 'control_plot.png'")
        plt.show()

    def plot_combined(self):
        """Plot combined visualization with trajectory and controls"""
        if self.trajectory_data is None or self.control_data is None:
            print("Error: Missing data! Need both trajectory and control files.")
            return

        fig = plt.figure(figsize=(16, 10))
        gs = fig.add_gridspec(2, 2, hspace=0.3, wspace=0.3)

        # Extract data
        x = self.trajectory_data[:, 0]
        y = self.trajectory_data[:, 1]
        theta = self.trajectory_data[:, 2]

        time = self.control_data[:, 0]
        velocity = self.control_data[:, 1]
        steering_deg = self.control_data[:, 3]

        # Plot 1: XY Trajectory (Large)
        ax1 = fig.add_subplot(gs[:, 0])
        ax1.set_title('AGV Trajectory', fontsize=16, fontweight='bold')
        ax1.set_xlabel('X (m)', fontsize=12)
        ax1.set_ylabel('Y (m)', fontsize=12)
        ax1.grid(True, alpha=0.3)
        ax1.axis('equal')

        ax1.plot(x, y, 'b-', linewidth=3, label='Trajectory', alpha=0.8)
        ax1.plot(x[0], y[0], 'go', markersize=15, label='Start',
                markeredgecolor='black', markeredgewidth=2)
        ax1.plot(x[-1], y[-1], 'ro', markersize=15, label='End',
                markeredgecolor='black', markeredgewidth=2)

        # Plot heading arrows
        step = max(1, len(x) // 15)
        for i in range(0, len(x), step):
            arrow_length = 0.4
            dx = arrow_length * np.cos(theta[i])
            dy = arrow_length * np.sin(theta[i])
            ax1.arrow(x[i], y[i], dx, dy,
                     head_width=0.2, head_length=0.15,
                     fc='red', ec='red', alpha=0.6)

        ax1.legend(fontsize=11, loc='best')

        # Plot 2: Velocity
        ax2 = fig.add_subplot(gs[0, 1])
        ax2.set_title('Velocity Profile', fontsize=14, fontweight='bold')
        ax2.set_xlabel('Time (s)', fontsize=11)
        ax2.set_ylabel('Velocity (m/s)', fontsize=11)
        ax2.grid(True, alpha=0.3)
        ax2.plot(time, velocity, 'b-', linewidth=2)
        ax2.fill_between(time, 0, velocity, alpha=0.3)

        avg_vel = np.mean(velocity)
        ax2.axhline(y=avg_vel, color='r', linestyle='--', linewidth=1.5,
                   label=f'Avg: {avg_vel:.3f} m/s')
        ax2.legend(fontsize=9)

        # Plot 3: Steering Angle
        ax3 = fig.add_subplot(gs[1, 1])
        ax3.set_title('Steering Angle', fontsize=14, fontweight='bold')
        ax3.set_xlabel('Time (s)', fontsize=11)
        ax3.set_ylabel('Steering Angle (°)', fontsize=11)
        ax3.grid(True, alpha=0.3)
        ax3.plot(time, steering_deg, 'g-', linewidth=2)
        ax3.fill_between(time, 0, steering_deg, alpha=0.3,
                        where=(steering_deg >= 0), color='green')
        ax3.fill_between(time, 0, steering_deg, alpha=0.3,
                        where=(steering_deg < 0), color='red')
        ax3.axhline(y=0, color='k', linestyle='-', alpha=0.3)

        avg_steer = np.mean(steering_deg)
        ax3.axhline(y=avg_steer, color='r', linestyle='--', linewidth=1.5,
                   label=f'Avg: {avg_steer:.3f}°')
        ax3.legend(fontsize=9)

        plt.savefig('combined_plot.png', dpi=300, bbox_inches='tight')
        print("Combined plot saved as 'combined_plot.png'")
        plt.show()

    def animate_agv(self, interval=100):
        """Create animation of AGV movement with pause/play functionality"""
        if self.trajectory_data is None:
            print("Error: No trajectory data loaded!")
            return

        x = self.trajectory_data[:, 0]
        y = self.trajectory_data[:, 1]
        theta = self.trajectory_data[:, 2]

        fig, ax = plt.subplots(figsize=(10, 10))
        plt.subplots_adjust(bottom=0.15)  # Make room for button
        ax.set_title('AGV Animation', fontsize=16, fontweight='bold')
        ax.set_xlabel('X (m)', fontsize=12)
        ax.set_ylabel('Y (m)', fontsize=12)
        ax.grid(True, alpha=0.3)
        ax.axis('equal')

        # Plot full trajectory
        ax.plot(x, y, 'b--', linewidth=1, alpha=0.5, label='Reference Path')
        ax.plot(x[0], y[0], 'go', markersize=12, label='Start')
        ax.plot(x[-1], y[-1], 'ro', markersize=12, label='End')

        # Initialize animated elements
        line, = ax.plot([], [], 'r-', linewidth=2, label='Traveled Path')
        agv_body = Rectangle((0, 0), 0.6, 0.3, fill=True,
                            facecolor='green', edgecolor='black', linewidth=2)
        ax.add_patch(agv_body)
        heading_arrow = ax.arrow(0, 0, 0, 0, head_width=0.15,
                                head_length=0.1, fc='red', ec='red')

        ax.legend(fontsize=10)

        # Set axis limits
        margin = 1.0
        ax.set_xlim(min(x) - margin, max(x) + margin)
        ax.set_ylim(min(y) - margin, max(y) + margin)

        # Animation state
        anim_running = [True]  # Use list to allow modification in nested function

        # Create pause/play button
        button_ax = plt.axes([0.45, 0.05, 0.1, 0.04])
        button = Button(button_ax, 'Pause', color='lightgray', hovercolor='gray')

        def toggle_pause(event):
            """Toggle animation pause/play state"""
            if anim_running[0]:
                anim.event_source.stop()
                button.label.set_text('Play')
                anim_running[0] = False
            else:
                anim.event_source.start()
                button.label.set_text('Pause')
                anim_running[0] = True
            plt.draw()

        button.on_clicked(toggle_pause)

        def init():
            line.set_data([], [])
            return line, agv_body, heading_arrow

        def animate(frame):
            # Update traveled path
            line.set_data(x[:frame+1], y[:frame+1])

            # Update AGV position and orientation
            agv_x = x[frame]
            agv_y = y[frame]
            agv_theta = theta[frame]

            # AGV body dimensions
            agv_length = 0.6
            agv_width = 0.3

            # Create transformation matrix
            import matplotlib.transforms as transforms
            t = transforms.Affine2D().rotate(agv_theta).translate(agv_x, agv_y) + ax.transData
            agv_body.set_transform(t)
            agv_body.set_xy((-agv_length/2, -agv_width/2))

            # Remove old arrow and create new one
            nonlocal heading_arrow
            if heading_arrow in ax.patches:
                heading_arrow.remove()

            arrow_length = 0.5
            dx = arrow_length * np.cos(agv_theta)
            dy = arrow_length * np.sin(agv_theta)
            heading_arrow = ax.arrow(agv_x, agv_y, dx, dy,
                                    head_width=0.2, head_length=0.15,
                                    fc='red', ec='red', alpha=0.8)

            return line, agv_body, heading_arrow

        anim = FuncAnimation(fig, animate, init_func=init,
                           frames=len(x), interval=interval,
                           blit=False, repeat=True)

        plt.show()


def main():
    """Main function"""
    print("=" * 60)
    print("AGV Path Tracking Visualization Tool")
    print("=" * 60)

    # Default file paths
    default_trajectory = "trajectory.csv"
    default_control = "control_sequence.csv"

    # Check if files exist in current directory
    trajectory_file = default_trajectory if os.path.exists(default_trajectory) else None
    control_file = default_control if os.path.exists(default_control) else None

    # Command line arguments
    if len(sys.argv) > 1:
        trajectory_file = sys.argv[1]
    if len(sys.argv) > 2:
        control_file = sys.argv[2]

    if trajectory_file is None and control_file is None:
        print("\nError: No data files found!")
        print("\nUsage:")
        print("  python visualize.py [trajectory_file] [control_file]")
        print("\nDefault files:")
        print(f"  - {default_trajectory}")
        print(f"  - {default_control}")
        print("\nPlease run the AGV demo first to generate data files.")
        return

    # Create visualizer
    visualizer = AGVVisualizer(trajectory_file, control_file)
    visualizer.load_data()

    # Show menu
    while True:
        print("\n" + "=" * 60)
        print("Visualization Options:")
        print("=" * 60)
        print("1. Plot Trajectory")
        print("2. Plot Control Sequence")
        print("3. Plot Combined View")
        print("4. Animate AGV Movement")
        print("5. Exit")
        print("=" * 60)

        choice = input("Select option (1-5): ").strip()

        if choice == '1':
            visualizer.plot_trajectory()
        elif choice == '2':
            visualizer.plot_controls()
        elif choice == '3':
            visualizer.plot_combined()
        elif choice == '4':
            visualizer.animate_agv(interval=50)
        elif choice == '5':
            print("Exiting...")
            break
        else:
            print("Invalid option! Please select 1-5.")


if __name__ == "__main__":
    main()
