app/main.py

"""
HVAC Calculator Code Documentation
"""

import streamlit as st
import pandas as pd
import numpy as np
import plotly.express as px
import json
import pycountry
import os
import sys
from typing import Dict, List, Any, Optional, Tuple

# Import application modules
from app.building_info_form import BuildingInfoForm
from app.component_selection import ComponentSelectionInterface, Orientation, ComponentType, Wall, Roof, Floor, Window, Door
from app.results_display import ResultsDisplay
from app.data_validation import DataValidation
from app.data_persistence import DataPersistence
from app.data_export import DataExport

# Import data modules
from data.reference_data import ReferenceData
from data.climate_data import ClimateData, ClimateLocation
from data.ashrae_tables import ASHRAETables
from data.building_components import Wall as WallModel, Roof as RoofModel

# Import utility modules
from utils.u_value_calculator import UValueCalculator
from utils.shading_system import ShadingSystem
from utils.area_calculation_system import AreaCalculationSystem
from utils.psychrometrics import Psychrometrics
from utils.heat_transfer import HeatTransferCalculations
from utils.cooling_load import CoolingLoadCalculator
from utils.heating_load import HeatingLoadCalculator
from utils.component_visualization import ComponentVisualization
from utils.scenario_comparison import ScenarioComparisonVisualization
from utils.psychrometric_visualization import PsychrometricVisualization
from utils.time_based_visualization import TimeBasedVisualization

class HVACCalculator:
    def __init__(self):
        st.set_page_config(
            page_title="HVAC Load Calculator",
            page_icon="🌡️",
            layout="wide",
            initial_sidebar_state="expanded"
        )
        
        # Initialize session state
        if 'page' not in st.session_state:
            st.session_state.page = 'Building Information'
        
        if 'building_info' not in st.session_state:
            st.session_state.building_info = {"project_name": ""}
        
        if 'components' not in st.session_state:
            st.session_state.components = {
                'walls': [],
                'roofs': [],
                'floors': [],
                'windows': [],
                'doors': []
            }
        
        if 'internal_loads' not in st.session_state:
            st.session_state.internal_loads = {
                'people': [],
                'lighting': [],
                'equipment': []
            }
        
        if 'calculation_results' not in st.session_state:
            st.session_state.calculation_results = {
                'cooling': {},
                'heating': {}
            }
        
        if 'saved_scenarios' not in st.session_state:
            st.session_state.saved_scenarios = {}
        
        if 'climate_data' not in st.session_state:
            st.session_state.climate_data = {}
        
        # Initialize modules
        self.building_info_form = BuildingInfoForm()
        self.component_selection = ComponentSelectionInterface()
        self.results_display = ResultsDisplay()
        self.data_validation = DataValidation()
        self.data_persistence = DataPersistence()
        self.data_export = DataExport()
        self.cooling_calculator = CoolingLoadCalculator()
        self.heating_calculator = HeatingLoadCalculator()
        self.climate_data = ClimateData()
        
        self.setup_layout()
    
    def setup_layout(self):
        st.sidebar.title("HVAC Load Calculator")
        st.sidebar.markdown("---")
        
        st.sidebar.subheader("Navigation")
        pages = [
            "Building Information",
            "Climate Data",
            "Building Components",
            "Internal Loads",
            "Calculation Results",
            "Export Data"
        ]
        
        selected_page = st.sidebar.radio("Go to", pages, index=pages.index(st.session_state.page))
        
        if selected_page != st.session_state.page:
            st.session_state.page = selected_page
        
        self.display_page(st.session_state.page)
        
        st.sidebar.markdown("---")
        st.sidebar.info(
            "HVAC Load Calculator v1.0.1\n\n"
            "Based on ASHRAE steady-state calculation methods\n\n"
            "Developed by: Dr Majed Abuseif\n\n"
            "School of Architecture and Built Environment\n\n"
            "Deakin University\n\n"
            "© 2025"
        )
    
    def display_page(self, page: str):
        if page == "Building Information":
            self.building_info_form.display_building_info_form(st.session_state)
        elif page == "Climate Data":
            self.climate_data.display_climate_input(st.session_state)
        elif page == "Building Components":
            self.component_selection.display_component_selection(st.session_state)
        elif page == "Internal Loads":
            self.display_internal_loads()
        elif page == "Calculation Results":
            self.display_calculation_results()
        elif page == "Export Data":
            self.data_export.display()
    
    def validate_internal_load(self, load_type: str, new_load: Dict) -> Tuple[bool, str]:
        """Validate if a new internal load is unique and within limits."""
        loads = st.session_state.internal_loads.get(load_type, [])
        max_loads = 50
        
        if len(loads) >= max_loads:
            return False, f"Maximum of {max_loads} {load_type} loads reached."
        
        # Check for duplicates based on key attributes
        for existing_load in loads:
            if load_type == 'people':
                if (existing_load['name'] == new_load['name'] and
                    existing_load['num_people'] == new_load['num_people'] and
                    existing_load['activity_level'] == new_load['activity_level'] and
                    existing_load['zone_type'] == new_load['zone_type'] and
                    existing_load['hours_in_operation'] == new_load['hours_in_operation']):
                    return False, f"Duplicate people load '{new_load['name']}' already exists."
            elif load_type == 'lighting':
                if (existing_load['name'] == new_load['name'] and
                    existing_load['power'] == new_load['power'] and
                    existing_load['usage_factor'] == new_load['usage_factor'] and
                    existing_load['zone_type'] == new_load['zone_type'] and
                    existing_load['hours_in_operation'] == new_load['hours_in_operation']):
                    return False, f"Duplicate lighting load '{new_load['name']}' already exists."
            elif load_type == 'equipment':
                if (existing_load['name'] == new_load['name'] and
                    existing_load['power'] == new_load['power'] and
                    existing_load['usage_factor'] == new_load['usage_factor'] and
                    existing_load['radiation_fraction'] == new_load['radiation_fraction'] and
                    existing_load['zone_type'] == new_load['zone_type'] and
                    existing_load['hours_in_operation'] == new_load['hours_in_operation']):
                    return False, f"Duplicate equipment load '{new_load['name']}' already exists."
        
        return True, "Valid load."
    
    def display_internal_loads(self):
        st.title("Internal Loads")
        
        # Reset button for all internal loads
        if st.button("Reset All Internal Loads"):
            st.session_state.internal_loads = {'people': [], 'lighting': [], 'equipment': []}
            st.success("All internal loads reset!")
            st.rerun()
        
        tabs = st.tabs(["People", "Lighting", "Equipment"])
        
        with tabs[0]:
            st.subheader("People")
            with st.form("people_form"):
                num_people = st.number_input("Number of People", min_value=0, value=0, step=1)
                activity_level = st.selectbox(
                    "Activity Level",
                    ["Seated/Resting", "Light Work", "Moderate Work", "Heavy Work"]
                )
                zone_type = st.selectbox("Zone Type", ["Office", "Classroom", "Retail", "Residential"])
                hours_in_operation = st.number_input(
                    "Hours in Operation",
                    min_value=0.0,
                    max_value=24.0,
                    value=8.0,
                    step=0.5
                )
                people_name = st.text_input("Name", value="Occupants")
                
                if st.form_submit_button("Add People Load"):
                    people_load = {
                        "id": f"people_{len(st.session_state.internal_loads['people'])}",
                        "name": people_name,
                        "num_people": num_people,
                        "activity_level": activity_level,
                        "zone_type": zone_type,
                        "hours_in_operation": hours_in_operation
                    }
                    is_valid, message = self.validate_internal_load('people', people_load)
                    if is_valid:
                        st.session_state.internal_loads['people'].append(people_load)
                        st.success("People load added!")
                        st.rerun()
                    else:
                        st.error(message)
            
            if st.session_state.internal_loads['people']:
                people_df = pd.DataFrame(st.session_state.internal_loads['people'])
                st.dataframe(people_df, use_container_width=True)
                
                selected_people = st.multiselect(
                    "Select People Loads to Delete",
                    [load['id'] for load in st.session_state.internal_loads['people']]
                )
                if st.button("Delete Selected People Loads"):
                    st.session_state.internal_loads['people'] = [
                        load for load in st.session_state.internal_loads['people']
                        if load['id'] not in selected_people
                    ]
                    st.success("Selected people loads deleted!")
                    st.rerun()
        
        with tabs[1]:
            st.subheader("Lighting")
            with st.form("lighting_form"):
                power = st.number_input("Power (W)", min_value=0.0, value=1000.0, step=100.0)
                usage_factor = st.number_input(
                    "Usage Factor",
                    min_value=0.0,
                    max_value=1.0,
                    value=0.8,
                    step=0.1
                )
                zone_type = st.selectbox("Zone Type", ["Office", "Classroom", "Retail", "Residential"])
                hours_in_operation = st.number_input(
                    "Hours in Operation",
                    min_value=0.0,
                    max_value=24.0,
                    value=8.0,
                    step=0.5
                )
                lighting_name = st.text_input("Name", value="General Lighting")
                
                if st.form_submit_button("Add Lighting Load"):
                    lighting_load = {
                        "id": f"lighting_{len(st.session_state.internal_loads['lighting'])}",
                        "name": lighting_name,
                        "power": power,
                        "usage_factor": usage_factor,
                        "zone_type": zone_type,
                        "hours_in_operation": hours_in_operation
                    }
                    is_valid, message = self.validate_internal_load('lighting', lighting_load)
                    if is_valid:
                        st.session_state.internal_loads['lighting'].append(lighting_load)
                        st.success("Lighting load added!")
                        st.rerun()
                    else:
                        st.error(message)
            
            if st.session_state.internal_loads['lighting']:
                lighting_df = pd.DataFrame(st.session_state.internal_loads['lighting'])
                st.dataframe(lighting_df, use_container_width=True)
                
                selected_lighting = st.multiselect(
                    "Select Lighting Loads to Delete",
                    [load['id'] for load in st.session_state.internal_loads['lighting']]
                )
                if st.button("Delete Selected Lighting Loads"):
                    st.session_state.internal_loads['lighting'] = [
                        load for load in st.session_state.internal_loads['lighting']
                        if load['id'] not in selected_lighting
                    ]
                    st.success("Selected lighting loads deleted!")
                    st.rerun()
        
        with tabs[2]:
            st.subheader("Equipment")
            with st.form("equipment_form"):
                power = st.number_input("Power (W)", min_value=0.0, value=500.0, step=100.0)
                usage_factor = st.number_input(
                    "Usage Factor",
                    min_value=0.0,
                    max_value=1.0,
                    value=0.7,
                    step=0.1
                )
                radiation_fraction = st.number_input(
                    "Radiation Fraction",
                    min_value=0.0,
                    max_value=1.0,
                    value=0.3,
                    step=0.1
                )
                zone_type = st.selectbox("Zone Type", ["Office", "Classroom", "Retail", "Residential"])
                hours_in_operation = st.number_input(
                    "Hours in Operation",
                    min_value=0.0,
                    max_value=24.0,
                    value=8.0,
                    step=0.5
                )
                equipment_name = st.text_input("Name", value="Office Equipment")
                
                if st.form_submit_button("Add Equipment Load"):
                    equipment_load = {
                        "id": f"equipment_{len(st.session_state.internal_loads['equipment'])}",
                        "name": equipment_name,
                        "power": power,
                        "usage_factor": usage_factor,
                        "radiation_fraction": radiation_fraction,
                        "zone_type": zone_type,
                        "hours_in_operation": hours_in_operation
                    }
                    is_valid, message = self.validate_internal_load('equipment', equipment_load)
                    if is_valid:
                        st.session_state.internal_loads['equipment'].append(equipment_load)
                        st.success("Equipment load added!")
                        st.rerun()
                    else:
                        st.error(message)
            
            if st.session_state.internal_loads['equipment']:
                equipment_df = pd.DataFrame(st.session_state.internal_loads['equipment'])
                st.dataframe(equipment_df, use_container_width=True)
                
                selected_equipment = st.multiselect(
                    "Select Equipment Loads to Delete",
                    [load['id'] for load in st.session_state.internal_loads['equipment']]
                )
                if st.button("Delete Selected Equipment Loads"):  # Fixed typo
                    st.session_state.internal_loads['equipment'] = [
                        load for load in st.session_state.internal_loads['equipment']
                        if load['id'] not in selected_equipment
                    ]
                    st.success("Selected equipment loads deleted!")
                    st.rerun()
        
        col1, col2 = st.columns(2)
        with col1:
            st.button(
                "Back to Building Components",
                on_click=lambda: setattr(st.session_state, "page", "Building Components")
            )
        with col2:
            st.button(
                "Continue to Calculation Results",
                on_click=lambda: setattr(st.session_state, "page", "Calculation Results")
            )
    
    def calculate_cooling(self) -> Tuple[bool, str, Dict]:
        """
        Calculate cooling loads using CoolingLoadCalculator.
        Returns: (success, message, results)
        """
        try:
            # Gather inputs
            building_components = st.session_state.get('components', {})
            internal_loads = st.session_state.get('internal_loads', {})
            building_info = st.session_state.get('building_info', {})
            
            # Validate inputs
            if not building_info.get('floor_area'):
                return False, "Floor area is missing in building information.", {}
            if not building_components.get('walls') and not building_components.get('roofs') and not building_components.get('windows'):
                return False, "No building components defined. Please add walls, roofs, or windows.", {}
            
            # Extract climate data using ClimateData.get_location_by_id
            climate_id = f"{building_info.get('country', 'XX')[:2].upper()}-{building_info.get('city', 'XXX')[:3].upper()}"
            location = self.climate_data.get_location_by_id(climate_id, st.session_state)
            if not location:
                return False, f"No climate data found for {climate_id}. Please provide valid climate data.", {}
            
            # Validate climate data
            if not all(k in location for k in ['summer_design_temp_db', 'summer_design_temp_wb', 'monthly_temps', 'latitude']):
                return False, f"Invalid climate data for {climate_id}. Missing required fields.", {}
            
            # Format conditions
            outdoor_conditions = {
                'temperature': location['summer_design_temp_db'],
                'relative_humidity': location['monthly_humidity'].get('Jul', 50.0),
                'ground_temperature': location['monthly_temps'].get('Jul', 20.0),
                'month': 'Jul',
                'latitude': f"{location['latitude']}N" if location['latitude'] >= 0 else f"{abs(location['latitude'])}S",
                'wind_speed': 4.0,  # Default as not provided in climate data
                'day_of_year': 204  # Approx. July 23
            }
            indoor_conditions = {
                'temperature': building_info.get('indoor_temp', 24.0),
                'relative_humidity': building_info.get('indoor_rh', 50.0)
            }
            
            # Debug: Log inputs
            st.write("Debug: Cooling Input State", {
                'climate_id': climate_id,
                'outdoor_conditions': outdoor_conditions,
                'indoor_conditions': indoor_conditions,
                'components': {k: len(v) for k, v in building_components.items()},
                'internal_loads': {
                    'people': len(internal_loads.get('people', [])),
                    'lighting': len(internal_loads.get('lighting', [])),
                    'equipment': len(internal_loads.get('equipment', []))
                },
                'building_info': building_info
            })
            
            # Format internal loads
            formatted_internal_loads = {
                'people': {
                    'number': sum(load['num_people'] for load in internal_loads.get('people', [])),
                    'activity_level': internal_loads.get('people', [{}])[0].get('activity_level', 'Seated/Resting'),
                    'operating_hours': f"{internal_loads.get('people', [{}])[0].get('hours_in_operation', 8)}:00-{internal_loads.get('people', [{}])[0].get('hours_in_operation', 8)+10}:00"
                },
                'lights': {
                    'power': sum(load['power'] for load in internal_loads.get('lighting', [])),
                    'use_factor': internal_loads.get('lighting', [{}])[0].get('usage_factor', 0.8),
                    'special_allowance': 0.1,
                    'hours_operation': f"{internal_loads.get('lighting', [{}])[0].get('hours_in_operation', 8)}h"
                },
                'equipment': {
                    'power': sum(load['power'] for load in internal_loads.get('equipment', [])),
                    'use_factor': internal_loads.get('equipment', [{}])[0].get('usage_factor', 0.7),
                    'radiation_factor': internal_loads.get('equipment', [{}])[0].get('radiation_fraction', 0.3),
                    'hours_operation': f"{internal_loads.get('equipment', [{}])[0].get('hours_in_operation', 8)}h"
                },
                'infiltration': {
                    'flow_rate': building_info.get('infiltration_rate', 0.05),
                    'height': building_info.get('building_height', 3.0),
                    'crack_length': building_info.get('crack_length', 10.0)
                },
                'ventilation': {
                    'flow_rate': building_info.get('ventilation_rate', 0.1)
                },
                'operating_hours': building_info.get('operating_hours', '8:00-18:00')
            }
            
            # Calculate hourly loads
            hourly_loads = self.cooling_calculator.calculate_hourly_cooling_loads(
                building_components=building_components,
                outdoor_conditions=outdoor_conditions,
                indoor_conditions=indoor_conditions,
                internal_loads=formatted_internal_loads,
                building_volume=building_info.get('floor_area', 100.0) * building_info.get('building_height', 3.0)
            )
            if not hourly_loads:
                return False, "Cooling hourly loads calculation failed. Check input data.", {}
            
            # Get design loads
            design_loads = self.cooling_calculator.calculate_design_cooling_load(hourly_loads)
            if not design_loads:
                return False, "Cooling design loads calculation failed. Check input data.", {}
            
            # Get summary
            summary = self.cooling_calculator.calculate_cooling_load_summary(design_loads)
            if not summary:
                return False, "Cooling load summary calculation failed. Check input data.", {}
            
            # Format results for results_display.py
            floor_area = building_info.get('floor_area', 100.0) or 100.0
            results = {
                'total_load': summary['total'] / 1000,  # kW
                'sensible_load': summary['total_sensible'] / 1000,  # kW
                'latent_load': summary['total_latent'] / 1000,  # kW
                'load_per_area': summary['total'] / floor_area,  # W/m²
                'component_loads': {
                    'walls': design_loads['walls'] / 1000,
                    'roof': design_loads['roofs'] / 1000,
                    'windows': (design_loads['windows_conduction'] + design_loads['windows_solar']) / 1000,
                    'doors': design_loads['doors'] / 1000,
                    'people': (design_loads['people_sensible'] + design_loads['people_latent']) / 1000,
                    'lighting': design_loads['lights'] / 1000,
                    'equipment': (design_loads['equipment_sensible'] + design_loads['equipment_latent']) / 1000,
                    'infiltration': (design_loads['infiltration_sensible'] + design_loads['infiltration_latent']) / 1000,
                    'ventilation': (design_loads['ventilation_sensible'] + design_loads['ventilation_latent']) / 1000
                },
                'detailed_loads': {
                    'walls': [],
                    'roofs': [],
                    'windows': [],
                    'doors': [],
                    'internal': [],
                    'infiltration': {
                        'air_flow': formatted_internal_loads['infiltration']['flow_rate'],
                        'sensible_load': design_loads['infiltration_sensible'] / 1000,
                        'latent_load': design_loads['infiltration_latent'] / 1000,
                        'total_load': (design_loads['infiltration_sensible'] + design_loads['infiltration_latent']) / 1000
                    },
                    'ventilation': {
                        'air_flow': formatted_internal_loads['ventilation']['flow_rate'],
                        'sensible_load': design_loads['ventilation_sensible'] / 1000,
                        'latent_load': design_loads['ventilation_latent'] / 1000,
                        'total_load': (design_loads['ventilation_sensible'] + design_loads['ventilation_latent']) / 1000
                    }
                },
                'building_info': building_info
            }
            
            # Populate detailed loads
            for wall in building_components.get('walls', []):
                load = self.cooling_calculator.calculate_wall_cooling_load(
                    wall=wall,
                    outdoor_temp=outdoor_conditions['temperature'],
                    indoor_temp=indoor_conditions['temperature'],
                    month=outdoor_conditions['month'],
                    hour=design_loads['design_hour'],
                    latitude=outdoor_conditions['latitude']
                )
                results['detailed_loads']['walls'].append({
                    'name': wall.name,
                    'orientation': wall.orientation.value,
                    'area': wall.area,
                    'u_value': wall.u_value,
                    'cltd': self.cooling_calculator.ashrae_tables.calculate_corrected_cltd_wall(
                        wall_group=wall.wall_group,
                        orientation=wall.orientation.value,
                        hour=design_loads['design_hour'],
                        color='Dark',
                        month=outdoor_conditions['month'],
                        latitude=outdoor_conditions['latitude'],
                        indoor_temp=indoor_conditions['temperature'],
                        outdoor_temp=outdoor_conditions['temperature']
                    ),
                    'load': load / 1000
                })
            
            for roof in building_components.get('roofs', []):
                load = self.cooling_calculator.calculate_roof_cooling_load(
                    roof=roof,
                    outdoor_temp=outdoor_conditions['temperature'],
                    indoor_temp=indoor_conditions['temperature'],
                    month=outdoor_conditions['month'],
                    hour=design_loads['design_hour'],
                    latitude=outdoor_conditions['latitude']
                )
                results['detailed_loads']['roofs'].append({
                    'name': roof.name,
                    'orientation': roof.orientation.value,
                    'area': roof.area,
                    'u_value': roof.u_value,
                    'cltd': self.cooling_calculator.ashrae_tables.calculate_corrected_cltd_roof(
                        roof_group=roof.roof_group,
                        hour=design_loads['design_hour'],
                        color='Dark',
                        month=outdoor_conditions['month'],
                        latitude=outdoor_conditions['latitude'],
                        indoor_temp=indoor_conditions['temperature'],
                        outdoor_temp=outdoor_conditions['temperature']
                    ),
                    'load': load / 1000
                })
            
            for window in building_components.get('windows', []):
                load_dict = self.cooling_calculator.calculate_window_cooling_load(
                    window=window,
                    outdoor_temp=outdoor_conditions['temperature'],
                    indoor_temp=indoor_conditions['temperature'],
                    month=outdoor_conditions['month'],
                    hour=design_loads['design_hour'],
                    latitude=outdoor_conditions['latitude'],
                    shading_coefficient=window.shading_coefficient
                )
                scl_latitude = f"{float(outdoor_conditions['latitude'][:-1])}_{outdoor_conditions['month'].upper()}"
                results['detailed_loads']['windows'].append({
                    'name': window.name,
                    'orientation': window.orientation.value,
                    'area': window.area,
                    'u_value': window.u_value,
                    'shgc': window.shgc,
                    'shading_device': window.shading_device,
                    'shading_coefficient': window.shading_coefficient,
                    'scl': self.cooling_calculator.ashrae_tables.get_scl(
                        window.orientation.value,
                        design_loads['design_hour'],
                        scl_latitude
                    ),
                    'load': load_dict['total'] / 1000
                })
            
            for door in building_components.get('doors', []):
                load = self.cooling_calculator.calculate_door_cooling_load(
                    door=door,
                    outdoor_temp=outdoor_conditions['temperature'],
                    indoor_temp=indoor_conditions['temperature']
                )
                results['detailed_loads']['doors'].append({
                    'name': door.name,
                    'orientation': door.orientation.value,
                    'area': door.area,
                    'u_value': door.u_value,
                    'cltd': outdoor_conditions['temperature'] - indoor_conditions['temperature'],
                    'load': load / 1000
                })
            
            for load_type, key in [('people', 'people'), ('lighting', 'lights'), ('equipment', 'equipment')]:
                for load in internal_loads.get(key, []):
                    if load_type == 'people':
                        load_dict = self.cooling_calculator.calculate_people_cooling_load(
                            num_people=load['num_people'],
                            activity_level=load['activity_level'],
                            hour=design_loads['design_hour']
                        )
                    elif load_type == 'lighting':
                        load_dict = {'total': self.cooling_calculator.calculate_lights_cooling_load(
                            power=load['power'],
                            use_factor=load['usage_factor'],
                            special_allowance=0.1,
                            hour=design_loads['design_hour']
                        )}
                    else:
                        load_dict = self.cooling_calculator.calculate_equipment_cooling_load(
                            power=load['power'],
                            use_factor=load['usage_factor'],
                            radiation_factor=load['radiation_fraction'],
                            hour=design_loads['design_hour']
                        )
                    results['detailed_loads']['internal'].append({
                        'type': load_type.capitalize(),
                        'name': load['name'],
                        'quantity': load.get('num_people', load.get('power', 1)),
                        'heat_gain': load_dict.get('sensible', load_dict['total']),
                        'clf': self.cooling_calculator.ashrae_tables.get_clf_people(
                            design_loads['design_hour'],
                            f"{load['hours_in_operation']}h"
                        ) if load_type == 'people' else 1.0,
                        'load': load_dict['total'] / 1000
                    })
            
            # Debug: Log results
            st.write("Debug: Cooling Results", {
                'total_load': results.get('total_load', 'N/A'),
                'component_loads': results.get('component_loads', 'N/A'),
                'detailed_loads': {k: len(v) if isinstance(v, list) else v for k, v in results.get('detailed_loads', {}).items()}
            })
            
            return True, "Cooling calculation completed.", results
        except Exception as e:
            st.error(f"Cooling calculation error: {str(e)}")
            return False, f"Cooling calculation error: {str(e)}", {}
    
    def calculate_heating(self) -> Tuple[bool, str, Dict]:
        """
        Calculate heating loads using HeatingLoadCalculator.
        Returns: (success, message, results)
        """
        try:
            # Gather inputs
            building_components = st.session_state.get('components', {})
            internal_loads = st.session_state.get('internal_loads', {})
            building_info = st.session_state.get('building_info', {})
            
            # Validate inputs
            if not building_info.get('floor_area'):
                return False, "Floor area is missing in building information.", {}
            if not building_components.get('walls') and not building_components.get('roofs') and not building_components.get('windows'):
                return False, "No building components defined. Please add walls, roofs, or windows.", {}
            
            # Extract climate data using ClimateData.get_location_by_id
            climate_id = f"{building_info.get('country', 'XX')[:2].upper()}-{building_info.get('city', 'XXX')[:3].upper()}"
            location = self.climate_data.get_location_by_id(climate_id, st.session_state)
            if not location:
                return False, f"No climate data found for {climate_id}. Please provide valid climate data.", {}
            
            # Validate climate data
            if not all(k in location for k in ['winter_design_temp', 'monthly_temps', 'monthly_humidity']):
                return False, f"Invalid climate data for {climate_id}. Missing required fields.", {}
            
            # Format conditions
            outdoor_conditions = {
                'design_temperature': location['winter_design_temp'],
                'design_relative_humidity': location['monthly_humidity'].get('Jan', 80.0),
                'ground_temperature': location['monthly_temps'].get('Jan', 10.0),
                'wind_speed': 4.0  # Default as not provided in climate data
            }
            indoor_conditions = {
                'temperature': building_info.get('indoor_temp', 21.0),
                'relative_humidity': building_info.get('indoor_rh', 40.0)
            }
            
            # Debug: Log inputs
            st.write("Debug: Heating Input State", {
                'climate_id': climate_id,
                'outdoor_conditions': outdoor_conditions,
                'indoor_conditions': indoor_conditions,
                'components': {k: len(v) for k, v in building_components.items()},
                'internal_loads': {
                    'people': len(internal_loads.get('people', [])),
                    'lighting': len(internal_loads.get('lighting', [])),
                    'equipment': len(internal_loads.get('equipment', []))
                },
                'building_info': building_info
            })
            
            # Format internal loads
            formatted_internal_loads = {
                'people': {
                    'number': sum(load['num_people'] for load in internal_loads.get('people', [])),
                    'sensible_gain': 70,
                    'operating_hours': f"{internal_loads.get('people', [{}])[0].get('hours_in_operation', 8)}:00-{internal_loads.get('people', [{}])[0].get('hours_in_operation', 8)+10}:00"
                },
                'lights': {
                    'power': sum(load['power'] for load in internal_loads.get('lighting', [])),
                    'use_factor': internal_loads.get('lighting', [{}])[0].get('usage_factor', 0.8),
                    'hours_operation': f"{internal_loads.get('lighting', [{}])[0].get('hours_in_operation', 8)}h"
                },
                'equipment': {
                    'power': sum(load['power'] for load in internal_loads.get('equipment', [])),
                    'use_factor': internal_loads.get('equipment', [{}])[0].get('usage_factor', 0.7),
                    'hours_operation': f"{internal_loads.get('equipment', [{}])[0].get('hours_in_operation', 8)}h"
                },
                'infiltration': {
                    'flow_rate': building_info.get('infiltration_rate', 0.05),
                    'height': building_info.get('building_height', 3.0),
                    'crack_length': building_info.get('crack_length', 10.0)
                },
                'ventilation': {
                    'flow_rate': building_info.get('ventilation_rate', 0.1)
                },
                'usage_factor': 0.7,
                'operating_hours': building_info.get('operating_hours', '8:00-18:00')
            }
            
            # Calculate design loads
            design_loads = self.heating_calculator.calculate_design_heating_load(
                building_components=building_components,
                outdoor_conditions=outdoor_conditions,
                indoor_conditions=indoor_conditions,
                internal_loads=formatted_internal_loads
            )
            if not design_loads:
                return False, "Heating design loads calculation failed. Check input data.", {}
            
            # Get summary
            summary = self.heating_calculator.calculate_heating_load_summary(design_loads)
            if not summary:
                return False, "Heating load summary calculation failed. Check input data.", {}
            
            # Format results
            floor_area = building_info.get('floor_area', 100.0) or 100.0
            results = {
                'total_load': summary['total'] / 1000,  # kW
                'load_per_area': summary['total'] / floor_area,  # W/m²
                'design_heat_loss': summary['subtotal'] / 1000,  # kW
                'safety_factor': summary['safety_factor'] * 100,  # %
                'component_loads': {
                    'walls': design_loads['walls'] / 1000,
                    'roof': design_loads['roofs'] / 1000,
                    'floor': design_loads['floors'] / 1000,
                    'windows': design_loads['windows'] / 1000,
                    'doors': design_loads['doors'] / 1000,
                    'infiltration': (design_loads['infiltration_sensible'] + design_loads['infiltration_latent']) / 1000,
                    'ventilation': (design_loads['ventilation_sensible'] + design_loads['ventilation_latent']) / 1000
                },
                'detailed_loads': {
                    'walls': [],
                    'roofs': [],
                    'floors': [],
                    'windows': [],
                    'doors': [],
                    'infiltration': {
                        'air_flow': formatted_internal_loads['infiltration']['flow_rate'],
                        'delta_t': indoor_conditions['temperature'] - outdoor_conditions['design_temperature'],
                        'load': (design_loads['infiltration_sensible'] + design_loads['infiltration_latent']) / 1000
                    },
                    'ventilation': {
                        'air_flow': formatted_internal_loads['ventilation']['flow_rate'],
                        'delta_t': indoor_conditions['temperature'] - outdoor_conditions['design_temperature'],
                        'load': (design_loads['ventilation_sensible'] + design_loads['ventilation_latent']) / 1000
                    }
                },
                'building_info': building_info
            }
            
            # Populate detailed loads
            delta_t = indoor_conditions['temperature'] - outdoor_conditions['design_temperature']
            for wall in building_components.get('walls', []):
                load = self.heating_calculator.calculate_wall_heating_load(
                    wall=wall,
                    outdoor_temp=outdoor_conditions['design_temperature'],
                    indoor_temp=indoor_conditions['temperature']
                )
                results['detailed_loads']['walls'].append({
                    'name': wall.name,
                    'orientation': wall.orientation.value,
                    'area': wall.area,
                    'u_value': wall.u_value,
                    'delta_t': delta_t,
                    'load': load / 1000
                })
            
            for roof in building_components.get('roofs', []):
                load = self.heating_calculator.calculate_roof_heating_load(
                    roof=roof,
                    outdoor_temp=outdoor_conditions['design_temperature'],
                    indoor_temp=indoor_conditions['temperature']
                )
                results['detailed_loads']['roofs'].append({
                    'name': roof.name,
                    'orientation': roof.orientation.value,
                    'area': roof.area,
                    'u_value': roof.u_value,
                    'delta_t': delta_t,
                    'load': load / 1000
                })
            
            for floor in building_components.get('floors', []):
                load = self.heating_calculator.calculate_floor_heating_load(
                    floor=floor,
                    ground_temp=outdoor_conditions['ground_temperature'],
                    indoor_temp=indoor_conditions['temperature']
                )
                results['detailed_loads']['floors'].append({
                    'name': floor.name,
                    'area': floor.area,
                    'u_value': floor.u_value,
                    'delta_t': indoor_conditions['temperature'] - outdoor_conditions['ground_temperature'],
                    'load': load / 1000
                })
            
            for window in building_components.get('windows', []):
                load = self.heating_calculator.calculate_window_heating_load(
                    window=window,
                    outdoor_temp=outdoor_conditions['design_temperature'],
                    indoor_temp=indoor_conditions['temperature']
                )
                results['detailed_loads']['windows'].append({
                    'name': window.name,
                    'orientation': window.orientation.value,
                    'area': window.area,
                    'u_value': window.u_value,
                    'delta_t': delta_t,
                    'load': load / 1000
                })
            
            for door in building_components.get('doors', []):
                load = self.heating_calculator.calculate_door_heating_load(
                    door=door,
                    outdoor_temp=outdoor_conditions['design_temperature'],
                    indoor_temp=indoor_conditions['temperature']
                )
                results['detailed_loads']['doors'].append({
                    'name': door.name,
                    'orientation': door.orientation.value,
                    'area': door.area,
                    'u_value': door.u_value,
                    'delta_t': delta_t,
                    'load': load / 1000
                })
            
            # Debug: Log results
            st.write("Debug: Heating Results", {
                'total_load': results.get('total_load', 'N/A'),
                'component_loads': results.get('component_loads', 'N/A'),
                'detailed_loads': {k: len(v) if isinstance(v, list) else v for k, v in results.get('detailed_loads', {}).items()}
            })
            
            return True, "Heating calculation completed.", results
        except Exception as e:
            st.error(f"Heating calculation error: {str(e)}")
            return False, f"Heating calculation error: {str(e)}", {}
    
    def display_calculation_results(self):
        st.title("Calculation Results")
        
        col1, col2 = st.columns(2)
        with col1:
            calculate_button = st.button("Calculate Loads")
        with col2:
            debug = st.checkbox("Debug Mode", value=True)  # Enable by default for now
            
        if calculate_button:
            # Reset results
            st.session_state.calculation_results = {'cooling': {}, 'heating': {}}
            
            # Calculate cooling load
            cooling_success, cooling_message, cooling_results = self.calculate_cooling()
            if cooling_success:
                st.session_state.calculation_results['cooling'] = cooling_results
                st.success(cooling_message)
            else:
                st.error(cooling_message)
            
            # Calculate heating load
            heating_success, heating_message, heating_results = self.calculate_heating()
            if heating_success:
                st.session_state.calculation_results['heating'] = heating_results
                st.success(heating_message)
            else:
                st.error(heating_message)
        
        # Display results
        self.results_display.display_results(st.session_state)
        
        # Navigation
        col1, col2 = st.columns(2)
        with col1:
            st.button(
                "Back to Internal Loads",
                on_click=lambda: setattr(st.session_state, "page", "Internal Loads")
            )
        with col2:
            st.button(
                "Continue to Export Data",
                on_click=lambda: setattr(st.session_state, "page", "Export Data")
            )

if __name__ == "__main__":
    app = HVACCalculator()