utils/scenario_integration.py

"""
Integration module for scenario comparison functionality in the HVAC calculator.

This module provides UI components and integration code for the scenario comparison
functionality in both the cooling and heating calculators.
"""

import streamlit as st
import pandas as pd
import matplotlib.pyplot as plt
from utils.scenario_manager import ScenarioManager

def add_scenario_ui(calculator_type):
    """
    Add scenario saving UI to the calculator.
    
    Args:
        calculator_type (str): Type of calculator ('cooling' or 'heating')
    
    Returns:
        bool: True if a scenario was saved, False otherwise
    """
    st.subheader("Save Current Calculation as Scenario")
    
    # Check if results exist
    results_key = f"{calculator_type}_results"
    if results_key not in st.session_state or not st.session_state[results_key]:
        st.warning("Please complete a calculation before saving a scenario.")
        return False
    
    # Get form data key
    form_data_key = f"{calculator_type}_form_data"
    
    # Create input fields for scenario name and description
    scenario_name = st.text_input("Scenario Name", f"My {calculator_type.capitalize()} Scenario")
    scenario_description = st.text_area("Scenario Description", "Description of this scenario...")
    
    # Save button
    if st.button("Save Scenario"):
        if not scenario_name:
            st.error("Please provide a name for the scenario.")
            return False
        
        # Initialize scenario manager
        scenario_manager = ScenarioManager()
        
        # Save the scenario
        try:
            path = scenario_manager.save_scenario(
                scenario_name,
                scenario_description,
                calculator_type,
                st.session_state[form_data_key],
                st.session_state[results_key]
            )
            
            st.success(f"Scenario saved successfully: {scenario_name}")
            return True
        except Exception as e:
            st.error(f"Error saving scenario: {str(e)}")
            return False
    
    return False

def scenario_comparison_page():
    """
    Create a page for comparing saved scenarios.
    """
    st.title("Scenario Comparison")
    
    # Initialize scenario manager
    scenario_manager = ScenarioManager()
    
    # Get available scenarios
    all_scenarios = scenario_manager.get_available_scenarios()
    
    if not all_scenarios:
        st.warning("No saved scenarios found. Please save some scenarios first.")
        return
    
    # Filter by calculator type
    calculator_type = st.radio(
        "Calculator Type",
        ["cooling", "heating", "all"],
        format_func=lambda x: x.capitalize() if x != "all" else "All"
    )
    
    filtered_scenarios = all_scenarios
    if calculator_type != "all":
        filtered_scenarios = [s for s in all_scenarios if s["calculator_type"] == calculator_type]
    
    if not filtered_scenarios:
        st.warning(f"No {calculator_type} scenarios found.")
        return
    
    # Create a DataFrame for display
    scenarios_df = pd.DataFrame([
        {
            "Name": s["name"],
            "Description": s["description"],
            "Type": s["calculator_type"].capitalize(),
            "Date": s["timestamp"]
        }
        for s in filtered_scenarios
    ])
    
    st.subheader("Available Scenarios")
    st.dataframe(scenarios_df)
    
    # Select scenarios to compare
    st.subheader("Select Scenarios to Compare")
    st.info("Select at least two scenarios to compare. The first selected scenario will be used as the base scenario.")
    
    # Create a mapping of display names to paths
    scenario_options = {f"{s['name']} ({s['calculator_type'].capitalize()})": s["path"] for s in filtered_scenarios}
    
    # Multi-select for scenarios
    selected_scenario_names = st.multiselect(
        "Select Scenarios",
        options=list(scenario_options.keys())
    )
    
    if len(selected_scenario_names) < 2:
        st.warning("Please select at least two scenarios to compare.")
        return
    
    # Get the paths for selected scenarios
    selected_scenario_paths = [scenario_options[name] for name in selected_scenario_names]
    
    # Compare button
    if st.button("Compare Selected Scenarios"):
        # Compare scenarios
        comparison = scenario_manager.compare_scenarios(selected_scenario_paths)
        
        if "error" in comparison:
            st.error(comparison["error"])
            return
        
        # Generate charts
        charts = scenario_manager.generate_comparison_charts(comparison)
        
        # Display comparison
        scenario_manager.display_comparison_in_streamlit(comparison, charts)
        
        # Add export options
        st.subheader("Export Comparison")
        
        # Export as CSV
        comparison_data = []
        
        # Add total loads
        for load in comparison["total_loads"]:
            row = {
                "Scenario": load["name"],
                "Metric": "Total Load",
                "Value": load["total_load_kw"],
                "Unit": "kW"
            }
            comparison_data.append(row)
            
            row = {
                "Scenario": load["name"],
                "Metric": "Recommended Size",
                "Value": load["recommended_size_kw"],
                "Unit": "kW"
            }
            comparison_data.append(row)
        
        # Add breakdown percentages
        for breakdown in comparison["breakdown"]:
            for category in breakdown:
                if category != "name":
                    row = {
                        "Scenario": breakdown["name"],
                        "Metric": f"{category.capitalize()} Percentage",
                        "Value": breakdown[category],
                        "Unit": "%"
                    }
                    comparison_data.append(row)
        
        # Add differences
        base_scenario = comparison["scenarios"][0]
        for scenario_name, diff in comparison["differences"].items():
            row = {
                "Scenario": scenario_name,
                "Metric": f"Absolute Difference from {base_scenario}",
                "Value": diff["absolute_diff_kw"],
                "Unit": "kW"
            }
            comparison_data.append(row)
            
            row = {
                "Scenario": scenario_name,
                "Metric": f"Percentage Difference from {base_scenario}",
                "Value": diff["percentage_diff"],
                "Unit": "%"
            }
            comparison_data.append(row)
        
        # Create DataFrame
        export_df = pd.DataFrame(comparison_data)
        
        # Convert to CSV
        csv = export_df.to_csv(index=False)
        
        # Create download button
        st.download_button(
            label="Download Comparison as CSV",
            data=csv,
            file_name="scenario_comparison.csv",
            mime="text/csv"
        )

def add_scenario_comparison_to_app(app_module):
    """
    Add scenario comparison functionality to the main app.
    
    Args:
        app_module: The main app module to modify
    """
    # Add scenario comparison page to the app
    app_module.pages["scenario_comparison"] = scenario_comparison_page
    
    # Update the main function to include the scenario comparison page
    original_main = app_module.main
    
    def new_main():
        st.sidebar.title("HVAC Load Calculator")
        
        # Add attribution
        st.sidebar.markdown("""
        **Created by:** Dr. Majed Abuseif
        
        This tool was created to facilitate HVAC calculation and understanding for Deakin University students, but has been enhanced to cover wider aspects to allow professionals and energy and HVAC enthusiasts to use it.
        """)
        
        # Add page selection
        page = st.sidebar.radio(
            "Select Calculator",
            ["Cooling Load Calculator", "Heating Load Calculator", "Scenario Comparison"]
        )
        
        if page == "Cooling Load Calculator":
            app_module.cooling_calculator()
        elif page == "Heating Load Calculator":
            app_module.heating_calculator()
        elif page == "Scenario Comparison":
            scenario_comparison_page()
    
    # Replace the main function
    app_module.main = new_main
    
    return app_module

def add_scenario_ui_to_calculator(calculator_module, calculator_type):
    """
    Add scenario saving UI to a calculator module.
    
    Args:
        calculator_module: The calculator module to modify
        calculator_type (str): Type of calculator ('cooling' or 'heating')
    """
    # Find the results section in the calculator
    if calculator_type == "cooling":
        original_results_form = calculator_module.results_form
        
        def new_results_form(ref_data):
            # Call the original results form
            original_results_form(ref_data)
            
            # Add scenario UI
            st.markdown("---")
            add_scenario_ui(calculator_type)
        
        # Replace the results form
        calculator_module.results_form = new_results_form
    
    elif calculator_type == "heating":
        original_results_form = calculator_module.results_form
        
        def new_results_form(ref_data):
            # Call the original results form
            original_results_form(ref_data)
            
            # Add scenario UI
            st.markdown("---")
            add_scenario_ui(calculator_type)
        
        # Replace the results form
        calculator_module.results_form = new_results_form
    
    return calculator_module