utils/scenario_manager.py

"""
Utility module for saving and comparing calculation scenarios.

This module provides functionality for saving calculation results as scenarios,
loading saved scenarios, and comparing multiple scenarios to analyze differences.
"""

import os
import json
import pandas as pd
import matplotlib.pyplot as plt
import streamlit as st
from datetime import datetime

class ScenarioManager:
    """
    Manager for saving, loading, and comparing calculation scenarios.
    """
    
    def __init__(self, base_path="scenarios"):
        """
        Initialize the scenario manager.
        
        Args:
            base_path (str): Base directory for storing scenarios
        """
        self.base_path = base_path
        os.makedirs(base_path, exist_ok=True)
    
    def save_scenario(self, name, description, calculator_type, form_data, results):
        """
        Save a calculation scenario.
        
        Args:
            name (str): Name of the scenario
            description (str): Description of the scenario
            calculator_type (str): Type of calculator ('cooling' or 'heating')
            form_data (dict): Form data used for the calculation
            results (dict): Calculation results
            
        Returns:
            str: Path to the saved scenario file
        """
        # Create a timestamp
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        # Create a safe filename
        safe_name = name.replace(" ", "_").lower()
        filename = f"{safe_name}_{timestamp}.json"
        
        # Create the full path
        calculator_dir = os.path.join(self.base_path, calculator_type)
        os.makedirs(calculator_dir, exist_ok=True)
        full_path = os.path.join(calculator_dir, filename)
        
        # Create the scenario data
        scenario_data = {
            "name": name,
            "description": description,
            "calculator_type": calculator_type,
            "timestamp": timestamp,
            "form_data": form_data,
            "results": results
        }
        
        # Save the scenario
        with open(full_path, "w") as f:
            json.dump(scenario_data, f, indent=2)
        
        return full_path
    
    def load_scenario(self, path):
        """
        Load a saved scenario.
        
        Args:
            path (str): Path to the scenario file
            
        Returns:
            dict: Scenario data or None if loading fails
        """
        try:
            with open(path, "r") as f:
                scenario_data = json.load(f)
            return scenario_data
        except Exception as e:
            print(f"Error loading scenario: {e}")
            return None
    
    def get_available_scenarios(self, calculator_type=None):
        """
        Get a list of available scenarios.
        
        Args:
            calculator_type (str, optional): Filter by calculator type
            
        Returns:
            list: List of scenario information dictionaries
        """
        scenarios = []
        
        # Determine which directories to search
        if calculator_type:
            dirs_to_search = [os.path.join(self.base_path, calculator_type)]
        else:
            dirs_to_search = [os.path.join(self.base_path, d) for d in ["cooling", "heating"]]
        
        # Search for scenario files
        for directory in dirs_to_search:
            if not os.path.exists(directory):
                continue
                
            for filename in os.listdir(directory):
                if filename.endswith(".json"):
                    path = os.path.join(directory, filename)
                    scenario = self.load_scenario(path)
                    if scenario:
                        scenarios.append({
                            "path": path,
                            "name": scenario["name"],
                            "description": scenario["description"],
                            "calculator_type": scenario["calculator_type"],
                            "timestamp": scenario["timestamp"]
                        })
        
        # Sort by timestamp (newest first)
        scenarios.sort(key=lambda x: x["timestamp"], reverse=True)
        
        return scenarios
    
    def compare_scenarios(self, scenario_paths):
        """
        Compare multiple scenarios.
        
        Args:
            scenario_paths (list): List of paths to scenario files
            
        Returns:
            dict: Comparison results
        """
        if not scenario_paths or len(scenario_paths) < 2:
            return {"error": "At least two scenarios are required for comparison"}
        
        # Load scenarios
        scenarios = []
        for path in scenario_paths:
            scenario = self.load_scenario(path)
            if scenario:
                scenarios.append(scenario)
        
        # Check if all scenarios are of the same type
        calculator_types = set(s["calculator_type"] for s in scenarios)
        if len(calculator_types) > 1:
            return {"error": "Cannot compare scenarios of different calculator types"}
        
        calculator_type = scenarios[0]["calculator_type"]
        
        # Prepare comparison data
        comparison = {
            "calculator_type": calculator_type,
            "scenarios": [s["name"] for s in scenarios],
            "total_loads": [],
            "breakdown": [],
            "differences": {}
        }
        
        # Extract key metrics for comparison
        for scenario in scenarios:
            results = scenario["results"]
            
            # Add total load
            comparison["total_loads"].append({
                "name": scenario["name"],
                "total_load_kw": results["total_load_kw"],
                "recommended_size_kw": results["recommended_size_kw"]
            })
            
            # Add breakdown percentages
            breakdown = {
                "name": scenario["name"]
            }
            
            if calculator_type == "cooling":
                breakdown.update({
                    "transmission": results["breakdown_percentage"]["transmission"],
                    "solar": results["breakdown_percentage"]["solar"],
                    "ventilation": results["breakdown_percentage"]["ventilation"],
                    "internal": results["breakdown_percentage"]["internal"]
                })
            else:  # heating
                breakdown.update({
                    "transmission": results["breakdown_percentage"]["transmission"],
                    "ventilation": results["breakdown_percentage"]["ventilation"]
                })
            
            comparison["breakdown"].append(breakdown)
        
        # Calculate differences between scenarios
        base_scenario = scenarios[0]
        base_load = base_scenario["results"]["total_load_kw"]
        
        for i, scenario in enumerate(scenarios[1:], 1):
            scenario_load = scenario["results"]["total_load_kw"]
            absolute_diff = scenario_load - base_load
            percentage_diff = (absolute_diff / base_load) * 100 if base_load > 0 else 0
            
            comparison["differences"][scenario["name"]] = {
                "absolute_diff_kw": absolute_diff,
                "percentage_diff": percentage_diff
            }
        
        return comparison
    
    def generate_comparison_charts(self, comparison):
        """
        Generate charts for scenario comparison.
        
        Args:
            comparison (dict): Comparison data from compare_scenarios
            
        Returns:
            dict: Dictionary of matplotlib figures
        """
        if "error" in comparison:
            return {"error": comparison["error"]}
        
        charts = {}
        
        # Total load comparison chart
        fig_total, ax_total = plt.subplots(figsize=(10, 6))
        scenario_names = [load["name"] for load in comparison["total_loads"]]
        total_loads = [load["total_load_kw"] for load in comparison["total_loads"]]
        recommended_sizes = [load["recommended_size_kw"] for load in comparison["total_loads"]]
        
        x = range(len(scenario_names))
        bar_width = 0.35
        
        ax_total.bar([i - bar_width/2 for i in x], total_loads, bar_width, label='Total Load (kW)')
        ax_total.bar([i + bar_width/2 for i in x], recommended_sizes, bar_width, label='Recommended Size (kW)')
        
        ax_total.set_xlabel('Scenario')
        ax_total.set_ylabel('Load (kW)')
        ax_total.set_title('Total Load Comparison')
        ax_total.set_xticks(x)
        ax_total.set_xticklabels(scenario_names, rotation=45, ha='right')
        ax_total.legend()
        
        plt.tight_layout()
        charts["total_load"] = fig_total
        
        # Breakdown comparison chart
        fig_breakdown, ax_breakdown = plt.subplots(figsize=(12, 6))
        
        # Determine categories based on calculator type
        if comparison["calculator_type"] == "cooling":
            categories = ["transmission", "solar", "ventilation", "internal"]
            category_labels = ["Transmission", "Solar", "Ventilation", "Internal"]
        else:  # heating
            categories = ["transmission", "ventilation"]
            category_labels = ["Transmission", "Ventilation"]
        
        # Prepare data for grouped bar chart
        bar_positions = []
        bar_heights = []
        bar_labels = []
        bar_colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd']
        
        for i, category in enumerate(categories):
            positions = [i + j * (len(categories) + 1) for j in range(len(comparison["breakdown"]))]
            bar_positions.extend(positions)
            
            heights = [breakdown[category] for breakdown in comparison["breakdown"]]
            bar_heights.extend(heights)
            
            for scenario_name in [breakdown["name"] for breakdown in comparison["breakdown"]]:
                bar_labels.append(scenario_name)
        
        # Create the grouped bar chart
        bars = ax_breakdown.bar(bar_positions, bar_heights, width=0.8)
        
        # Color the bars by scenario
        for i, bar in enumerate(bars):
            scenario_index = i % len(comparison["breakdown"])
            bar.set_color(bar_colors[scenario_index % len(bar_colors)])
        
        # Set labels and title
        ax_breakdown.set_xlabel('Category')
        ax_breakdown.set_ylabel('Percentage (%)')
        ax_breakdown.set_title('Load Breakdown Comparison')
        
        # Set x-ticks at the center of each group
        group_centers = [i + (len(comparison["breakdown"]) - 1) / 2 for i in range(0, len(bar_positions), len(comparison["breakdown"]))]
        ax_breakdown.set_xticks(group_centers)
        ax_breakdown.set_xticklabels(category_labels)
        
        # Add a legend
        scenario_names = [breakdown["name"] for breakdown in comparison["breakdown"]]
        legend_handles = [plt.Rectangle((0, 0), 1, 1, color=bar_colors[i % len(bar_colors)]) for i in range(len(scenario_names))]
        ax_breakdown.legend(legend_handles, scenario_names, loc='upper right')
        
        plt.tight_layout()
        charts["breakdown"] = fig_breakdown
        
        # Differences chart (if there are differences)
        if comparison["differences"]:
            fig_diff, ax_diff = plt.subplots(figsize=(10, 6))
            
            scenario_names = list(comparison["differences"].keys())
            absolute_diffs = [diff["absolute_diff_kw"] for diff in comparison["differences"].values()]
            percentage_diffs = [diff["percentage_diff"] for diff in comparison["differences"].values()]
            
            x = range(len(scenario_names))
            
            # Create two y-axes
            ax_abs = ax_diff
            ax_pct = ax_abs.twinx()
            
            # Plot data
            bars = ax_abs.bar(x, absolute_diffs, width=0.6, color='#1f77b4', alpha=0.7, label='Absolute Difference (kW)')
            line = ax_pct.plot(x, percentage_diffs, 'ro-', label='Percentage Difference (%)')
            
            # Add labels and title
            ax_abs.set_xlabel('Scenario')
            ax_abs.set_ylabel('Absolute Difference (kW)')
            ax_pct.set_ylabel('Percentage Difference (%)')
            ax_abs.set_title(f'Differences Compared to Base Scenario ({comparison["scenarios"][0]})')
            
            # Set x-ticks
            ax_abs.set_xticks(x)
            ax_abs.set_xticklabels(scenario_names, rotation=45, ha='right')
            
            # Add legends
            lines, labels = ax_abs.get_legend_handles_labels()
            lines2, labels2 = ax_pct.get_legend_handles_labels()
            ax_abs.legend(lines + lines2, labels + labels2, loc='upper left')
            
            plt.tight_layout()
            charts["differences"] = fig_diff
        
        return charts
    
    def display_comparison_in_streamlit(self, comparison, charts=None):
        """
        Display scenario comparison in Streamlit.
        
        Args:
            comparison (dict): Comparison data from compare_scenarios
            charts (dict, optional): Charts from generate_comparison_charts
        """
        if "error" in comparison:
            st.error(comparison["error"])
            return
        
        # Display total loads
        st.subheader("Total Load Comparison")
        
        # Create a DataFrame for the total loads
        total_loads_df = pd.DataFrame(comparison["total_loads"])
        total_loads_df = total_loads_df.rename(columns={
            "name": "Scenario",
            "total_load_kw": "Total Load (kW)",
            "recommended_size_kw": "Recommended Size (kW)"
        })
        
        st.dataframe(total_loads_df)
        
        # Display the total load chart
        if charts and "total_load" in charts:
            st.pyplot(charts["total_load"])
        
        # Display breakdown
        st.subheader("Load Breakdown Comparison")
        
        # Create a DataFrame for the breakdown
        breakdown_df = pd.DataFrame(comparison["breakdown"])
        
        # Rename columns for better display
        column_mapping = {
            "name": "Scenario",
            "transmission": "Transmission (%)",
            "solar": "Solar (%)",
            "ventilation": "Ventilation (%)",
            "internal": "Internal (%)"
        }
        
        breakdown_df = breakdown_df.rename(columns={k: v for k, v in column_mapping.items() if k in breakdown_df.columns})
        
        st.dataframe(breakdown_df)
        
        # Display the breakdown chart
        if charts and "breakdown" in charts:
            st.pyplot(charts["breakdown"])
        
        # Display differences
        if comparison["differences"]:
            st.subheader(f"Differences Compared to Base Scenario ({comparison['scenarios'][0]})")
            
            # Create a DataFrame for the differences
            diff_data = []
            for scenario_name, diff in comparison["differences"].items():
                diff_data.append({
                    "Scenario": scenario_name,
                    "Absolute Difference (kW)": diff["absolute_diff_kw"],
                    "Percentage Difference (%)": diff["percentage_diff"]
                })
            
            diff_df = pd.DataFrame(diff_data)
            st.dataframe(diff_df)
            
            # Display the differences chart
            if charts and "differences" in charts:
                st.pyplot(charts["differences"])
        
        # Display interpretation
        st.subheader("Interpretation")
        
        base_scenario = comparison["scenarios"][0]
        
        if comparison["calculator_type"] == "cooling":
            st.write(f"""
            ### Key Observations:
            
            - The base scenario ({base_scenario}) has a total cooling load of {comparison['total_loads'][0]['total_load_kw']:.2f} kW.
            - The recommended cooling system size for the base scenario is {comparison['total_loads'][0]['recommended_size_kw']:.2f} kW.
            """)
            
            if comparison["differences"]:
                for scenario_name, diff in comparison["differences"].items():
                    if diff["absolute_diff_kw"] > 0:
                        st.write(f"- {scenario_name} has a **higher** cooling load by {abs(diff['absolute_diff_kw']):.2f} kW ({abs(diff['percentage_diff']):.1f}%) compared to the base scenario.")
                    else:
                        st.write(f"- {scenario_name} has a **lower** cooling load by {abs(diff['absolute_diff_kw']):.2f} kW ({abs(diff['percentage_diff']):.1f}%) compared to the base scenario.")
        else:  # heating
            st.write(f"""
            ### Key Observations:
            
            - The base scenario ({base_scenario}) has a total heating load of {comparison['total_loads'][0]['total_load_kw']:.2f} kW.
            - The recommended heating system size for the base scenario is {comparison['total_loads'][0]['recommended_size_kw']:.2f} kW.
            """)
            
            if comparison["differences"]:
                for scenario_name, diff in comparison["differences"].items():
                    if diff["absolute_diff_kw"] > 0:
                        st.write(f"- {scenario_name} has a **higher** heating load by {abs(diff['absolute_diff_kw']):.2f} kW ({abs(diff['percentage_diff']):.1f}%) compared to the base scenario.")
                    else:
                        st.write(f"- {scenario_name} has a **lower** heating load by {abs(diff['absolute_diff_kw']):.2f} kW ({abs(diff['percentage_diff']):.1f}%) compared to the base scenario.")