from __future__ import annotations

from typing import Any, Dict, List

import matplotlib.pyplot as plt
import numpy as np
import yaml


class ElevatorPlotter:

    def __init__(self, yaml_file_path: str):
        self.yaml_file_path = yaml_file_path
        self.data = self._load_yaml_data()
        self.ticks = self.data['analysis']['ticks']
        self.total_ticks = self.data['analysis']['total_ticks']
        self.total_passengers = self.data['analysis']['total_passengers']
        self.elevator_logic = self.data['scenario_config']['elevator_logic']
        self.n_floors = self.data['scenario_config']['n_floors']
        self.n_elevators = self.data['scenario_config']['n_elevators']

    def _load_yaml_data(self) -> Dict[str, Any]:
        try:
            with open(self.yaml_file_path, 'r', encoding='utf-8') as file:
                return yaml.safe_load(file)
        except FileNotFoundError as exc:
            raise FileNotFoundError(f"Simulation results file not found: {self.yaml_file_path}") from exc
        except yaml.YAMLError as e:
            raise ValueError(f"Error parsing YAML file: {e}") from e

    def _extract_floor_number(self, floor_id: str) -> int:
        """Extract numeric floor number from floor ID (e.g., 'F__0' -> 0)."""
        return int(floor_id.split('__')[1])

    def _extract_elevator_data(self) -> Dict[str, List[tuple]]:
        elevator_data = {}

        for tick_str, tick_data in self.ticks.items():
            tick = int(tick_str)
            elevators = tick_data.get('elevators', {})

            for elevator_id, elevator_info in elevators.items():
                if elevator_id not in elevator_data:
                    elevator_data[elevator_id] = []

                current_floor_id = elevator_info['current_floor']
                floor_number = self._extract_floor_number(current_floor_id)
                elevator_data[elevator_id].append((tick, floor_number))

        return elevator_data

    def _get_elevator_colors(self) -> List[str]:
        colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd',
                 '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf']
        return colors

    def _calculate_passenger_statistics(self) -> Dict[str, float]:
        """Calculate passenger statistics including total waiting time."""
        total_wait_time = 0
        total_ride_time = 0
        total_serve_time = 0
        passenger_count = 0

        for _, tick_data in self.ticks.items():
            floors = tick_data.get('floors', {})
            for _, floor_data in floors.items():
                disembarked = floor_data.get('dismbarked_passengers', [])
                for passenger in disembarked:
                    stats = passenger.get('stats')
                    if stats:
                        total_wait_time += stats.get('n_wait_ticks', 0)
                        total_ride_time += stats.get('n_ride_ticks', 0)
                        total_serve_time += stats.get('n_serve_ticks', 0)
                        passenger_count += 1

        avg_wait_time = total_wait_time / passenger_count if passenger_count > 0 else 0
        avg_ride_time = total_ride_time / passenger_count if passenger_count > 0 else 0
        avg_serve_time = total_serve_time / passenger_count if passenger_count > 0 else 0

        throughput = self.total_passengers / self.total_ticks if self.total_ticks > 0 else 0
        elevator_capacity = self.data['scenario_config'].get('elevator_capacity', 'N/A')

        total_elevator_utilization = 0
        max_possible_utilization = 0

        for _, tick_data in self.ticks.items():
            elevators = tick_data.get('elevators', {})
            for _, elevator_data in elevators.items():
                passenger_count_tick = elevator_data.get('onboard_passengers_count', 0)
                capacity = elevator_data.get('capacity', elevator_capacity)
                total_elevator_utilization += passenger_count_tick
                max_possible_utilization += capacity if isinstance(capacity, (int, float)) else 0

        utilization_rate = (total_elevator_utilization / max_possible_utilization * 100) if max_possible_utilization > 0 else 0

        return {
            'total_wait_time': total_wait_time,
            'total_ride_time': total_ride_time,
            'total_serve_time': total_serve_time,
            'avg_wait_time': avg_wait_time,
            'avg_ride_time': avg_ride_time,
            'avg_serve_time': avg_serve_time,
            'passenger_count': passenger_count,
            'throughput': throughput,
            'utilization_rate': utilization_rate,
            'elevator_capacity': elevator_capacity
        }

    def _plot_waiting_passengers(self, ax, show_passenger_ids: bool = False):
        """Plot waiting passengers information on each floor as a heatmap."""
        # Extract waiting passengers data
        ticks = sorted(self.ticks.keys())
        floor_ids = sorted(set(floor_id for tick_data in self.ticks.values()
                              for floor_id in tick_data.get('floors', {}).keys()))

        # Sort floor_ids by floor number (ascending) and reverse for proper display order
        floor_ids = sorted(floor_ids, key=self._extract_floor_number, reverse=True)

        # Create matrix for waiting passengers count
        waiting_data = np.zeros((len(floor_ids), len(ticks)))

        for tick_idx, tick in enumerate(ticks):
            tick_data = self.ticks[tick]
            floors = tick_data.get('floors', {})

            for floor_idx, floor_id in enumerate(floor_ids):
                if floor_id in floors:
                    waiting_count = floors[floor_id].get('waiting_passengers_count', 0)
                    waiting_data[floor_idx, tick_idx] = waiting_count

        # Create heatmap
        ax.imshow(waiting_data, cmap='Reds', aspect='auto', interpolation='nearest',
                       vmin=0, vmax=max(1, np.max(waiting_data)))

        # Customize the waiting passengers plot
        ax.set_xlabel('Time (Ticks)', fontsize=12, fontweight='bold')
        ax.set_ylabel('Floor', fontsize=12, fontweight='bold')
        ax.set_title('Waiting Passengers per Floor', fontsize=14, fontweight='bold')

        # Set ticks and labels
        ax.set_xticks(range(len(ticks)))
        ax.set_xticklabels(ticks)
        ax.set_yticks(range(len(floor_ids)))
        ax.set_yticklabels(floor_ids)

        # Add text annotations for passenger counts and optionally IDs
        for floor_idx, floor_id in enumerate(floor_ids): # pylint: disable=too-many-nested-blocks
            for tick_idx, tick in enumerate(ticks):
                count = waiting_data[floor_idx, tick_idx]
                if count > 0:
                    passenger_text = str(int(count))

                    if show_passenger_ids:
                        tick = ticks[tick_idx]
                        floor_id = floor_ids[floor_idx] #pylint: disable=unnecessary-list-index-lookup
                        tick_data = self.ticks[tick]
                        floors = tick_data.get('floors', {})

                        if floor_id in floors:
                            waiting_passengers = floors[floor_id].get('waiting_passengers', [])
                            if waiting_passengers:
                                passenger_ids = [p['id'].replace('P__', 'P') for p in waiting_passengers]
                                passenger_text = '\n'.join(passenger_ids)

                    text_color = 'white' if count > np.max(waiting_data) / 2 else 'black'
                    ax.text(tick_idx, floor_idx, passenger_text,
                           ha='center', va='center', color=text_color,
                           fontweight='bold', fontsize=8)

        # Style the subplot
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        ax.spines['left'].set_linewidth(1.5)
        ax.spines['bottom'].set_linewidth(1.5)

    def _add_passenger_annotations(self, ax):
        """Add passenger ID annotations to show where passengers are at each tick."""
        # Track passenger positions for each tick
        for tick_str, tick_data in self.ticks.items():
            tick = int(tick_str)
            elevators = tick_data.get('elevators', {})
            floors = tick_data.get('floors', {})

            # Annotate passengers in elevators
            for _, elevator_data in elevators.items():
                elevator_floor = self._extract_floor_number(elevator_data['current_floor'])
                onboard_passengers = elevator_data.get('onboard_passengers', [])

                if onboard_passengers:
                    passenger_ids = [p['id'] for p in onboard_passengers]
                    passenger_text = ','.join([pid.replace('P__', 'P') for pid in passenger_ids])
                    ax.annotate(passenger_text,
                              xy=(tick, elevator_floor),
                              xytext=(5, 5),
                              textcoords='offset points',
                              fontsize=8,
                              color='darkblue',
                              bbox={"boxstyle": 'round,pad=0.2', 'facecolor': 'lightblue', 'alpha': 0.7})

            # Annotate waiting passengers on floors
            for floor_id, floor_data in floors.items():
                floor_num = self._extract_floor_number(floor_id)
                waiting_passengers = floor_data.get('waiting_passengers', [])

                if waiting_passengers:
                    passenger_ids = [p['id'] for p in waiting_passengers]
                    passenger_text = ','.join([pid.replace('P__', 'P') for pid in passenger_ids])
                    ax.annotate(passenger_text,
                              xy=(tick, floor_num),
                              xytext=(-15, -15),
                              textcoords='offset points',
                              fontsize=8,
                              color='darkred',
                              bbox={"boxstyle": 'round,pad=0.2', 'facecolor': 'lightyellow', 'alpha': 0.7})

    def plot_elevator_movements(self, save_path: str = None, show_plot: bool = True, show_passenger_ids: bool = False):

        elevator_data = self._extract_elevator_data()
        colors = self._get_elevator_colors()

        # Calculate dynamic figure width based on number of ticks
        num_ticks = len(self.ticks)
        # Base width of 8, plus 0.15 per tick, with minimum of 10 and maximum of 25
        width = max(10, min(25, 8 + num_ticks * 0.15))

        # Create figure with subplots (elevator movement + waiting passengers)
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(width, 10), height_ratios=[3, 1])

        # Plot each elevator's movement (top subplot)
        for i, (elevator_id, positions) in enumerate(elevator_data.items()):
            if not positions:
                continue

            ticks, floors = zip(*positions)
            color = colors[i % len(colors)]

            # Plot line connecting positions
            ax1.plot(ticks, floors, color=color, linewidth=2, alpha=0.7,
                   label=f'Elevator {elevator_id}')

            # Plot dots for each position
            ax1.scatter(ticks, floors, color=color, s=50, zorder=5, alpha=0.8)

        # Add passenger information annotations if requested
        if show_passenger_ids:
            self._add_passenger_annotations(ax1)

        # Customize the elevator movement plot
        self._customize_plot(ax1)

        # Add legend for elevator movement
        if elevator_data:
            ax1.legend(loc='upper right', framealpha=0.9)

        # Plot waiting passengers (bottom subplot)
        self._plot_waiting_passengers(ax2, show_passenger_ids)

        # Adjust layout
        plt.tight_layout()

        # Save plot if path provided
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"Plot saved to: {save_path}")

        # Show plot if requested
        if show_plot:
            plt.show()

        return fig, (ax1, ax2)

    def _customize_plot(self, ax):
        """Customize the plot appearance and labels."""
        stats = self._calculate_passenger_statistics()

        title = (f"Elevator Traffic Analysis - {self.elevator_logic.replace('_', ' ').title()}\n"
                f"Duration: {self.total_ticks} ticks | "
                f"Throughput: {stats['throughput']:.2f} passengers/tick | "
                f"Utilization: {stats['utilization_rate']:.1f}%")
        ax.set_title(title, fontsize=14, fontweight='bold', pad=20)

        ax.set_xlabel('Time (Ticks)', fontsize=12, fontweight='bold')
        ax.set_ylabel('Floor', fontsize=12, fontweight='bold')

        ax.set_xlim(-0.5, self.total_ticks + 0.5)
        ax.set_ylim(-0.5, self.n_floors - 0.5)

        ax.set_xticks(range(0, self.total_ticks + 1))
        ax.set_yticks(range(0, self.n_floors))
        ax.set_yticklabels([f'F__{i}' for i in range(0, self.n_floors)])

        ax.grid(True, alpha=0.3, linestyle='--')

        combined_text = (f"Building Configuration:\n"
                        f"• {self.n_floors} floors, {self.n_elevators} elevators\n"
                        f"• Capacity: {stats['elevator_capacity']} passengers/elevator\n"
                        f"• Total passengers: {self.total_passengers}\n"
                        f"\nPerformance Metrics:\n"
                        f"• Avg Wait Time: {stats['avg_wait_time']:.1f} ticks\n"
                        f"• Avg Ride Time: {stats['avg_ride_time']:.1f} ticks\n"
                        f"• Avg Total Service: {stats['avg_serve_time']:.1f} ticks\n"
                        f"• Total Wait Time: {stats['total_wait_time']:.0f} ticks")

        combined_props = {"boxstyle": 'round,pad=0.5', "facecolor": 'lightblue', "alpha": 0.8}
        ax.text(0.02, 0.98, combined_text, transform=ax.transAxes, fontsize=9,
                verticalalignment='top', bbox=combined_props)

        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        ax.spines['left'].set_linewidth(1.5)
        ax.spines['bottom'].set_linewidth(1.5)
