app/main.py

"""
HVAC Calculator Code Documentation
Updated 2025-04-27: Enhanced climate ID generation, input validation, debug mode, and error handling.
Updated 2025-04-28: Added activity-level-based internal gains, ground temperature validation, ASHRAE 62.1 ventilation rates, negative load prevention, and improved usability.
"""

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

# NEW: ASHRAE 62.1 Ventilation Rates (Table 6.1)
VENTILATION_RATES = {
    "Office": {"people_rate": 2.5, "area_rate": 0.3},  # L/s/person, L/s/m²
    "Classroom": {"people_rate": 5.0, "area_rate": 0.9},
    "Retail": {"people_rate": 3.8, "area_rate": 0.9},
    "Restaurant": {"people_rate": 5.0, "area_rate": 1.8},
    "Custom": {"people_rate": 0.0, "area_rate": 0.0}
}

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 = {}
        
        if 'debug_mode' not in st.session_state:
            st.session_state.debug_mode = False
        
        # 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()
        
        # Persist ClimateData in session_state
        if 'climate_data_obj' not in st.session_state:
            st.session_state.climate_data_obj = ClimateData()
        self.climate_data = st.session_state.climate_data_obj
        
        # Load default climate data if locations are empty
        try:
            if not self.climate_data.locations:
                self.climate_data = ClimateData.from_json("/home/user/app/climate_data.json")
                st.session_state.climate_data_obj = self.climate_data
        except FileNotFoundError:
            st.warning("Default climate data file not found. Please enter climate data manually.")
        
        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 generate_climate_id(self, country: str, city: str) -> str:
        """Generate a climate ID from country and city names."""
        try:
            country = country.strip().title()
            city = city.strip().title()
            if len(country) < 2 or len(city) < 3:
                raise ValueError("Country and city names must be at least 2 and 3 characters long, respectively.")
            return f"{country[:2].upper()}-{city[:3].upper()}"
        except Exception as e:
            raise ValueError(f"Invalid country or city name: {str(e)}")
    
    def validate_calculation_inputs(self) -> Tuple[bool, str]:
        """Validate inputs for cooling and heating calculations."""
        building_info = st.session_state.get('building_info', {})
        components = st.session_state.get('components', {})
        climate_data = st.session_state.get('climate_data', {})
        
        # Check building info
        if not building_info.get('floor_area', 0) > 0:
            return False, "Floor area must be positive."
        if not any(components.get(key, []) for key in ['walls', 'roofs', 'windows']):
            return False, "At least one wall, roof, or window must be defined."
        
        # NEW: Validate climate data using climate_data.py
        if not climate_data:
            return False, "Climate data is missing."
        if not self.climate_data.validate_climate_data(climate_data):
            return False, "Invalid climate data format or values."
        
        # Validate components
        for component_type in ['walls', 'roofs', 'windows', 'doors', 'floors']:
            for comp in components.get(component_type, []):
                if comp.area <= 0:
                    return False, f"Invalid area for {component_type}: {comp.name}"
                if comp.u_value <= 0:
                    return False, f"Invalid U-value for {component_type}: {comp.name}"
                # NEW: Validate ground temperature for floors
                if component_type == 'floors' and getattr(comp, 'ground_contact', False):
                    if not -10 <= comp.ground_temperature_c <= 40:
                        return False, f"Ground temperature for {comp.name} must be between -10°C and 40°C"
                    # NEW: Validate perimeter
                    if getattr(comp, 'perimeter', 0) < 0:
                        return False, f"Perimeter for {comp.name} cannot be negative"
        
        # NEW: Validate ventilation rate
        if building_info.get('ventilation_rate', 0) < 0:
            return False, "Ventilation rate cannot be negative"
        if building_info.get('zone_type', '') == 'Custom' and building_info.get('ventilation_rate', 0) == 0:
            return False, "Custom ventilation rate must be specified"
        
        return True, "Inputs valid."
    
    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", "Ventilation"])  # NEW: Added Ventilation tab
        
        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,
                    help="Total number of occupants in the building"
                )
                activity_level = st.selectbox(
                    "Activity Level",
                    ["Seated/Resting", "Light Work", "Moderate Work", "Heavy Work"],
                    help="Select typical activity level (affects internal heat gains per ASHRAE)"
                )
                zone_type = st.selectbox(
                    "Zone Type",
                    ["Office", "Classroom", "Retail", "Residential"],
                    help="Select zone type for occupancy characteristics"
                )
                hours_in_operation = st.number_input(
                    "Hours in Operation",
                    min_value=0.0,
                    max_value=24.0,
                    value=8.0,
                    step=0.5,
                    help="Daily hours of occupancy"
                )
                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,
                    help="Total lighting power consumption"
                )
                usage_factor = st.number_input(
                    "Usage Factor",
                    min_value=0.0,
                    max_value=1.0,
                    value=0.8,
                    step=0.1,
                    help="Fraction of time lighting is in use (0 to 1)"
                )
                zone_type = st.selectbox(
                    "Zone Type",
                    ["Office", "Classroom", "Retail", "Residential"],
                    help="Select zone type for lighting characteristics"
                )
                hours_in_operation = st.number_input(
                    "Hours in Operation",
                    min_value=0.0,
                    max_value=24.0,
                    value=8.0,
                    step=0.5,
                    help="Daily hours of lighting operation"
                )
                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,
                    help="Total equipment power consumption"
                )
                usage_factor = st.number_input(
                    "Usage Factor",
                    min_value=0.0,
                    max_value=1.0,
                    value=0.7,
                    step=0.1,
                    help="Fraction of time equipment is in use (0 to 1)"
                )
                radiation_fraction = st.number_input(
                    "Radiation Fraction",
                    min_value=0.0,
                    max_value=1.0,
                    value=0.3,
                    step=0.1,
                    help="Fraction of heat gain radiated to surroundings"
                )
                zone_type = st.selectbox(
                    "Zone Type",
                    ["Office", "Classroom", "Retail", "Residential"],
                    help="Select zone type for equipment characteristics"
                )
                hours_in_operation = st.number_input(
                    "Hours in Operation",
                    min_value=0.0,
                    max_value=24.0,
                    value=8.0,
                    step=0.5,
                    help="Daily hours of equipment operation"
                )
                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"):
                    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()
        
        with tabs[3]:  # NEW: Ventilation tab
            st.subheader("Ventilation Requirements (ASHRAE 62.1)")
            with st.form("ventilation_form"):
                col1, col2 = st.columns(2)
                with col1:
                    zone_type = st.selectbox(
                        "Zone Type",
                        ["Office", "Classroom", "Retail", "Restaurant", "Custom"],
                        help="Select building zone type for ASHRAE 62.1 ventilation rates"
                    )
                    ventilation_method = st.selectbox(
                        "Ventilation Method",
                        ["Constant Volume", "Demand-Controlled"],
                        help="Constant Volume uses fixed rate; Demand-Controlled adjusts based on occupancy"
                    )
                with col2:
                    if zone_type == "Custom":
                        people_rate = st.number_input(
                            "Ventilation Rate per Person (L/s/person)",
                            min_value=0.0,
                            value=2.5,
                            step=0.1,
                            help="Custom ventilation rate per person (ASHRAE 62.1)"
                        )
                        area_rate = st.number_input(
                            "Ventilation Rate per Area (L/s/m²)",
                            min_value=0.0,
                            value=0.3,
                            step=0.1,
                            help="Custom ventilation rate per floor area (ASHRAE 62.1)"
                        )
                    else:
                        people_rate = VENTILATION_RATES[zone_type]["people_rate"]
                        area_rate = VENTILATION_RATES[zone_type]["area_rate"]
                        st.write(f"People Rate: {people_rate} L/s/person (ASHRAE 62.1)")
                        st.write(f"Area Rate: {area_rate} L/s/m² (ASHRAE 62.1)")
                
                if st.form_submit_button("Save Ventilation Settings"):
                    total_people = sum(load['num_people'] for load in st.session_state.internal_loads.get('people', []))
                    floor_area = st.session_state.building_info.get('floor_area', 100.0)
                    ventilation_rate = (
                        (total_people * people_rate + floor_area * area_rate) / 1000  # Convert L/s to m³/s
                    )
                    if ventilation_method == 'Demand-Controlled':
                        ventilation_rate *= 0.75  # Reduce by 25% for DCV
                    st.session_state.building_info.update({
                        'zone_type': zone_type,
                        'ventilation_method': ventilation_method,
                        'ventilation_rate': ventilation_rate
                    })
                    st.success(f"Ventilation settings saved! Total rate: {ventilation_rate:.3f} m³/s")
        
        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:
            # Validate inputs
            valid, message = self.validate_calculation_inputs()
            if not valid:
                return False, message, {}
            
            # 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', {})
            
            # Check climate data
            if "climate_data" not in st.session_state or not st.session_state["climate_data"]:
                return False, "Please enter climate data in the 'Climate Data' page.", {}
            
            # Extract climate data
            country = building_info.get('country', '').strip().title()
            city = building_info.get('city', '').strip().title()
            if not country or not city:
                return False, "Country and city must be set in Building Information.", {}
            climate_id = self.generate_climate_id(country, city)
            location = self.climate_data.get_location_by_id(climate_id, st.session_state)
            if not location:
                available_locations = list(self.climate_data.locations.keys())[:5]
                return False, f"No climate data for {climate_id}. Available locations: {', '.join(available_locations)}...", {}
            
            # 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': location['latitude'],  # Pass raw latitude value, validation will happen in cooling_load.py
                'wind_speed': building_info.get('wind_speed', 4.0),
                '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)
            }
            
            if st.session_state.get('debug_mode', False):
                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.", {}
            
            # Ensure summary has all required keys
            if 'total' not in summary:
                # Calculate total if missing
                if 'total_sensible' in summary and 'total_latent' in summary:
                    summary['total'] = summary['total_sensible'] + summary['total_latent']
                else:
                    # Fallback to sum of design loads if needed
                    total_load = sum(value for key, value in design_loads.items() if key != 'design_hour')
                    summary = {
                        'total_sensible': total_load * 0.7,  # Approximate sensible ratio
                        'total_latent': total_load * 0.3,    # Approximate latent ratio
                        'total': total_load
                    }
            
            # 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['infiltration_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
                )
                # Ensure load_dict has a 'total' key
                if 'total' not in load_dict:
                    if 'conduction' in load_dict and 'solar' in load_dict:
                        load_dict['total'] = load_dict['conduction'] + load_dict['solar']
                    else:
                        load_dict['total'] = window.u_value * window.area * (outdoor_conditions['temperature'] - indoor_conditions['temperature'])
                
                # Pass latitude directly to get_scl method which has its own validation
                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(
                        latitude=outdoor_conditions['latitude'],
                        month=outdoor_conditions['month'].title(),
                        orientation=window.orientation.value,
                        hour=design_loads['design_hour']
                    ),
                    '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']
                        )
                        # Ensure load_dict has a 'total' key
                        if 'total' not in load_dict and ('sensible' in load_dict or 'latent' in load_dict):
                            load_dict['total'] = load_dict.get('sensible', 0) + load_dict.get('latent', 0)
                    elif load_type == 'lighting':
                        light_load = self.cooling_calculator.calculate_lights_cooling_load(
                            power=load['power'],
                            use_factor=load['usage_factor'],
                            special_allowance=0.1,
                            hour=design_loads['design_hour']
                        )
                        load_dict = {'total': light_load if light_load is not None else 0}
                    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']
                        )
                        # Ensure load_dict has a 'total' key
                        if 'total' not in load_dict and ('sensible' in load_dict or 'latent' in load_dict):
                            load_dict['total'] = load_dict.get('sensible', 0) + load_dict.get('latent', 0)
                    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.get('total', 0)),
                        'clf': self.cooling_calculator.ashrae_tables.get_clf_people(
                            zone_type='A',
                            hours_occupied='6h',  # Using valid '6h' instead of dynamic value that might not exist
                            hour=design_loads['design_hour']
                        ) if load_type == 'people' else 1.0,
                        'load': load_dict.get('total', 0) / 1000
                    })
            
            if st.session_state.get('debug_mode', False):
                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 ValueError as ve:
            st.error(f"Input error in cooling calculation: {str(ve)}")
            return False, f"Input error: {str(ve)}", {}
        except KeyError as ke:
            st.error(f"Missing data in cooling calculation: {str(ke)}")
            return False, f"Missing data: {str(ke)}", {}
        except Exception as e:
            st.error(f"Unexpected error in cooling calculation: {str(e)}")
            return False, f"Unexpected error: {str(e)}", {}
    
    def calculate_heating(self) -> Tuple[bool, str, Dict]:
        """
        Calculate heating loads using HeatingLoadCalculator.
        Returns: (success, message, results)
        """
        try:
            # Validate inputs
            valid, message = self.validate_calculation_inputs()
            if not valid:
                return False, message, {}
            
            # 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', {})
            
            # Check climate data
            if "climate_data" not in st.session_state or not st.session_state["climate_data"]:
                return False, "Please enter climate data in the 'Climate Data' page.", {}
            
            # Extract climate data
            country = building_info.get('country', '').strip().title()
            city = building_info.get('city', '').strip().title()
            if not country or not city:
                return False, "Country and city must be set in Building Information.", {}
            climate_id = self.generate_climate_id(country, city)
            location = self.climate_data.get_location_by_id(climate_id, st.session_state)
            if not location:
                available_locations = list(self.climate_data.locations.keys())[:5]
                return False, f"No climate data for {climate_id}. Available locations: {', '.join(available_locations)}...", {}
            
            # 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.", {}
            
            # NEW: Calculate ground temperature from floors or fallback to climate data
            ground_contact_floors = [f for f in building_components.get('floors', []) if getattr(f, 'ground_contact', False)]
            ground_temperature = (
                sum(f.ground_temperature_c for f in ground_contact_floors) / len(ground_contact_floors)
                if ground_contact_floors else
                location['monthly_temps'].get('Jan', 10.0)
            )
            if not -10 <= ground_temperature <= 40:
                return False, f"Invalid ground temperature: {ground_temperature}°C", {}
            
            # NEW: Skip heating calculation if outdoor temp exceeds indoor temp
            indoor_temp = building_info.get('indoor_temp', 21.0)
            outdoor_temp = location['winter_design_temp']
            if outdoor_temp >= indoor_temp:
                results = {
                    'total_load': 0.0,
                    'load_per_area': 0.0,
                    'design_heat_loss': 0.0,
                    'safety_factor': 115.0,
                    'component_loads': {
                        'walls': 0.0,
                        'roof': 0.0,
                        'floor': 0.0,
                        'windows': 0.0,
                        'doors': 0.0,
                        'infiltration': 0.0,
                        'ventilation': 0.0
                    },
                    'detailed_loads': {
                        'walls': [],
                        'roofs': [],
                        'floors': [],
                        'windows': [],
                        'doors': [],
                        'infiltration': {'air_flow': 0.0, 'delta_t': 0.0, 'load': 0.0},
                        'ventilation': {'air_flow': 0.0, 'delta_t': 0.0, 'load': 0.0}
                    },
                    'building_info': building_info
                }
                return True, "No heating required (outdoor temp exceeds indoor temp).", results
            
            # Format conditions
            outdoor_conditions = {
                'design_temperature': location['winter_design_temp'],
                'design_relative_humidity': location['monthly_humidity'].get('Jan', 80.0),
                'ground_temperature': ground_temperature,
                'wind_speed': building_info.get('wind_speed', 4.0)
            }
            indoor_conditions = {
                'temperature': indoor_temp,
                'relative_humidity': building_info.get('indoor_rh', 40.0)
            }
            
            if st.session_state.get('debug_mode', False):
                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
                })
            
            # NEW: Activity-level-based sensible gains
            ACTIVITY_GAINS = {
                'Seated/Resting': 70.0,  # W/person
                'Light Work': 85.0,
                'Moderate Work': 100.0,
                'Heavy Work': 150.0
            }
            
            # Format internal loads
            formatted_internal_loads = {
                'people': {
                    'number': sum(load['num_people'] for load in internal_loads.get('people', [])),
                    'sensible_gain': ACTIVITY_GAINS.get(
                        internal_loads.get('people', [{}])[0].get('activity_level', 'Seated/Resting'),
                        70.0
                    ),
                    '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
                })
            
            if st.session_state.get('debug_mode', False):
                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 ValueError as ve:
            st.error(f"Input error in heating calculation: {str(ve)}")
            return False, f"Input error: {str(ve)}", {}
        except KeyError as ke:
            st.error(f"Missing data in heating calculation: {str(ke)}")
            return False, f"Missing data: {str(ke)}", {}
        except Exception as e:
            st.error(f"Unexpected error in heating calculation: {str(e)}")
            return False, f"Unexpected 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:
            st.session_state.debug_mode = st.checkbox("Debug Mode", value=st.session_state.get('debug_mode', False))
            
        if calculate_button:
            # Reset results
            st.session_state.calculation_results = {'cooling': {}, 'heating': {}}
            
            with st.spinner("Calculating loads..."):
                # 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()