"""
Scenario comparison visualization module for HVAC Load Calculator.
This module provides visualization tools for comparing different scenarios.
"""

import streamlit as st
import pandas as pd
import numpy as np
import plotly.graph_objects as go
import plotly.express as px
from typing import Dict, List, Any, Optional, Tuple
import math

# Import calculation modules
from utils.cooling_load import CoolingLoad
from utils.heating_load import HeatingLoad


class ScenarioComparison:
    """Class for scenario comparison visualization."""
    
    @staticmethod
    def create_scenario_summary_table(scenarios: Dict[str, Dict[str, Any]]) -> pd.DataFrame:
        """
        Create a summary table of different scenarios.
        
        Args:
            scenarios: Dictionary with scenario data
            
        Returns:
            DataFrame with scenario summary
        """
        # Initialize data
        data = []
        
        # Process scenarios
        for scenario_name, scenario_data in scenarios.items():
            # Extract cooling and heating loads
            cooling_loads = scenario_data.get("cooling_loads", {})
            heating_loads = scenario_data.get("heating_loads", {})
            
            # Create summary row
            row = {
                "Scenario": scenario_name,
                "Cooling Load (W)": cooling_loads.get("total", 0),
                "Sensible Heat Ratio": cooling_loads.get("sensible_heat_ratio", 0),
                "Heating Load (W)": heating_loads.get("total", 0)
            }
            
            # Add to data
            data.append(row)
        
        # Create DataFrame
        df = pd.DataFrame(data)
        
        return df
    
    @staticmethod
    def create_load_comparison_chart(scenarios: Dict[str, Dict[str, Any]], load_type: str = "cooling") -> go.Figure:
        """
        Create a bar chart comparing loads across scenarios.
        
        Args:
            scenarios: Dictionary with scenario data
            load_type: Type of load to compare ("cooling" or "heating")
            
        Returns:
            Plotly figure with load comparison
        """
        # Initialize data
        scenario_names = []
        total_loads = []
        component_loads = {}
        
        # Process scenarios
        for scenario_name, scenario_data in scenarios.items():
            # Extract loads based on load type
            if load_type == "cooling":
                loads = scenario_data.get("cooling_loads", {})
                components = ["walls", "roofs", "floors", "windows_conduction", "windows_solar", 
                             "doors", "infiltration_sensible", "infiltration_latent", 
                             "people_sensible", "people_latent", "lights", "equipment_sensible", "equipment_latent"]
            else:  # heating
                loads = scenario_data.get("heating_loads", {})
                components = ["walls", "roofs", "floors", "windows", "doors", 
                             "infiltration_sensible", "infiltration_latent", 
                             "ventilation_sensible", "ventilation_latent"]
            
            # Add scenario name
            scenario_names.append(scenario_name)
            
            # Add total load
            total_loads.append(loads.get("total", 0))
            
            # Add component loads
            for component in components:
                if component not in component_loads:
                    component_loads[component] = []
                
                component_loads[component].append(loads.get(component, 0))
        
        # Create figure
        fig = go.Figure()
        
        # Add total load bars
        fig.add_trace(go.Bar(
            x=scenario_names,
            y=total_loads,
            name="Total Load",
            marker_color="rgba(55, 83, 109, 0.7)",
            opacity=0.7
        ))
        
        # Add component load bars
        for component, loads in component_loads.items():
            # Skip components with zero loads
            if sum(loads) == 0:
                continue
            
            # Format component name for display
            display_name = component.replace("_", " ").title()
            
            fig.add_trace(go.Bar(
                x=scenario_names,
                y=loads,
                name=display_name,
                visible="legendonly"
            ))
        
        # Update layout
        title = f"{load_type.title()} Load Comparison"
        y_title = f"{load_type.title()} Load (W)"
        
        fig.update_layout(
            title=title,
            xaxis_title="Scenario",
            yaxis_title=y_title,
            barmode="group",
            height=500,
            legend=dict(
                orientation="h",
                yanchor="bottom",
                y=1.02,
                xanchor="right",
                x=1
            )
        )
        
        return fig
    
    @staticmethod
    def create_percentage_difference_chart(scenarios: Dict[str, Dict[str, Any]], 
                                         baseline_scenario: str,
                                         load_type: str = "cooling") -> go.Figure:
        """
        Create a bar chart showing percentage differences from a baseline scenario.
        
        Args:
            scenarios: Dictionary with scenario data
            baseline_scenario: Name of the baseline scenario
            load_type: Type of load to compare ("cooling" or "heating")
            
        Returns:
            Plotly figure with percentage difference chart
        """
        # Check if baseline scenario exists
        if baseline_scenario not in scenarios:
            raise ValueError(f"Baseline scenario '{baseline_scenario}' not found in scenarios")
        
        # Get baseline loads
        if load_type == "cooling":
            baseline_loads = scenarios[baseline_scenario].get("cooling_loads", {})
            components = ["walls", "roofs", "floors", "windows_conduction", "windows_solar", 
                         "doors", "infiltration_sensible", "infiltration_latent", 
                         "people_sensible", "people_latent", "lights", "equipment_sensible", "equipment_latent"]
        else:  # heating
            baseline_loads = scenarios[baseline_scenario].get("heating_loads", {})
            components = ["walls", "roofs", "floors", "windows", "doors", 
                         "infiltration_sensible", "infiltration_latent", 
                         "ventilation_sensible", "ventilation_latent"]
        
        baseline_total = baseline_loads.get("total", 0)
        
        # Initialize data
        scenario_names = []
        percentage_diffs = []
        component_diffs = {}
        
        # Process scenarios (excluding baseline)
        for scenario_name, scenario_data in scenarios.items():
            if scenario_name == baseline_scenario:
                continue
            
            # Extract loads based on load type
            if load_type == "cooling":
                loads = scenario_data.get("cooling_loads", {})
            else:  # heating
                loads = scenario_data.get("heating_loads", {})
            
            # Add scenario name
            scenario_names.append(scenario_name)
            
            # Calculate percentage difference for total load
            scenario_total = loads.get("total", 0)
            if baseline_total != 0:
                percentage_diff = (scenario_total - baseline_total) / baseline_total * 100
            else:
                percentage_diff = 0
            
            percentage_diffs.append(percentage_diff)
            
            # Calculate percentage differences for components
            for component in components:
                if component not in component_diffs:
                    component_diffs[component] = []
                
                baseline_component = baseline_loads.get(component, 0)
                scenario_component = loads.get(component, 0)
                
                if baseline_component != 0:
                    component_diff = (scenario_component - baseline_component) / baseline_component * 100
                else:
                    component_diff = 0
                
                component_diffs[component].append(component_diff)
        
        # Create figure
        fig = go.Figure()
        
        # Add total percentage difference bars
        fig.add_trace(go.Bar(
            x=scenario_names,
            y=percentage_diffs,
            name="Total Load",
            marker_color="rgba(55, 83, 109, 0.7)",
            opacity=0.7
        ))
        
        # Add component percentage difference bars
        for component, diffs in component_diffs.items():
            # Skip components with zero differences
            if sum([abs(diff) for diff in diffs]) == 0:
                continue
            
            # Format component name for display
            display_name = component.replace("_", " ").title()
            
            fig.add_trace(go.Bar(
                x=scenario_names,
                y=diffs,
                name=display_name,
                visible="legendonly"
            ))
        
        # Update layout
        title = f"{load_type.title()} Load Percentage Difference from {baseline_scenario}"
        y_title = "Percentage Difference (%)"
        
        fig.update_layout(
            title=title,
            xaxis_title="Scenario",
            yaxis_title=y_title,
            barmode="group",
            height=500,
            legend=dict(
                orientation="h",
                yanchor="bottom",
                y=1.02,
                xanchor="right",
                x=1
            )
        )
        
        return fig
    
    @staticmethod
    def create_radar_chart(scenarios: Dict[str, Dict[str, Any]], 
                         metrics: List[str],
                         load_type: str = "cooling") -> go.Figure:
        """
        Create a radar chart comparing multiple metrics across scenarios.
        
        Args:
            scenarios: Dictionary with scenario data
            metrics: List of metrics to compare
            load_type: Type of load to compare ("cooling" or "heating")
            
        Returns:
            Plotly figure with radar chart
        """
        # Initialize data
        data = []
        
        # Process scenarios
        for scenario_name, scenario_data in scenarios.items():
            # Extract loads based on load type
            if load_type == "cooling":
                loads = scenario_data.get("cooling_loads", {})
            else:  # heating
                loads = scenario_data.get("heating_loads", {})
            
            # Extract metric values
            values = []
            for metric in metrics:
                values.append(loads.get(metric, 0))
            
            # Add trace
            data.append(go.Scatterpolar(
                r=values,
                theta=metrics,
                fill="toself",
                name=scenario_name
            ))
        
        # Create figure
        fig = go.Figure(data=data)
        
        # Update layout
        title = f"{load_type.title()} Load Metrics Comparison"
        
        fig.update_layout(
            title=title,
            polar=dict(
                radialaxis=dict(
                    visible=True,
                    range=[0, max([max(trace["r"]) for trace in data]) * 1.1]
                )
            ),
            showlegend=True,
            height=600
        )
        
        return fig
    
    @staticmethod
    def create_parallel_coordinates_chart(scenarios: Dict[str, Dict[str, Any]], 
                                        metrics: List[str],
                                        load_type: str = "cooling") -> go.Figure:
        """
        Create a parallel coordinates chart comparing multiple metrics across scenarios.
        
        Args:
            scenarios: Dictionary with scenario data
            metrics: List of metrics to compare
            load_type: Type of load to compare ("cooling" or "heating")
            
        Returns:
            Plotly figure with parallel coordinates chart
        """
        # Initialize data
        data = []
        
        # Process scenarios
        for scenario_name, scenario_data in scenarios.items():
            # Extract loads based on load type
            if load_type == "cooling":
                loads = scenario_data.get("cooling_loads", {})
            else:  # heating
                loads = scenario_data.get("heating_loads", {})
            
            # Create row
            row = {"Scenario": scenario_name}
            
            # Add metrics
            for metric in metrics:
                # Format metric name for display
                display_name = metric.replace("_", " ").title()
                row[display_name] = loads.get(metric, 0)
            
            # Add to data
            data.append(row)
        
        # Create DataFrame
        df = pd.DataFrame(data)
        
        # Create figure
        fig = px.parallel_coordinates(
            df,
            color="Scenario",
            dimensions=[col for col in df.columns if col != "Scenario"],
            title=f"{load_type.title()} Load Metrics Comparison"
        )
        
        # Update layout
        fig.update_layout(
            height=600
        )
        
        return fig
    
    @staticmethod
    def calculate_scenario_differences(scenarios: Dict[str, Dict[str, Any]], 
                                     baseline_scenario: str,
                                     load_type: str = "cooling") -> pd.DataFrame:
        """
        Calculate differences between scenarios and a baseline scenario.
        
        Args:
            scenarios: Dictionary with scenario data
            baseline_scenario: Name of the baseline scenario
            load_type: Type of load to compare ("cooling" or "heating")
            
        Returns:
            DataFrame with scenario differences
        """
        # Check if baseline scenario exists
        if baseline_scenario not in scenarios:
            raise ValueError(f"Baseline scenario '{baseline_scenario}' not found in scenarios")
        
        # Get baseline loads
        if load_type == "cooling":
            baseline_loads = scenarios[baseline_scenario].get("cooling_loads", {})
            components = ["walls", "roofs", "floors", "windows_conduction", "windows_solar", 
                         "doors", "infiltration_sensible", "infiltration_latent", 
                         "people_sensible", "people_latent", "lights", "equipment_sensible", "equipment_latent"]
        else:  # heating
            baseline_loads = scenarios[baseline_scenario].get("heating_loads", {})
            components = ["walls", "roofs", "floors", "windows", "doors", 
                         "infiltration_sensible", "infiltration_latent", 
                         "ventilation_sensible", "ventilation_latent"]
        
        # Initialize data
        data = []
        
        # Process scenarios (excluding baseline)
        for scenario_name, scenario_data in scenarios.items():
            if scenario_name == baseline_scenario:
                continue
            
            # Extract loads based on load type
            if load_type == "cooling":
                loads = scenario_data.get("cooling_loads", {})
            else:  # heating
                loads = scenario_data.get("heating_loads", {})
            
            # Create row
            row = {"Scenario": scenario_name}
            
            # Add absolute differences
            for component in components:
                baseline_value = baseline_loads.get(component, 0)
                scenario_value = loads.get(component, 0)
                
                # Format component name for display
                display_name = component.replace("_", " ").title()
                
                # Calculate absolute difference
                row[f"{display_name} (W)"] = scenario_value - baseline_value
                
                # Calculate percentage difference
                if baseline_value != 0:
                    row[f"{display_name} (%)"] = (scenario_value - baseline_value) / baseline_value * 100
                else:
                    row[f"{display_name} (%)"] = 0
            
            # Add total difference
            baseline_total = baseline_loads.get("total", 0)
            scenario_total = loads.get("total", 0)
            
            row["Total (W)"] = scenario_total - baseline_total
            
            if baseline_total != 0:
                row["Total (%)"] = (scenario_total - baseline_total) / baseline_total * 100
            else:
                row["Total (%)"] = 0
            
            # Add to data
            data.append(row)
        
        # Create DataFrame
        df = pd.DataFrame(data)
        
        return df