utils/scenario_comparison.py

"""
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 CoolingLoadCalculator
from utils.heating_load import HeatingLoadCalculator


class ScenarioComparisonVisualization:
    """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
            )
        )
        
        # Add zero line
        fig.add_shape(
            type="line",
            x0=-0.5,
            x1=len(scenario_names) - 0.5,
            y0=0,
            y1=0,
            line=dict(
                color="black",
                width=1,
                dash="dash"
            )
        )
        
        return fig
    
    @staticmethod
    def create_radar_chart(scenarios: Dict[str, Dict[str, Any]], load_type: str = "cooling") -> go.Figure:
        """
        Create a radar chart comparing key metrics across scenarios.
        
        Args:
            scenarios: Dictionary with scenario data
            load_type: Type of load to compare ("cooling" or "heating")
            
        Returns:
            Plotly figure with radar chart
        """
        # Define metrics based on load type
        if load_type == "cooling":
            metrics = [
                "total",
                "total_sensible",
                "total_latent",
                "walls",
                "roofs",
                "windows_conduction",
                "windows_solar",
                "infiltration_sensible",
                "people_sensible",
                "lights",
                "equipment_sensible"
            ]
            metric_names = [
                "Total Load",
                "Sensible Load",
                "Latent Load",
                "Walls",
                "Roofs",
                "Windows (Conduction)",
                "Windows (Solar)",
                "Infiltration",
                "People",
                "Lights",
                "Equipment"
            ]
        else:  # heating
            metrics = [
                "total",
                "walls",
                "roofs",
                "floors",
                "windows",
                "doors",
                "infiltration_sensible",
                "ventilation_sensible"
            ]
            metric_names = [
                "Total Load",
                "Walls",
                "Roofs",
                "Floors",
                "Windows",
                "Doors",
                "Infiltration",
                "Ventilation"
            ]
        
        # Initialize figure
        fig = go.Figure()
        
        # 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 = [loads.get(metric, 0) for metric in metrics]
            
            # Add trace
            fig.add_trace(go.Scatterpolar(
                r=values,
                theta=metric_names,
                fill="toself",
                name=scenario_name
            ))
        
        # Update layout
        title = f"{load_type.title()} Load Comparison (Radar Chart)"
        
        fig.update_layout(
            title=title,
            polar=dict(
                radialaxis=dict(
                    visible=True,
                    range=[0, max([max([scenarios[s].get(f"{load_type}_loads", {}).get(m, 0) for m in metrics]) for s in scenarios]) * 1.1]
                )
            ),
            height=600,
            showlegend=True
        )
        
        return fig
    
    @staticmethod
    def create_parallel_coordinates_chart(scenarios: Dict[str, Dict[str, Any]]) -> go.Figure:
        """
        Create a parallel coordinates chart comparing scenarios.
        
        Args:
            scenarios: Dictionary with scenario data
            
        Returns:
            Plotly figure with parallel coordinates chart
        """
        # 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 data point
            point = {
                "Scenario": scenario_name,
                "Cooling Load (W)": cooling_loads.get("total", 0),
                "Heating Load (W)": heating_loads.get("total", 0),
                "Sensible Heat Ratio": cooling_loads.get("sensible_heat_ratio", 0),
                "Walls (Cooling)": cooling_loads.get("walls", 0),
                "Windows (Cooling)": cooling_loads.get("windows_conduction", 0) + cooling_loads.get("windows_solar", 0),
                "Internal Gains (Cooling)": cooling_loads.get("people_sensible", 0) + cooling_loads.get("lights", 0) + cooling_loads.get("equipment_sensible", 0),
                "Walls (Heating)": heating_loads.get("walls", 0),
                "Windows (Heating)": heating_loads.get("windows", 0),
                "Infiltration (Heating)": heating_loads.get("infiltration_sensible", 0)
            }
            
            # Add to data
            data.append(point)
        
        # Create DataFrame
        df = pd.DataFrame(data)
        
        # Create figure
        fig = px.parallel_coordinates(
            df,
            color="Cooling Load (W)",
            labels={
                "Scenario": "Scenario",
                "Cooling Load (W)": "Cooling Load (W)",
                "Heating Load (W)": "Heating Load (W)",
                "Sensible Heat Ratio": "Sensible Heat Ratio",
                "Walls (Cooling)": "Walls (Cooling)",
                "Windows (Cooling)": "Windows (Cooling)",
                "Internal Gains (Cooling)": "Internal Gains (Cooling)",
                "Walls (Heating)": "Walls (Heating)",
                "Windows (Heating)": "Windows (Heating)",
                "Infiltration (Heating)": "Infiltration (Heating)"
            },
            color_continuous_scale=px.colors.sequential.Viridis
        )
        
        # Update layout
        fig.update_layout(
            title="Scenario Comparison (Parallel Coordinates)",
            height=600
        )
        
        return fig
    
    @staticmethod
    def display_scenario_comparison(scenarios: Dict[str, Dict[str, Any]]) -> None:
        """
        Display scenario comparison visualization in Streamlit.
        
        Args:
            scenarios: Dictionary with scenario data
        """
        st.header("Scenario Comparison Visualization")
        
        # Check if scenarios exist
        if not scenarios:
            st.warning("No scenarios available for comparison.")
            return
        
        # Create tabs for different visualizations
        tab1, tab2, tab3, tab4, tab5 = st.tabs([
            "Scenario Summary", 
            "Load Comparison", 
            "Percentage Difference", 
            "Radar Chart",
            "Parallel Coordinates"
        ])
        
        with tab1:
            st.subheader("Scenario Summary")
            df = ScenarioComparisonVisualization.create_scenario_summary_table(scenarios)
            st.dataframe(df, use_container_width=True)
            
            # Add download button for CSV
            csv = df.to_csv(index=False).encode('utf-8')
            st.download_button(
                label="Download Scenario Summary as CSV",
                data=csv,
                file_name="scenario_summary.csv",
                mime="text/csv"
            )
        
        with tab2:
            st.subheader("Load Comparison")
            
            # Add load type selector
            load_type = st.radio(
                "Select Load Type",
                ["cooling", "heating"],
                horizontal=True,
                key="load_comparison_type"
            )
            
            # Create and display chart
            fig = ScenarioComparisonVisualization.create_load_comparison_chart(scenarios, load_type)
            st.plotly_chart(fig, use_container_width=True)
        
        with tab3:
            st.subheader("Percentage Difference")
            
            # Add baseline scenario selector
            baseline_scenario = st.selectbox(
                "Select Baseline Scenario",
                list(scenarios.keys()),
                key="baseline_scenario"
            )
            
            # Add load type selector
            load_type = st.radio(
                "Select Load Type",
                ["cooling", "heating"],
                horizontal=True,
                key="percentage_diff_type"
            )
            
            # Create and display chart
            try:
                fig = ScenarioComparisonVisualization.create_percentage_difference_chart(
                    scenarios, baseline_scenario, load_type
                )
                st.plotly_chart(fig, use_container_width=True)
            except ValueError as e:
                st.error(str(e))
        
        with tab4:
            st.subheader("Radar Chart")
            
            # Add load type selector
            load_type = st.radio(
                "Select Load Type",
                ["cooling", "heating"],
                horizontal=True,
                key="radar_chart_type"
            )
            
            # Create and display chart
            fig = ScenarioComparisonVisualization.create_radar_chart(scenarios, load_type)
            st.plotly_chart(fig, use_container_width=True)
        
        with tab5:
            st.subheader("Parallel Coordinates")
            
            # Create and display chart
            fig = ScenarioComparisonVisualization.create_parallel_coordinates_chart(scenarios)
            st.plotly_chart(fig, use_container_width=True)


# Create a singleton instance
scenario_comparison = ScenarioComparisonVisualization()

# Example usage
if __name__ == "__main__":
    import streamlit as st
    
    # Create sample scenarios
    scenarios = {
        "Base Case": {
            "cooling_loads": {
                "total": 5000,
                "total_sensible": 4000,
                "total_latent": 1000,
                "sensible_heat_ratio": 0.8,
                "walls": 1000,
                "roofs": 800,
                "floors": 200,
                "windows_conduction": 500,
                "windows_solar": 800,
                "doors": 100,
                "infiltration_sensible": 300,
                "infiltration_latent": 200,
                "people_sensible": 300,
                "people_latent": 200,
                "lights": 400,
                "equipment_sensible": 400,
                "equipment_latent": 600
            },
            "heating_loads": {
                "total": 6000,
                "walls": 1500,
                "roofs": 1000,
                "floors": 500,
                "windows": 1200,
                "doors": 200,
                "infiltration_sensible": 800,
                "infiltration_latent": 0,
                "ventilation_sensible": 800,
                "ventilation_latent": 0,
                "internal_gains_offset": 1000
            }
        },
        "Improved Insulation": {
            "cooling_loads": {
                "total": 4200,
                "total_sensible": 3500,
                "total_latent": 700,
                "sensible_heat_ratio": 0.83,
                "walls": 600,
                "roofs": 500,
                "floors": 150,
                "windows_conduction": 500,
                "windows_solar": 800,
                "doors": 100,
                "infiltration_sensible": 300,
                "infiltration_latent": 200,
                "people_sensible": 300,
                "people_latent": 200,
                "lights": 400,
                "equipment_sensible": 400,
                "equipment_latent": 300
            },
            "heating_loads": {
                "total": 4500,
                "walls": 900,
                "roofs": 600,
                "floors": 300,
                "windows": 1200,
                "doors": 200,
                "infiltration_sensible": 800,
                "infiltration_latent": 0,
                "ventilation_sensible": 800,
                "ventilation_latent": 0,
                "internal_gains_offset": 1000
            }
        },
        "Better Windows": {
            "cooling_loads": {
                "total": 4000,
                "total_sensible": 3300,
                "total_latent": 700,
                "sensible_heat_ratio": 0.83,
                "walls": 1000,
                "roofs": 800,
                "floors": 200,
                "windows_conduction": 250,
                "windows_solar": 400,
                "doors": 100,
                "infiltration_sensible": 300,
                "infiltration_latent": 200,
                "people_sensible": 300,
                "people_latent": 200,
                "lights": 400,
                "equipment_sensible": 400,
                "equipment_latent": 300
            },
            "heating_loads": {
                "total": 5000,
                "walls": 1500,
                "roofs": 1000,
                "floors": 500,
                "windows": 600,
                "doors": 200,
                "infiltration_sensible": 800,
                "infiltration_latent": 0,
                "ventilation_sensible": 800,
                "ventilation_latent": 0,
                "internal_gains_offset": 1000
            }
        }
    }
    
    # Display scenario comparison
    scenario_comparison.display_scenario_comparison(scenarios)