import pytest
import tempfile
import os
import yaml
import numpy as np
from unittest.mock import patch, MagicMock
from elevator_traffic_analysis.plotter import ElevatorPlotter


class TestElevatorPlotterInit:
    """Tests for ElevatorPlotter initialization"""

    def test_plotter_initialization_with_valid_file(self):
        """
        Given: A valid simulation results file
        When: ElevatorPlotter is initialized
        Then: All configuration properties are correctly loaded
        """
        results = {
            'scenario_config': {
                'n_floors': 5,
                'n_elevators': 2,
                'elevator_capacity': 10,
                'elevator_logic': 'first_come_first_served'
            },
            'analysis': {
                'total_ticks': 10,
                'total_passengers': 3,
                'ticks': {
                    0: {
                        'elevators': {'E__0': {'current_floor': 'F__0', 'capacity': 10,
                                               'onboard_passengers_count': 0, 'onboard_passengers': []}},
                        'floors': {'F__0': {'waiting_passengers_count': 0, 'waiting_passengers': [],
                                           'dismbarked_passengers_count': 0, 'dismbarked_passengers': []}}
                    }
                }
            }
        }

        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
            yaml.dump(results, f)
            results_file = f.name

        try:
            plotter = ElevatorPlotter(results_file)
            assert plotter.yaml_file_path == results_file
            assert plotter.total_ticks == 10
            assert plotter.total_passengers == 3
            assert plotter.elevator_logic == 'first_come_first_served'
            assert plotter.n_floors == 5
            assert plotter.n_elevators == 2
        finally:
            os.unlink(results_file)

    def test_plotter_initialization_file_not_found(self):
        """
        Given: A non-existent results file path
        When: ElevatorPlotter is initialized
        Then: FileNotFoundError is raised
        """
        with pytest.raises(FileNotFoundError):
            ElevatorPlotter("nonexistent.yaml")

    def test_plotter_initialization_invalid_yaml(self):
        """
        Given: A file with invalid YAML content
        When: ElevatorPlotter is initialized
        Then: ValueError is raised
        """
        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
            f.write("invalid: yaml: content:")
            results_file = f.name

        try:
            with pytest.raises(ValueError):
                ElevatorPlotter(results_file)
        finally:
            os.unlink(results_file)


class TestExtractFloorNumber:
    """Tests for _extract_floor_number method"""

    def test_extract_floor_number_ground(self):
        """
        Given: A plotter with initialized results
        When: _extract_floor_number is called with ground floor ID 'F__0'
        Then: Floor number 0 is returned
        """
        results = self._create_minimal_results()
        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
            yaml.dump(results, f)
            results_file = f.name

        try:
            plotter = ElevatorPlotter(results_file)
            assert plotter._extract_floor_number('F__0') == 0
        finally:
            os.unlink(results_file)

    def test_extract_floor_number_high_floor(self):
        """
        Given: A plotter with initialized results
        When: _extract_floor_number is called with high floor ID 'F__42'
        Then: Floor number 42 is returned
        """
        results = self._create_minimal_results()
        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
            yaml.dump(results, f)
            results_file = f.name

        try:
            plotter = ElevatorPlotter(results_file)
            assert plotter._extract_floor_number('F__42') == 42
        finally:
            os.unlink(results_file)

    @staticmethod
    def _create_minimal_results():
        """Helper to create minimal valid results structure"""
        return {
            'scenario_config': {
                'n_floors': 5,
                'n_elevators': 1,
                'elevator_logic': 'first_come_first_served'
            },
            'analysis': {
                'total_ticks': 1,
                'total_passengers': 0,
                'ticks': {
                    0: {
                        'elevators': {'E__0': {'current_floor': 'F__0', 'capacity': 10,
                                               'onboard_passengers_count': 0, 'onboard_passengers': []}},
                        'floors': {'F__0': {'waiting_passengers_count': 0, 'waiting_passengers': [],
                                           'dismbarked_passengers_count': 0, 'dismbarked_passengers': []}}
                    }
                }
            }
        }


class TestExtractElevatorData:
    """Tests for _extract_elevator_data method"""

    def test_extract_elevator_data_single_elevator(self):
        """
        Given: Simulation results with a single elevator moving through floors
        When: _extract_elevator_data is called
        Then: Elevator position data is correctly extracted for all ticks
        """
        results = {
            'scenario_config': {
                'n_floors': 5,
                'n_elevators': 1,
                'elevator_logic': 'first_come_first_served'
            },
            'analysis': {
                'total_ticks': 3,
                'total_passengers': 0,
                'ticks': {
                    0: {'elevators': {'E__0': {'current_floor': 'F__0'}}, 'floors': {}},
                    1: {'elevators': {'E__0': {'current_floor': 'F__1'}}, 'floors': {}},
                    2: {'elevators': {'E__0': {'current_floor': 'F__2'}}, 'floors': {}}
                }
            }
        }

        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
            yaml.dump(results, f)
            results_file = f.name

        try:
            plotter = ElevatorPlotter(results_file)
            elevator_data = plotter._extract_elevator_data()
            
            assert 'E__0' in elevator_data
            assert len(elevator_data['E__0']) == 3
            assert elevator_data['E__0'] == [(0, 0), (1, 1), (2, 2)]
        finally:
            os.unlink(results_file)

    def test_extract_elevator_data_multiple_elevators(self):
        """
        Given: Simulation results with multiple elevators
        When: _extract_elevator_data is called
        Then: Position data for all elevators is correctly extracted
        """
        results = {
            'scenario_config': {
                'n_floors': 5,
                'n_elevators': 2,
                'elevator_logic': 'first_come_first_served'
            },
            'analysis': {
                'total_ticks': 2,
                'total_passengers': 0,
                'ticks': {
                    0: {'elevators': {'E__0': {'current_floor': 'F__0'},
                                     'E__1': {'current_floor': 'F__0'}}, 'floors': {}},
                    1: {'elevators': {'E__0': {'current_floor': 'F__1'},
                                     'E__1': {'current_floor': 'F__2'}}, 'floors': {}}
                }
            }
        }

        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
            yaml.dump(results, f)
            results_file = f.name

        try:
            plotter = ElevatorPlotter(results_file)
            elevator_data = plotter._extract_elevator_data()
            
            assert 'E__0' in elevator_data
            assert 'E__1' in elevator_data
            assert elevator_data['E__0'] == [(0, 0), (1, 1)]
            assert elevator_data['E__1'] == [(0, 0), (1, 2)]
        finally:
            os.unlink(results_file)


class TestGetElevatorColors:
    """Tests for _get_elevator_colors method"""

    def test_get_elevator_colors_returns_list(self):
        """
        Given: A plotter with initialized results
        When: _get_elevator_colors is called
        Then: A list of colors is returned
        """
        results = TestExtractFloorNumber._create_minimal_results()
        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
            yaml.dump(results, f)
            results_file = f.name

        try:
            plotter = ElevatorPlotter(results_file)
            colors = plotter._get_elevator_colors()
            assert isinstance(colors, list)
            assert len(colors) > 0
        finally:
            os.unlink(results_file)

    def test_get_elevator_colors_valid_format(self):
        """
        Given: A plotter with initialized results
        When: _get_elevator_colors is called
        Then: All colors are in valid hex format (#RRGGBB)
        """
        results = TestExtractFloorNumber._create_minimal_results()
        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
            yaml.dump(results, f)
            results_file = f.name

        try:
            plotter = ElevatorPlotter(results_file)
            colors = plotter._get_elevator_colors()
            for color in colors:
                assert color.startswith('#')
                assert len(color) == 7
        finally:
            os.unlink(results_file)


class TestCalculatePassengerStatistics:
    """Tests for _calculate_passenger_statistics method"""

    def test_calculate_passenger_statistics_no_passengers(self):
        """
        Given: Simulation results with no passengers
        When: _calculate_passenger_statistics is called
        Then: All statistics return zero values
        """
        results = {
            'scenario_config': {
                'n_floors': 5,
                'n_elevators': 1,
                'elevator_capacity': 10,
                'elevator_logic': 'first_come_first_served'
            },
            'analysis': {
                'total_ticks': 10,
                'total_passengers': 0,
                'ticks': {
                    0: {
                        'elevators': {'E__0': {'current_floor': 'F__0', 'capacity': 10,
                                               'onboard_passengers_count': 0}},
                        'floors': {'F__0': {'waiting_passengers_count': 0, 'waiting_passengers': [],
                                           'dismbarked_passengers_count': 0, 'dismbarked_passengers': []}}
                    }
                }
            }
        }

        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
            yaml.dump(results, f)
            results_file = f.name

        try:
            plotter = ElevatorPlotter(results_file)
            stats = plotter._calculate_passenger_statistics()
            
            assert stats['avg_wait_time'] == 0
            assert stats['avg_ride_time'] == 0
            assert stats['avg_serve_time'] == 0
            assert stats['passenger_count'] == 0
        finally:
            os.unlink(results_file)

    def test_calculate_passenger_statistics_with_passengers(self):
        """
        Given: Simulation results with passenger data
        When: _calculate_passenger_statistics is called
        Then: Average wait, ride, and serve times are correctly calculated
        """
        results = {
            'scenario_config': {
                'n_floors': 5,
                'n_elevators': 1,
                'elevator_capacity': 10,
                'elevator_logic': 'first_come_first_served'
            },
            'analysis': {
                'total_ticks': 20,
                'total_passengers': 2,
                'ticks': {
                    0: {
                        'elevators': {'E__0': {'current_floor': 'F__0', 'capacity': 10,
                                               'onboard_passengers_count': 0}},
                        'floors': {
                            'F__0': {
                                'waiting_passengers_count': 0,
                                'waiting_passengers': [],
                                'dismbarked_passengers_count': 1,
                                'dismbarked_passengers': [
                                    {
                                        'id': 'P__0',
                                        'stats': {
                                            'n_wait_ticks': 5,
                                            'n_ride_ticks': 10,
                                            'n_serve_ticks': 15
                                        }
                                    }
                                ]
                            },
                            'F__3': {
                                'waiting_passengers_count': 0,
                                'waiting_passengers': [],
                                'dismbarked_passengers_count': 1,
                                'dismbarked_passengers': [
                                    {
                                        'id': 'P__1',
                                        'stats': {
                                            'n_wait_ticks': 3,
                                            'n_ride_ticks': 6,
                                            'n_serve_ticks': 9
                                        }
                                    }
                                ]
                            }
                        }
                    }
                }
            }
        }

        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
            yaml.dump(results, f)
            results_file = f.name

        try:
            plotter = ElevatorPlotter(results_file)
            stats = plotter._calculate_passenger_statistics()
            
            assert stats['avg_wait_time'] == 4.0  # (5 + 3) / 2
            assert stats['avg_ride_time'] == 8.0  # (10 + 6) / 2
            assert stats['avg_serve_time'] == 12.0  # (15 + 9) / 2
            assert stats['passenger_count'] == 2
            assert stats['total_wait_time'] == 8
            assert stats['total_ride_time'] == 16
            assert stats['total_serve_time'] == 24
        finally:
            os.unlink(results_file)

    def test_calculate_passenger_statistics_throughput(self):
        """
        Given: Simulation results with multiple passengers and known ticks
        When: _calculate_passenger_statistics is called
        Then: Throughput (passengers per tick) is correctly calculated
        """
        results = {
            'scenario_config': {
                'n_floors': 5,
                'n_elevators': 1,
                'elevator_capacity': 10,
                'elevator_logic': 'first_come_first_served'
            },
            'analysis': {
                'total_ticks': 10,
                'total_passengers': 5,
                'ticks': {
                    0: {
                        'elevators': {'E__0': {'current_floor': 'F__0', 'capacity': 10,
                                               'onboard_passengers_count': 0}},
                        'floors': {'F__0': {'dismbarked_passengers': []}}
                    }
                }
            }
        }

        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
            yaml.dump(results, f)
            results_file = f.name

        try:
            plotter = ElevatorPlotter(results_file)
            stats = plotter._calculate_passenger_statistics()
            
            assert stats['throughput'] == 0.5  # 5 passengers / 10 ticks
        finally:
            os.unlink(results_file)


class TestPlotElevatorMovements:
    """Tests for plot_elevator_movements method"""

    @patch('matplotlib.pyplot.show')
    def test_plot_elevator_movements_basic(self, mock_show):
        """
        Given: Valid simulation results
        When: plot_elevator_movements is called
        Then: A plot is generated and saved successfully
        """
        results = {
            'scenario_config': {
                'n_floors': 5,
                'n_elevators': 1,
                'elevator_capacity': 10,
                'elevator_logic': 'first_come_first_served'
            },
            'analysis': {
                'total_ticks': 3,
                'total_passengers': 0,
                'ticks': {
                    0: {
                        'elevators': {'E__0': {'current_floor': 'F__0', 'capacity': 10,
                                               'onboard_passengers_count': 0, 'onboard_passengers': []}},
                        'floors': {'F__0': {'waiting_passengers_count': 0, 'waiting_passengers': [],
                                           'dismbarked_passengers_count': 0, 'dismbarked_passengers': []}}
                    },
                    1: {
                        'elevators': {'E__0': {'current_floor': 'F__1', 'capacity': 10,
                                               'onboard_passengers_count': 0, 'onboard_passengers': []}},
                        'floors': {'F__1': {'waiting_passengers_count': 0, 'waiting_passengers': [],
                                           'dismbarked_passengers_count': 0, 'dismbarked_passengers': []}}
                    },
                    2: {
                        'elevators': {'E__0': {'current_floor': 'F__2', 'capacity': 10,
                                               'onboard_passengers_count': 0, 'onboard_passengers': []}},
                        'floors': {'F__2': {'waiting_passengers_count': 0, 'waiting_passengers': [],
                                           'dismbarked_passengers_count': 0, 'dismbarked_passengers': []}}
                    }
                }
            }
        }

        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
            yaml.dump(results, f)
            results_file = f.name

        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as f:
            plot_file = f.name

        try:
            plotter = ElevatorPlotter(results_file)
            fig, (ax1, ax2) = plotter.plot_elevator_movements(
                save_path=plot_file,
                show_plot=False,
                show_passenger_ids=False
            )
            
            assert os.path.exists(plot_file)
            assert fig is not None
            assert ax1 is not None
            assert ax2 is not None
        finally:
            os.unlink(results_file)
            if os.path.exists(plot_file):
                os.unlink(plot_file)

    @patch('matplotlib.pyplot.show')
    def test_plot_elevator_movements_with_passenger_ids(self, mock_show):
        """
        Given: Simulation results with passengers
        When: plot_elevator_movements is called with show_passenger_ids=True
        Then: A plot with passenger IDs is generated successfully
        """
        results = {
            'scenario_config': {
                'n_floors': 5,
                'n_elevators': 1,
                'elevator_capacity': 10,
                'elevator_logic': 'first_come_first_served'
            },
            'analysis': {
                'total_ticks': 2,
                'total_passengers': 1,
                'ticks': {
                    0: {
                        'elevators': {
                            'E__0': {
                                'current_floor': 'F__0',
                                'capacity': 10,
                                'onboard_passengers_count': 1,
                                'onboard_passengers': [{'id': 'P__0', 'destination_floor': 'F__3'}]
                            }
                        },
                        'floors': {
                            'F__0': {
                                'waiting_passengers_count': 0,
                                'waiting_passengers': [],
                                'dismbarked_passengers_count': 0,
                                'dismbarked_passengers': []
                            }
                        }
                    },
                    1: {
                        'elevators': {
                            'E__0': {
                                'current_floor': 'F__1',
                                'capacity': 10,
                                'onboard_passengers_count': 1,
                                'onboard_passengers': [{'id': 'P__0', 'destination_floor': 'F__3'}]
                            }
                        },
                        'floors': {
                            'F__1': {
                                'waiting_passengers_count': 0,
                                'waiting_passengers': [],
                                'dismbarked_passengers_count': 0,
                                'dismbarked_passengers': []
                            }
                        }
                    }
                }
            }
        }

        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
            yaml.dump(results, f)
            results_file = f.name

        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as f:
            plot_file = f.name

        try:
            plotter = ElevatorPlotter(results_file)
            fig, _ = plotter.plot_elevator_movements(
                save_path=plot_file,
                show_plot=False,
                show_passenger_ids=True
            )
            
            assert os.path.exists(plot_file)
            assert fig is not None
        finally:
            os.unlink(results_file)
            if os.path.exists(plot_file):
                os.unlink(plot_file)

    @patch('matplotlib.pyplot.show')
    def test_plot_elevator_movements_no_save(self, mock_show):
        """
        Given: Valid simulation results
        When: plot_elevator_movements is called without save_path
        Then: A plot is generated without being saved to file
        """
        results = TestExtractFloorNumber._create_minimal_results()

        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
            yaml.dump(results, f)
            results_file = f.name

        try:
            plotter = ElevatorPlotter(results_file)
            fig, _ = plotter.plot_elevator_movements(
                save_path=None,
                show_plot=False
            )
            
            assert fig is not None
        finally:
            os.unlink(results_file)


class TestPlotWaitingPassengers:
    """Tests for _plot_waiting_passengers method"""

    @patch('matplotlib.pyplot.show')
    def test_plot_waiting_passengers_basic(self, mock_show):
        """
        Given: Simulation results with waiting passengers on different floors
        When: _plot_waiting_passengers is called
        Then: The waiting passengers plot is rendered without errors
        """
        results = {
            'scenario_config': {
                'n_floors': 3,
                'n_elevators': 1,
                'elevator_logic': 'first_come_first_served'
            },
            'analysis': {
                'total_ticks': 2,
                'total_passengers': 1,
                'ticks': {
                    0: {
                        'elevators': {'E__0': {'current_floor': 'F__0'}},
                        'floors': {
                            'F__0': {
                                'waiting_passengers_count': 2,
                                'waiting_passengers': [
                                    {'id': 'P__0'},
                                    {'id': 'P__1'}
                                ]
                            },
                            'F__1': {
                                'waiting_passengers_count': 0,
                                'waiting_passengers': []
                            }
                        }
                    },
                    1: {
                        'elevators': {'E__0': {'current_floor': 'F__0'}},
                        'floors': {
                            'F__0': {
                                'waiting_passengers_count': 0,
                                'waiting_passengers': []
                            },
                            'F__1': {
                                'waiting_passengers_count': 1,
                                'waiting_passengers': [{'id': 'P__2'}]
                            }
                        }
                    }
                }
            }
        }

        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
            yaml.dump(results, f)
            results_file = f.name

        try:
            plotter = ElevatorPlotter(results_file)
            import matplotlib.pyplot as plt
            fig, ax = plt.subplots()
            
            # Should not raise any errors
            plotter._plot_waiting_passengers(ax, show_passenger_ids=False)
        finally:
            os.unlink(results_file)
