pages/heating_calculator.py

"""
Heating Load Calculator Page

This module implements the heating load calculator interface for the HVAC Load Calculator web application.
It provides a step-by-step form for inputting building information and calculates heating loads
using the ASHRAE method.
"""

import streamlit as st
import pandas as pd
import numpy as np
import plotly.express as px
import plotly.graph_objects as go
import json
import os
import sys
from pathlib import Path
from datetime import datetime

# Add the parent directory to sys.path to import modules
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# Import custom modules
from heating_load import HeatingLoadCalculator
from reference_data import ReferenceData
from utils.validation import validate_input, ValidationWarning
from utils.export import export_data


def load_session_state():
    """Initialize or load session state variables."""
    # Initialize session state for form data
    if 'heating_form_data' not in st.session_state:
        st.session_state.heating_form_data = {
            'building_info': {},
            'building_envelope': {},
            'windows': {},
            'ventilation': {},
            'occupancy': {},
            'results': {}
        }
    
    # Initialize session state for validation warnings
    if 'heating_warnings' not in st.session_state:
        st.session_state.heating_warnings = {
            'building_info': [],
            'building_envelope': [],
            'windows': [],
            'ventilation': [],
            'occupancy': []
        }
    
    # Initialize session state for form completion status
    if 'heating_completed' not in st.session_state:
        st.session_state.heating_completed = {
            'building_info': False,
            'building_envelope': False,
            'windows': False,
            'ventilation': False,
            'occupancy': False
        }
    
    # Initialize session state for calculation results
    if 'heating_results' not in st.session_state:
        st.session_state.heating_results = None


def building_info_form(ref_data):
    """
    Form for building information.
    
    Args:
        ref_data: Reference data object
    """
    st.subheader("Building Information")
    st.write("Enter general building information, location, and design temperatures.")
    
    # Get location options from reference data
    location_options = {loc_id: loc_data['name'] for loc_id, loc_data in ref_data.locations.items()}
    
    col1, col2 = st.columns(2)
    
    with col1:
        # Building name
        building_name = st.text_input(
            "Building Name",
            value=st.session_state.heating_form_data['building_info'].get('building_name', ''),
            help="Enter a name for this building or project"
        )
        
        # Location selection
        location = st.selectbox(
            "Location",
            options=list(location_options.keys()),
            format_func=lambda x: location_options[x],
            index=list(location_options.keys()).index(st.session_state.heating_form_data['building_info'].get('location', 'sydney')) if st.session_state.heating_form_data['building_info'].get('location') in location_options else 0,
            help="Select the location of the building"
        )
        
        # Get climate data for selected location
        location_data = ref_data.get_location_data(location)
        
        # Indoor design temperature
        indoor_temp = st.number_input(
            "Indoor Design Temperature (°C)",
            value=float(st.session_state.heating_form_data['building_info'].get('indoor_temp', 21.0)),
            min_value=15.0,
            max_value=25.0,
            step=0.5,
            help="Recommended indoor design temperature for heating is 21°C for living areas and 17°C for bedrooms"
        )
    
    with col2:
        # Building type
        building_type = st.selectbox(
            "Building Type",
            options=["Residential", "Small Office", "Educational", "Other"],
            index=["Residential", "Small Office", "Educational", "Other"].index(st.session_state.heating_form_data['building_info'].get('building_type', 'Residential')),
            help="Select the type of building"
        )
        
        # Outdoor design temperature (with default from location data)
        outdoor_temp = st.number_input(
            "Outdoor Design Temperature (°C)",
            value=float(st.session_state.heating_form_data['building_info'].get('outdoor_temp', location_data['winter_design_temp'])),
            min_value=-10.0,
            max_value=15.0,
            step=0.5,
            help=f"Default value is based on selected location ({location_data['name']})"
        )
    
    # Building dimensions
    st.subheader("Building Dimensions")
    
    col1, col2, col3 = st.columns(3)
    
    with col1:
        length = st.number_input(
            "Length (m)",
            value=float(st.session_state.heating_form_data['building_info'].get('length', 10.0)),
            min_value=1.0,
            step=0.1,
            help="Building length in meters"
        )
    
    with col2:
        width = st.number_input(
            "Width (m)",
            value=float(st.session_state.heating_form_data['building_info'].get('width', 8.0)),
            min_value=1.0,
            step=0.1,
            help="Building width in meters"
        )
    
    with col3:
        height = st.number_input(
            "Height (m)",
            value=float(st.session_state.heating_form_data['building_info'].get('height', 2.7)),
            min_value=1.0,
            step=0.1,
            help="Floor-to-ceiling height in meters"
        )
    
    # Calculate floor area and volume
    floor_area = length * width
    volume = floor_area * height
    
    st.info(f"Floor Area: {floor_area:.2f} m² | Volume: {volume:.2f} m³")
    
    # Save form data to session state
    form_data = {
        'building_name': building_name,
        'building_type': building_type,
        'location': location,
        'location_name': location_data['name'],
        'indoor_temp': indoor_temp,
        'outdoor_temp': outdoor_temp,
        'length': length,
        'width': width,
        'height': height,
        'floor_area': floor_area,
        'volume': volume,
        'temp_diff': indoor_temp - outdoor_temp
    }
    
    # Validate inputs
    warnings = []
    
    # Check if building name is provided
    if not building_name:
        warnings.append(ValidationWarning("Building name is empty", "Consider adding a building name for reference"))
    
    # Check if temperature difference is reasonable
    if form_data['temp_diff'] <= 0:
        warnings.append(ValidationWarning(
            "Invalid temperature difference",
            "Indoor temperature should be higher than outdoor temperature for heating load calculation",
            is_critical=False  # Changed to non-critical to allow proceeding with warnings
        ))
    
    # Check if dimensions are reasonable
    if floor_area > 500:
        warnings.append(ValidationWarning(
            "Large floor area",
            "Floor area exceeds 500 m², verify if this is correct for a residential building"
        ))
    
    if height < 2.4 or height > 3.5:
        warnings.append(ValidationWarning(
            "Unusual ceiling height",
            "Typical residential ceiling heights are between 2.4m and 3.5m"
        ))
    
    # Save warnings to session state
    st.session_state.heating_warnings['building_info'] = warnings
    
    # Display warnings if any
    if warnings:
        st.warning("Please review the following warnings:")
        for warning in warnings:
            st.write(f"- {warning.message}" + (" (Critical)" if warning.is_critical else ""))
            st.write(f"  Suggestion: {warning.suggestion}")
    
    # Save form data regardless of warnings
    st.session_state.heating_form_data['building_info'] = form_data
    
    # Mark this step as completed if there are no critical warnings
    st.session_state.heating_completed['building_info'] = not any(w.is_critical for w in warnings)
    
    # Navigation buttons
    col1, col2 = st.columns([1, 1])
    
    with col2:
        next_button = st.button("Next: Building Envelope →", key="heating_building_info_next")
        if next_button:
            st.session_state.heating_active_tab = "building_envelope"
            st.experimental_rerun()


def building_envelope_form(ref_data):
    """
    Form for building envelope information.
    
    Args:
        ref_data: Reference data object
    """
    st.subheader("Building Envelope")
    st.write("Enter information about walls, roof, and floor construction.")
    
    # Get building dimensions from previous step
    building_info = st.session_state.heating_form_data['building_info']
    length = building_info.get('length', 10.0)
    width = building_info.get('width', 8.0)
    height = building_info.get('height', 2.7)
    temp_diff = building_info.get('temp_diff', 16.5)
    
    # Calculate default areas
    default_wall_area = 2 * (length + width) * height
    default_roof_area = length * width
    default_floor_area = length * width
    
    # Initialize envelope data if not already in session state
    if 'walls' not in st.session_state.heating_form_data['building_envelope']:
        st.session_state.heating_form_data['building_envelope']['walls'] = []
    
    if 'roof' not in st.session_state.heating_form_data['building_envelope']:
        st.session_state.heating_form_data['building_envelope']['roof'] = {}
    
    if 'floor' not in st.session_state.heating_form_data['building_envelope']:
        st.session_state.heating_form_data['building_envelope']['floor'] = {}
    
    # Walls section
    st.write("### Walls")
    
    # Get wall material options from reference data
    wall_material_options = {mat_id: mat_data['name'] for mat_id, mat_data in ref_data.materials['walls'].items()}
    # Add custom option
    wall_material_options["custom_walls"] = "Custom Wall (User-defined)"
    
    # Display existing wall entries
    if st.session_state.heating_form_data['building_envelope']['walls']:
        st.write("Current walls:")
        walls_df = pd.DataFrame(st.session_state.heating_form_data['building_envelope']['walls'])
        walls_df['Material'] = walls_df['material_id'].map(lambda x: wall_material_options.get(x, "Unknown"))
        # Add orientation column with default value if not present
        walls_df['orientation'] = walls_df['orientation'].fillna('not specified')
        walls_df = walls_df[['name', 'Material', 'area', 'u_value', 'orientation']]
        walls_df.columns = ['Name', 'Material', 'Area (m²)', 'U-Value (W/m²°C)', 'Orientation']
        st.dataframe(walls_df)
    
    # Add new wall form
    st.write("Add a new wall:")
    
    col1, col2 = st.columns(2)
    
    with col1:
        wall_name = st.text_input("Wall Name", value="", key="new_wall_name_heating")
        wall_material = st.selectbox(
            "Wall Material",
            options=list(wall_material_options.keys()),
            format_func=lambda x: wall_material_options[x],
            key="new_wall_material_heating"
        )
        
        # Add wall orientation selection
        wall_orientation = st.selectbox(
            "Wall Orientation",
            options=["north", "east", "south", "west"],
            key="new_wall_orientation_heating"
        )
        
        # Get material properties
        material_data = ref_data.get_material_by_type("walls", wall_material)
        u_value = material_data['u_value']
        
        # Add custom U-value input if custom material is selected
        if wall_material == "custom_walls":
            u_value = st.number_input(
                "Custom U-Value (W/m²°C)",
                value=1.0,
                min_value=0.1,
                max_value=5.0,
                step=0.1,
                key="custom_wall_u_value_heating"
            )
            
            # Store custom material in session state
            if "custom_materials" not in st.session_state:
                st.session_state.custom_materials = {}
            
            st.session_state.custom_materials["walls"] = {
                "name": "Custom Wall",
                "u_value": u_value,
                "r_value": 1.0 / u_value if u_value > 0 else 1.0,
                "description": "Custom wall with user-defined properties"
            }
        
    with col2:
        wall_area = st.number_input(
            "Wall Area (m²)",
            value=default_wall_area / 4,  # Default to 1/4 of total wall area as a starting point
            min_value=0.1,
            step=0.1,
            key="new_wall_area_heating"
        )
        
        st.write(f"Material U-Value: {u_value} W/m²°C")
        st.write(f"Heat Loss: {u_value * wall_area * temp_diff:.2f} W")
    
    # Add wall button
    if st.button("Add Wall", key="add_wall_heating"):
        new_wall = {
            'name': wall_name if wall_name else f"Wall {len(st.session_state.heating_form_data['building_envelope']['walls']) + 1}",
            'material_id': wall_material,
            'area': wall_area,
            'u_value': u_value,
            'temp_diff': temp_diff,
            'orientation': wall_orientation  # Add orientation to wall data
        }
        st.session_state.heating_form_data['building_envelope']['walls'].append(new_wall)
        st.experimental_rerun()
    
    # Roof section
    st.write("### Roof")
    
    # Get roof material options from reference data
    roof_material_options = {mat_id: mat_data['name'] for mat_id, mat_data in ref_data.materials['roofs'].items()}
    # Add custom option
    roof_material_options["custom_roofs"] = "Custom Roof (User-defined)"
    
    col1, col2 = st.columns(2)
    
    with col1:
        roof_material = st.selectbox(
            "Roof Material",
            options=list(roof_material_options.keys()),
            format_func=lambda x: roof_material_options[x],
            index=list(roof_material_options.keys()).index(st.session_state.heating_form_data['building_envelope'].get('roof', {}).get('material_id', 'metal_deck_insulated')) if st.session_state.heating_form_data['building_envelope'].get('roof', {}).get('material_id') in roof_material_options else 0
        )
        
        # Get material properties
        material_data = ref_data.get_material_by_type("roofs", roof_material)
        roof_u_value = material_data['u_value']
        
        # Add custom U-value input if custom material is selected
        if roof_material == "custom_roofs":
            roof_u_value = st.number_input(
                "Custom Roof U-Value (W/m²°C)",
                value=1.0,
                min_value=0.1,
                max_value=5.0,
                step=0.1,
                key="custom_roof_u_value_heating"
            )
            
            # Store custom material in session state
            if "custom_materials" not in st.session_state:
                st.session_state.custom_materials = {}
            
            st.session_state.custom_materials["roofs"] = {
                "name": "Custom Roof",
                "u_value": roof_u_value,
                "r_value": 1.0 / roof_u_value if roof_u_value > 0 else 1.0,
                "description": "Custom roof with user-defined properties"
            }
        
    with col2:
        roof_area = st.number_input(
            "Roof Area (m²)",
            value=float(st.session_state.heating_form_data['building_envelope'].get('roof', {}).get('area', default_roof_area)),
            min_value=0.1,
            step=0.1,
            key="roof_area_heating"
        )
        
        st.write(f"Material U-Value: {roof_u_value} W/m²°C")
        st.write(f"Heat Loss: {roof_u_value * roof_area * temp_diff:.2f} W")
    
    # Save roof data
    st.session_state.heating_form_data['building_envelope']['roof'] = {
        'material_id': roof_material,
        'area': roof_area,
        'u_value': roof_u_value,
        'temp_diff': temp_diff
    }
    
    # Floor section
    st.write("### Floor")
    
    # Get floor material options from reference data
    floor_material_options = {mat_id: mat_data['name'] for mat_id, mat_data in ref_data.materials['floors'].items()}
    # Add custom option
    floor_material_options["custom_floors"] = "Custom Floor (User-defined)"
    
    col1, col2 = st.columns(2)
    
    with col1:
        floor_material = st.selectbox(
            "Floor Material",
            options=list(floor_material_options.keys()),
            format_func=lambda x: floor_material_options[x],
            index=list(floor_material_options.keys()).index(st.session_state.heating_form_data['building_envelope'].get('floor', {}).get('material_id', 'concrete_slab_ground')) if st.session_state.heating_form_data['building_envelope'].get('floor', {}).get('material_id') in floor_material_options else 0
        )
        
        # Get material properties
        material_data = ref_data.get_material_by_type("floors", floor_material)
        floor_u_value = material_data['u_value']
        
        # Add custom U-value input if custom material is selected
        if floor_material == "custom_floors":
            floor_u_value = st.number_input(
                "Custom Floor U-Value (W/m²°C)",
                value=1.0,
                min_value=0.1,
                max_value=5.0,
                step=0.1,
                key="custom_floor_u_value_heating"
            )
            
            # Store custom material in session state
            if "custom_materials" not in st.session_state:
                st.session_state.custom_materials = {}
            
            st.session_state.custom_materials["floors"] = {
                "name": "Custom Floor",
                "u_value": floor_u_value,
                "r_value": 1.0 / floor_u_value if floor_u_value > 0 else 1.0,
                "description": "Custom floor with user-defined properties"
            }
        
    with col2:
        floor_area = st.number_input(
            "Floor Area (m²)",
            value=float(st.session_state.heating_form_data['building_envelope'].get('floor', {}).get('area', default_floor_area)),
            min_value=0.1,
            step=0.1,
            key="floor_area_heating"
        )
        
        st.write(f"Material U-Value: {floor_u_value} W/m²°C")
        st.write(f"Heat Loss: {floor_u_value * floor_area * temp_diff:.2f} W")
    
    # Save floor data
    st.session_state.heating_form_data['building_envelope']['floor'] = {
        'material_id': floor_material,
        'area': floor_area,
        'u_value': floor_u_value,
        'temp_diff': temp_diff
    }
    
    # Validate inputs
    warnings = []
    
    # Check if walls are defined
    if not st.session_state.heating_form_data['building_envelope']['walls']:
        warnings.append(ValidationWarning(
            "No walls defined",
            "Add at least one wall to continue",
            is_critical=False  # Changed to non-critical to allow proceeding with warnings
        ))
    
    # Check if total wall area is reasonable
    total_wall_area = sum(wall['area'] for wall in st.session_state.heating_form_data['building_envelope']['walls'])
    expected_wall_area = 2 * (length + width) * height
    
    if total_wall_area < expected_wall_area * 0.8 or total_wall_area > expected_wall_area * 1.2:
        warnings.append(ValidationWarning(
            "Unusual wall area",
            f"Total wall area ({total_wall_area:.2f} m²) differs significantly from the expected area ({expected_wall_area:.2f} m²) based on building dimensions"
        ))
    
    # Check if roof area matches floor area
    if abs(roof_area - floor_area) > 1.0:
        warnings.append(ValidationWarning(
            "Roof area doesn't match floor area",
            "For a simple building, roof area should approximately match floor area"
        ))
    
    # Save warnings to session state
    st.session_state.heating_warnings['building_envelope'] = warnings
    
    # Display warnings if any
    if warnings:
        st.warning("Please review the following warnings:")
        for warning in warnings:
            st.write(f"- {warning.message}" + (" (Critical)" if warning.is_critical else ""))
            st.write(f"  Suggestion: {warning.suggestion}")
    
    # Mark this step as completed if there are no critical warnings
    st.session_state.heating_completed['building_envelope'] = not any(w.is_critical for w in warnings)
    
    # Navigation buttons
    col1, col2 = st.columns([1, 1])
    
    with col1:
        prev_button = st.button("← Back: Building Information", key="heating_building_envelope_prev")
        if prev_button:
            st.session_state.heating_active_tab = "building_info"
            st.experimental_rerun()
    
    with col2:
        next_button = st.button("Next: Windows & Doors →", key="heating_building_envelope_next")
        if next_button:
            st.session_state.heating_active_tab = "windows"
            st.experimental_rerun()


def windows_form(ref_data):
    """
    Form for windows and doors information.
    
    Args:
        ref_data: Reference data object
    """
    st.subheader("Windows & Doors")
    st.write("Enter information about windows and doors.")
    
    # Get temperature difference from building info
    temp_diff = st.session_state.heating_form_data['building_info'].get('temp_diff', 16.5)
    
    # Initialize windows data if not already in session state
    if 'windows' not in st.session_state.heating_form_data['windows']:
        st.session_state.heating_form_data['windows']['windows'] = []
    
    if 'doors' not in st.session_state.heating_form_data['windows']:
        st.session_state.heating_form_data['windows']['doors'] = []
    
    # Windows section
    st.write("### Windows")
    
    # Get glass type options from reference data
    glass_type_options = {glass_id: glass_data['name'] for glass_id, glass_data in ref_data.glass_types.items()}
    
    # Display existing window entries
    if st.session_state.heating_form_data['windows']['windows']:
        st.write("Current windows:")
        windows_df = pd.DataFrame(st.session_state.heating_form_data['windows']['windows'])
        windows_df['Glass Type'] = windows_df['glass_type'].map(lambda x: glass_type_options.get(x, "Unknown"))
        windows_df = windows_df[['name', 'orientation', 'Glass Type', 'area', 'u_value']]
        windows_df.columns = ['Name', 'Orientation', 'Glass Type', 'Area (m²)', 'U-Value (W/m²°C)']
        st.dataframe(windows_df)
    
    # Add new window form
    st.write("Add a new window:")
    
    col1, col2 = st.columns(2)
    
    with col1:
        window_name = st.text_input("Window Name", value="", key="new_window_name_heating")
        
        orientation = st.selectbox(
            "Orientation",
            options=["north", "east", "south", "west", "horizontal"],
            key="new_window_orientation_heating"
        )
        
        glass_type = st.selectbox(
            "Glass Type",
            options=list(glass_type_options.keys()),
            format_func=lambda x: glass_type_options[x],
            key="new_window_glass_type_heating"
        )
        
        # Get glass properties
        glass_data = ref_data.get_glass_type(glass_type)
        window_u_value = glass_data['u_value']
        
    with col2:
        window_area = st.number_input(
            "Window Area (m²)",
            value=2.0,
            min_value=0.1,
            step=0.1,
            key="new_window_area_heating"
        )
        
        st.write(f"Glass U-Value: {window_u_value} W/m²°C")
        st.write(f"Heat Loss: {window_u_value * window_area * temp_diff:.2f} W")
    
    # Add window button
    if st.button("Add Window", key="add_window_heating"):
        new_window = {
            'name': window_name if window_name else f"Window {len(st.session_state.heating_form_data['windows']['windows']) + 1}",
            'orientation': orientation,
            'glass_type': glass_type,
            'area': window_area,
            'u_value': window_u_value,
            'temp_diff': temp_diff
        }
        st.session_state.heating_form_data['windows']['windows'].append(new_window)
        st.experimental_rerun()
    
    # Doors section
    st.write("### Doors")
    
    # Display existing door entries
    if st.session_state.heating_form_data['windows']['doors']:
        st.write("Current doors:")
        doors_df = pd.DataFrame(st.session_state.heating_form_data['windows']['doors'])
        doors_df = doors_df[['name', 'type', 'area', 'u_value']]
        doors_df.columns = ['Name', 'Type', 'Area (m²)', 'U-Value (W/m²°C)']
        st.dataframe(doors_df)
    
    # Add new door form
    st.write("Add a new door:")
    
    col1, col2 = st.columns(2)
    
    with col1:
        door_name = st.text_input("Door Name", value="", key="new_door_name_heating")
        
        door_type = st.selectbox(
            "Door Type",
            options=["Solid wood", "Hollow core", "Glass", "Insulated"],
            key="new_door_type_heating"
        )
        
        # Set U-value based on door type
        door_u_values = {
            "Solid wood": 2.0,
            "Hollow core": 2.5,
            "Glass": 5.0,
            "Insulated": 1.2
        }
        door_u_value = door_u_values[door_type]
        
    with col2:
        door_area = st.number_input(
            "Door Area (m²)",
            value=2.0,
            min_value=0.1,
            step=0.1,
            key="new_door_area_heating"
        )
        
        st.write(f"Door U-Value: {door_u_value} W/m²°C")
        st.write(f"Heat Loss: {door_u_value * door_area * temp_diff:.2f} W")
    
    # Add door button
    if st.button("Add Door", key="add_door_heating"):
        new_door = {
            'name': door_name if door_name else f"Door {len(st.session_state.heating_form_data['windows']['doors']) + 1}",
            'type': door_type,
            'area': door_area,
            'u_value': door_u_value,
            'temp_diff': temp_diff
        }
        st.session_state.heating_form_data['windows']['doors'].append(new_door)
        st.experimental_rerun()
    
    # Validate inputs
    warnings = []
    
    # Check if windows are defined
    if not st.session_state.heating_form_data['windows']['windows']:
        warnings.append(ValidationWarning(
            "No windows defined",
            "Add at least one window to continue"
        ))
    
    # Check window-to-wall ratio
    if st.session_state.heating_form_data['windows']['windows']:
        total_window_area = sum(window['area'] for window in st.session_state.heating_form_data['windows']['windows'])
        total_wall_area = sum(wall['area'] for wall in st.session_state.heating_form_data['building_envelope']['walls'])
        window_wall_ratio = total_window_area / total_wall_area if total_wall_area > 0 else 0
        
        if window_wall_ratio > 0.6:
            warnings.append(ValidationWarning(
                "High window-to-wall ratio",
                f"Window-to-wall ratio is {window_wall_ratio:.2f}, which is unusually high. Typical ratios are 0.2-0.4."
            ))
    
    # Save warnings to session state
    st.session_state.heating_warnings['windows'] = warnings
    
    # Display warnings if any
    if warnings:
        st.warning("Please review the following warnings:")
        for warning in warnings:
            st.write(f"- {warning.message}" + (" (Critical)" if warning.is_critical else ""))
            st.write(f"  Suggestion: {warning.suggestion}")
    
    # Mark this step as completed if there are no critical warnings
    st.session_state.heating_completed['windows'] = not any(w.is_critical for w in warnings)
    
    # Navigation buttons
    col1, col2 = st.columns([1, 1])
    
    with col1:
        prev_button = st.button("← Back: Building Envelope", key="heating_windows_prev")
        if prev_button:
            st.session_state.heating_active_tab = "building_envelope"
            st.experimental_rerun()
    
    with col2:
        next_button = st.button("Next: Ventilation →", key="heating_windows_next")
        if next_button:
            st.session_state.heating_active_tab = "ventilation"
            st.experimental_rerun()


def ventilation_form(ref_data):
    """
    Form for ventilation and infiltration information.
    
    Args:
        ref_data: Reference data object
    """
    st.subheader("Ventilation & Infiltration")
    st.write("Enter information about ventilation and infiltration rates.")
    
    # Get building info
    building_info = st.session_state.heating_form_data['building_info']
    volume = building_info.get('volume', 216.0)
    temp_diff = building_info.get('temp_diff', 16.5)
    
    # Initialize ventilation data if not already in session state
    if 'infiltration' not in st.session_state.heating_form_data['ventilation']:
        st.session_state.heating_form_data['ventilation']['infiltration'] = {
            'air_changes': 0.5
        }
    
    if 'ventilation' not in st.session_state.heating_form_data['ventilation']:
        st.session_state.heating_form_data['ventilation']['ventilation'] = {
            'type': 'natural',
            'air_changes': 0.0
        }
    
    # Initialize internal loads data if not already in session state
    if 'internal_loads' not in st.session_state.heating_form_data:
        st.session_state.heating_form_data['internal_loads'] = {}
    
    if 'occupants' not in st.session_state.heating_form_data['internal_loads']:
        st.session_state.heating_form_data['internal_loads']['occupants'] = {
            'count': 4,
            'activity_level': 'seated_resting'
        }
    
    if 'lighting' not in st.session_state.heating_form_data['internal_loads']:
        st.session_state.heating_form_data['internal_loads']['lighting'] = {
            'type': 'led',
            'power_density': 5.0  # W/m²
        }
    
    if 'appliances' not in st.session_state.heating_form_data['internal_loads']:
        st.session_state.heating_form_data['internal_loads']['appliances'] = {
            'kitchen': True,
            'living_room': True,
            'bedroom': True,
            'office': False
        }
    
    # Infiltration section
    st.write("### Infiltration")
    st.write("Infiltration is the unintended air leakage through the building envelope.")
    
    infiltration_ach = st.slider(
        "Infiltration Rate (air changes per hour)",
        value=float(st.session_state.heating_form_data['ventilation']['infiltration'].get('air_changes', 0.5)),
        min_value=0.1,
        max_value=2.0,
        step=0.1,
        help="Typical values: 0.5 ACH for modern construction, 1.0 ACH for average construction, 1.5+ ACH for older buildings",
        key="infiltration_ach_heating"
    )
    
    # Calculate infiltration heat loss
    infiltration_heat_loss = 0.33 * volume * infiltration_ach * temp_diff
    
    st.write(f"Infiltration heat loss: {infiltration_heat_loss:.2f} W")
    
    # Save infiltration data
    st.session_state.heating_form_data['ventilation']['infiltration'] = {
        'air_changes': infiltration_ach,
        'volume': volume,
        'temp_diff': temp_diff,
        'heat_loss': infiltration_heat_loss
    }
    
    # Ventilation section
    st.write("### Ventilation")
    st.write("Ventilation is the intentional introduction of outside air into the building.")
    
    # Internal Loads section
    st.write("### Internal Loads")
    st.write("Internal loads are heat sources inside the building that reduce heating requirements.")
    
    # Occupants section
    st.write("#### Occupants")
    
    col1, col2 = st.columns(2)
    
    with col1:
        occupant_count = st.number_input(
            "Number of Occupants",
            value=int(st.session_state.heating_form_data['internal_loads']['occupants'].get('count', 4)),
            min_value=1,
            step=1,
            key="occupant_count_heating"
        )
    
    with col2:
        # Get activity level options from reference data
        activity_options = {act_id: act_data['name'] for act_id, act_data in ref_data.internal_loads['people'].items()}
        
        activity_level = st.selectbox(
            "Activity Level",
            options=list(activity_options.keys()),
            format_func=lambda x: activity_options[x],
            index=list(activity_options.keys()).index(st.session_state.heating_form_data['internal_loads']['occupants'].get('activity_level', 'seated_resting')) if st.session_state.heating_form_data['internal_loads']['occupants'].get('activity_level') in activity_options else 0,
            key="activity_level_heating"
        )
    
    # Get heat gain per person
    activity_data = ref_data.get_internal_load('people', activity_level)
    sensible_heat_pp = activity_data['sensible_heat']
    latent_heat_pp = activity_data['latent_heat']
    total_heat_pp = sensible_heat_pp + latent_heat_pp
    
    st.write(f"Heat gain per person: {total_heat_pp} W ({sensible_heat_pp} W sensible + {latent_heat_pp} W latent)")
    st.write(f"Total occupant heat gain: {total_heat_pp * occupant_count} W")
    
    # Save occupants data
    st.session_state.heating_form_data['internal_loads']['occupants'] = {
        'count': occupant_count,
        'activity_level': activity_level,
        'sensible_heat_pp': sensible_heat_pp,
        'latent_heat_pp': latent_heat_pp,
        'total_heat_gain': total_heat_pp * occupant_count
    }
    
    # Lighting section
    st.write("#### Lighting")
    
    col1, col2 = st.columns(2)
    
    with col1:
        # Get lighting type options from reference data
        lighting_options = {light_id: light_data['name'] for light_id, light_data in ref_data.internal_loads['lighting'].items()}
        
        lighting_type = st.selectbox(
            "Lighting Type",
            options=list(lighting_options.keys()),
            format_func=lambda x: lighting_options[x],
            index=list(lighting_options.keys()).index(st.session_state.heating_form_data['internal_loads']['lighting'].get('type', 'led')) if st.session_state.heating_form_data['internal_loads']['lighting'].get('type') in lighting_options else 0,
            key="lighting_type_heating"
        )
    
    with col2:
        lighting_power_density = st.number_input(
            "Lighting Power Density (W/m²)",
            value=float(st.session_state.heating_form_data['internal_loads']['lighting'].get('power_density', 5.0)),
            min_value=1.0,
            max_value=20.0,
            step=0.5,
            help="Typical values: Residential 5-10 W/m², Office 10-15 W/m²",
            key="lighting_power_density_heating"
        )
    
    # Get lighting heat factor
    lighting_data = ref_data.get_internal_load('lighting', lighting_type)
    lighting_heat_factor = lighting_data['heat_factor']
    
    # Calculate lighting heat gain
    floor_area = st.session_state.heating_form_data['building_info'].get('floor_area', 80.0)
    lighting_heat_gain = lighting_power_density * floor_area * lighting_heat_factor
    
    st.write(f"Lighting heat factor: {lighting_heat_factor}")
    st.write(f"Total lighting heat gain: {lighting_heat_gain:.2f} W")
    
    # Save lighting data
    st.session_state.heating_form_data['internal_loads']['lighting'] = {
        'type': lighting_type,
        'power_density': lighting_power_density,
        'heat_factor': lighting_heat_factor,
        'total_heat_gain': lighting_heat_gain
    }
    
    # Equipment section
    st.write("#### Equipment")
    st.write("Select the equipment present in your space:")
    
    col1, col2 = st.columns(2)
    
    with col1:
        has_kitchen = st.checkbox(
            "Kitchen Appliances",
            value=st.session_state.heating_form_data['internal_loads']['appliances'].get('kitchen', True),
            help="Refrigerator, stove, microwave, etc.",
            key="has_kitchen_heating"
        )
        
        has_living_room = st.checkbox(
            "Living Room Equipment",
            value=st.session_state.heating_form_data['internal_loads']['appliances'].get('living_room', True),
            help="TV, audio equipment, etc.",
            key="has_living_room_heating"
        )
    
    with col2:
        has_bedroom = st.checkbox(
            "Bedroom Equipment",
            value=st.session_state.heating_form_data['internal_loads']['appliances'].get('bedroom', True),
            help="TV, chargers, etc.",
            key="has_bedroom_heating"
        )
        
        has_office = st.checkbox(
            "Office Equipment",
            value=st.session_state.heating_form_data['internal_loads']['appliances'].get('office', False),
            help="Computer, printer, etc.",
            key="has_office_heating"
        )
    
    # Calculate equipment heat gain
    equipment_watts = 0
    
    if has_kitchen:
        equipment_watts += 1000  # Kitchen appliances
    if has_living_room:
        equipment_watts += 300   # Living room equipment
    if has_bedroom:
        equipment_watts += 150   # Bedroom equipment
    if has_office:
        equipment_watts += 450   # Office equipment
    
    st.write(f"Total equipment heat gain: {equipment_watts} W")
    
    # Save appliances data
    st.session_state.heating_form_data['internal_loads']['appliances'] = {
        'kitchen': has_kitchen,
        'living_room': has_living_room,
        'bedroom': has_bedroom,
        'office': has_office,
        'total_heat_gain': equipment_watts
    }
    
    # Calculate total internal heat gain
    total_internal_gain = (
        st.session_state.heating_form_data['internal_loads']['occupants']['total_heat_gain'] +
        st.session_state.heating_form_data['internal_loads']['lighting']['total_heat_gain'] +
        st.session_state.heating_form_data['internal_loads']['appliances']['total_heat_gain']
    )
    
    st.write(f"Total internal heat gain: {total_internal_gain:.2f} W")
    
    # Save total internal gain
    st.session_state.heating_form_data['internal_loads']['total_heat_gain'] = total_internal_gain
    
    col1, col2 = st.columns(2)
    
    with col1:
        ventilation_type = st.selectbox(
            "Ventilation Type",
            options=["natural", "mechanical", "mixed"],
            format_func=lambda x: x.capitalize(),
            index=["natural", "mechanical", "mixed"].index(st.session_state.heating_form_data['ventilation']['ventilation'].get('type', 'natural')),
            key="ventilation_type_heating"
        )
    
    with col2:
        ventilation_ach = st.number_input(
            "Ventilation Rate (air changes per hour)",
            value=float(st.session_state.heating_form_data['ventilation']['ventilation'].get('air_changes', 0.0)),
            min_value=0.0,
            max_value=5.0,
            step=0.1,
            help="Typical values: 0.35-1.0 ACH for residential buildings",
            key="ventilation_ach_heating"
        )
    
    # Calculate ventilation heat loss
    ventilation_heat_loss = 0.33 * volume * ventilation_ach * temp_diff
    
    st.write(f"Ventilation heat loss: {ventilation_heat_loss:.2f} W")
    
    # Save ventilation data
    st.session_state.heating_form_data['ventilation']['ventilation'] = {
        'type': ventilation_type,
        'air_changes': ventilation_ach,
        'volume': volume,
        'temp_diff': temp_diff,
        'heat_loss': ventilation_heat_loss
    }
    
    # Calculate total ventilation and infiltration heat loss
    total_ventilation_loss = infiltration_heat_loss + ventilation_heat_loss
    
    st.info(f"Total Ventilation & Infiltration Heat Loss: {total_ventilation_loss:.2f} W")
    
    # Save total ventilation loss
    st.session_state.heating_form_data['ventilation']['total_loss'] = total_ventilation_loss
    
    # Validate inputs
    warnings = []
    
    # Check if infiltration rate is reasonable
    if infiltration_ach < 0.3:
        warnings.append(ValidationWarning(
            "Low infiltration rate",
            "Infiltration rate below 0.3 ACH is unusually low for most buildings."
        ))
    elif infiltration_ach > 1.5:
        warnings.append(ValidationWarning(
            "High infiltration rate",
            "Infiltration rate above 1.5 ACH indicates a leaky building envelope."
        ))
    
    # Check if ventilation rate is reasonable
    if ventilation_ach > 0 and ventilation_ach < 0.35:
        warnings.append(ValidationWarning(
            "Low ventilation rate",
            "Ventilation rate below 0.35 ACH may not provide adequate fresh air."
        ))
    elif ventilation_ach > 2.0:
        warnings.append(ValidationWarning(
            "High ventilation rate",
            "Ventilation rate above 2.0 ACH is unusually high for residential buildings."
        ))
    
    # Save warnings to session state
    st.session_state.heating_warnings['ventilation'] = warnings
    
    # Display warnings if any
    if warnings:
        st.warning("Please review the following warnings:")
        for warning in warnings:
            st.write(f"- {warning.message}" + (" (Critical)" if warning.is_critical else ""))
            st.write(f"  Suggestion: {warning.suggestion}")
    
    # Mark this step as completed if there are no critical warnings
    st.session_state.heating_completed['ventilation'] = not any(w.is_critical for w in warnings)
    
    # Navigation buttons
    col1, col2 = st.columns([1, 1])
    
    with col1:
        prev_button = st.button("← Back: Windows & Doors", key="heating_ventilation_prev")
        if prev_button:
            st.session_state.heating_active_tab = "windows"
            st.experimental_rerun()
    
    with col2:
        next_button = st.button("Next: Occupancy →", key="heating_ventilation_next")
        if next_button:
            st.session_state.heating_active_tab = "occupancy"
            st.experimental_rerun()


def occupancy_form(ref_data):
    """
    Form for occupancy information.
    
    Args:
        ref_data: Reference data object
    """
    st.subheader("Occupancy Information")
    st.write("Enter information about occupancy patterns and heating degree days.")
    
    # Get location from building info
    location = st.session_state.heating_form_data['building_info'].get('location', 'sydney')
    location_name = st.session_state.heating_form_data['building_info'].get('location_name', 'Sydney')
    
    # Initialize occupancy data if not already in session state
    if 'occupancy_type' not in st.session_state.heating_form_data['occupancy']:
        st.session_state.heating_form_data['occupancy']['occupancy_type'] = 'continuous'
    
    if 'heating_degree_days' not in st.session_state.heating_form_data['occupancy']:
        # Get default HDD from reference data
        calculator = HeatingLoadCalculator()
        default_hdd = calculator.get_heating_degree_days(location)
        st.session_state.heating_form_data['occupancy']['heating_degree_days'] = default_hdd
    
    # Occupancy section
    st.write("### Occupancy Pattern")
    
    # Get occupancy options from reference data
    occupancy_options = {occ_id: occ_data['name'] for occ_id, occ_data in ref_data.occupancy_factors.items()}
    
    occupancy_type = st.selectbox(
        "Occupancy Type",
        options=list(occupancy_options.keys()),
        format_func=lambda x: occupancy_options[x],
        index=list(occupancy_options.keys()).index(st.session_state.heating_form_data['occupancy'].get('occupancy_type', 'continuous')) if st.session_state.heating_form_data['occupancy'].get('occupancy_type') in occupancy_options else 0,
        help="Select the occupancy pattern that best describes how the building is used"
    )
    
    # Get occupancy factor
    occupancy_data = ref_data.get_occupancy_factor(occupancy_type)
    occupancy_factor = occupancy_data['factor']
    
    st.write(f"Occupancy correction factor: {occupancy_factor}")
    st.write(f"Description: {occupancy_data['description']}")
    
    # Save occupancy data
    st.session_state.heating_form_data['occupancy']['occupancy_type'] = occupancy_type
    st.session_state.heating_form_data['occupancy']['occupancy_factor'] = occupancy_factor
    
    # Heating degree days section
    st.write("### Heating Degree Days")
    st.write("Heating degree days are used to estimate annual heating energy requirements.")
    
    col1, col2 = st.columns(2)
    
    with col1:
        base_temp = st.selectbox(
            "Base Temperature",
            options=[18, 15.5, 12],
            index=[18, 15.5, 12].index(st.session_state.heating_form_data['occupancy'].get('base_temp', 18)) if st.session_state.heating_form_data['occupancy'].get('base_temp') in [18, 15.5, 12] else 0,
            help="Base temperature for heating degree days calculation"
        )
    
    with col2:
        # Get default HDD from reference data
        calculator = HeatingLoadCalculator()
        default_hdd = calculator.get_heating_degree_days(location, base_temp)
        
        heating_degree_days = st.number_input(
            "Heating Degree Days",
            value=float(st.session_state.heating_form_data['occupancy'].get('heating_degree_days', default_hdd)),
            min_value=0.0,
            step=10.0,
            help=f"Default value for {location_name} at base {base_temp}°C: {default_hdd}"
        )
    
    st.write(f"Heating degree days represent the sum of daily temperature differences between the base temperature and the average daily temperature when it falls below the base temperature.")
    
    # Save heating degree days data
    st.session_state.heating_form_data['occupancy']['base_temp'] = base_temp
    st.session_state.heating_form_data['occupancy']['heating_degree_days'] = heating_degree_days
    
    # Validate inputs
    warnings = []
    
    # Check if heating degree days are reasonable
    if heating_degree_days == 0:
        warnings.append(ValidationWarning(
            "Zero heating degree days",
            "With zero heating degree days, annual heating energy will be zero."
        ))
    elif heating_degree_days < 100 and base_temp == 18:
        warnings.append(ValidationWarning(
            "Very low heating degree days",
            f"Heating degree days below 100 at base {base_temp}°C is unusually low for most locations."
        ))
    elif heating_degree_days > 3000:
        warnings.append(ValidationWarning(
            "Very high heating degree days",
            "Heating degree days above 3000 is unusually high for most locations."
        ))
    
    # Save warnings to session state
    st.session_state.heating_warnings['occupancy'] = warnings
    
    # Display warnings if any
    if warnings:
        st.warning("Please review the following warnings:")
        for warning in warnings:
            st.write(f"- {warning.message}" + (" (Critical)" if warning.is_critical else ""))
            st.write(f"  Suggestion: {warning.suggestion}")
    
    # Mark this step as completed if there are no critical warnings
    st.session_state.heating_completed['occupancy'] = not any(w.is_critical for w in warnings)
    
    # Navigation buttons
    col1, col2 = st.columns([1, 1])
    
    with col1:
        prev_button = st.button("← Back: Ventilation", key="heating_occupancy_prev")
        if prev_button:
            st.session_state.heating_active_tab = "ventilation"
            st.experimental_rerun()
    
    with col2:
        calculate_button = st.button("Calculate Results →", key="heating_occupancy_calculate")
        if calculate_button:
            # Calculate heating load
            calculate_heating_load()
            st.session_state.heating_active_tab = "results"
            st.experimental_rerun()


def calculate_heating_load():
    """Calculate heating load based on input data."""
    # Create calculator instance
    calculator = HeatingLoadCalculator()
    
    # Get form data
    form_data = st.session_state.heating_form_data
    
    # Prepare building components for calculation
    building_components = []
    
    # Add walls
    for wall in form_data['building_envelope'].get('walls', []):
        building_components.append({
            'name': wall['name'],
            'area': wall['area'],
            'u_value': wall['u_value'],
            'temp_diff': wall['temp_diff']
        })
    
    # Add roof
    roof = form_data['building_envelope'].get('roof', {})
    if roof:
        building_components.append({
            'name': 'Roof',
            'area': roof['area'],
            'u_value': roof['u_value'],
            'temp_diff': roof['temp_diff']
        })
    
    # Add floor
    floor = form_data['building_envelope'].get('floor', {})
    if floor:
        building_components.append({
            'name': 'Floor',
            'area': floor['area'],
            'u_value': floor['u_value'],
            'temp_diff': floor['temp_diff']
        })
    
    # Add windows
    for window in form_data['windows'].get('windows', []):
        building_components.append({
            'name': window['name'],
            'area': window['area'],
            'u_value': window['u_value'],
            'temp_diff': window['temp_diff']
        })
    
    # Add doors
    for door in form_data['windows'].get('doors', []):
        building_components.append({
            'name': door['name'],
            'area': door['area'],
            'u_value': door['u_value'],
            'temp_diff': door['temp_diff']
        })
    
    # Prepare infiltration data
    infiltration = form_data['ventilation'].get('infiltration', {})
    ventilation = form_data['ventilation'].get('ventilation', {})
    
    infiltration_data = {
        'volume': infiltration.get('volume', 0),
        'air_changes': infiltration.get('air_changes', 0) + ventilation.get('air_changes', 0),
        'temp_diff': infiltration.get('temp_diff', 0)
    }
    
    # Prepare internal loads data
    internal_loads = None
    if 'internal_loads' in form_data:
        internal_loads = {
            'num_people': form_data['internal_loads']['occupants'].get('count', 0),
            'has_kitchen': form_data['internal_loads']['appliances'].get('kitchen', False),
            'equipment_watts': form_data['internal_loads']['appliances'].get('total_heat_gain', 0)
        }
    
    # Calculate heating load
    results = calculator.calculate_total_heating_load(
        building_components=building_components,
        infiltration=infiltration_data,
        internal_gains=internal_loads
    )
    
    # Calculate annual heating requirement
    location = form_data['building_info'].get('location', 'sydney')
    occupancy_type = form_data['occupancy'].get('occupancy_type', 'continuous')
    base_temp = form_data['occupancy'].get('base_temp', 18)
    
    annual_results = calculator.calculate_annual_heating_requirement(
        results['total_load'],
        location,
        occupancy_type,
        base_temp
    )
    
    # Combine results
    combined_results = {**results, **annual_results}
    
    # Save results to session state
    st.session_state.heating_results = combined_results
    
    # Add timestamp
    st.session_state.heating_results['timestamp'] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    # Add building info
    st.session_state.heating_results['building_info'] = form_data['building_info']
    
    return combined_results


def results_page():
    """Display calculation results."""
    st.subheader("Heating Load Calculation Results")
    
    # Check if results are available
    if not st.session_state.heating_results:
        st.warning("No calculation results available. Please complete the input forms and calculate results.")
        return
    
    # Get results
    results = st.session_state.heating_results
    
    # Display summary
    st.write("### Summary")
    
    col1, col2 = st.columns(2)
    
    with col1:
        st.metric("Total Heating Load", f"{results['total_load']:.2f} W")
        
        # Convert to kW
        total_load_kw = results['total_load'] / 1000
        st.metric("Total Heating Load", f"{total_load_kw:.2f} kW")
        
        # Annual heating energy
        st.metric("Annual Heating Energy", f"{results['annual_energy_kwh']:.2f} kWh")
    
    with col2:
        # Calculate heating load per area
        floor_area = results['building_info'].get('floor_area', 80.0)
        heating_load_per_area = results['total_load'] / floor_area
        st.metric("Heating Load per Area", f"{heating_load_per_area:.2f} W/m²")
        
        # Annual heating energy per area
        annual_energy_per_area = results['annual_energy_kwh'] / floor_area
        st.metric("Annual Heating Energy per Area", f"{annual_energy_per_area:.2f} kWh/m²")
        
        # Equipment sizing recommendation
        # Add 10% safety factor
        recommended_size = total_load_kw * 1.1
        st.metric("Recommended Equipment Size", f"{recommended_size:.2f} kW")
    
    # Display load breakdown
    st.write("### Load Breakdown")
    
    # Prepare data for pie chart
    component_losses = results['component_losses']
    
    # Create pie chart for component losses
    fig = px.pie(
        values=list(component_losses.values()),
        names=list(component_losses.keys()),
        title="Heating Load Components",
        color_discrete_sequence=px.colors.sequential.Viridis,
        hole=0.4,  # Create a donut chart for better readability
        labels={'label': 'Component', 'value': 'Heat Loss (W)'}
    )
    
    # Improve layout and formatting
    fig.update_traces(
        textposition='inside',
        textinfo='percent+label',
        hoverinfo='label+percent+value',
        marker=dict(line=dict(color='#FFFFFF', width=2))
    )
    
    # Improve layout
    fig.update_layout(
        legend_title_text='Building Components',
        font=dict(size=14),
        title_font=dict(size=18),
        title_x=0.5,  # Center the title
        margin=dict(t=50, b=50, l=50, r=50)
    )
    
    st.plotly_chart(fig)
    
    # Display load components in a table
    load_components = {
        'Conduction (Building Envelope)': results['total_conduction_loss'] - results.get('infiltration_loss', 0),
        'Infiltration & Ventilation': results.get('infiltration_loss', 0)
    }
    
    # Add internal gains and solar gains if available
    if 'internal_gain' in results and results['internal_gain'] > 0:
        load_components['Internal Gains (reduction)'] = -results['internal_gain']
    
    if 'wall_solar_gain' in results and results['wall_solar_gain'] > 0:
        load_components['Solar Gains (reduction)'] = -results['wall_solar_gain']
    
    load_df = pd.DataFrame({
        'Component': list(load_components.keys()),
        'Load (W)': list(load_components.values()),
        'Percentage (%)': [abs(value) / results['total_load'] * 100 for value in load_components.values()]
    })
    
    st.dataframe(load_df.style.format({
        'Load (W)': '{:.2f}',
        'Percentage (%)': '{:.2f}'
    }))
    
    # Display detailed results
    st.write("### Detailed Results")
    
    # Create tabs for different result sections
    tabs = st.tabs([
        "Building Components", 
        "Ventilation", 
        "Annual Energy"
    ])
    
    with tabs[0]:
        st.subheader("Building Component Heat Losses")
        
        # Create dataframe from component losses
        components_data = []
        for name, loss in component_losses.items():
            # Find the component in the original data to get area and U-value
            component = None
            for comp in st.session_state.heating_form_data['building_envelope'].get('walls', []):
                if comp['name'] == name:
                    component = comp
                    break
            
            if name == 'Roof':
                component = st.session_state.heating_form_data['building_envelope'].get('roof', {})
            elif name == 'Floor':
                component = st.session_state.heating_form_data['building_envelope'].get('floor', {})
            
            # Check windows and doors
            if not component:
                for window in st.session_state.heating_form_data['windows'].get('windows', []):
                    if window['name'] == name:
                        component = window
                        break
            
            if not component:
                for door in st.session_state.heating_form_data['windows'].get('doors', []):
                    if door['name'] == name:
                        component = door
                        break
            
            if component:
                components_data.append({
                    'Component': name,
                    'Area (m²)': component.get('area', 0),
                    'U-Value (W/m²°C)': component.get('u_value', 0),
                    'Temperature Difference (°C)': component.get('temp_diff', 0),
                    'Heat Loss (W)': loss
                })
            else:
                components_data.append({
                    'Component': name,
                    'Area (m²)': 0,
                    'U-Value (W/m²°C)': 0,
                    'Temperature Difference (°C)': 0,
                    'Heat Loss (W)': loss
                })
        
        # Create dataframe
        components_df = pd.DataFrame(components_data)
        
        # Display table
        st.dataframe(components_df.style.format({
            'Area (m²)': '{:.2f}',
            'U-Value (W/m²°C)': '{:.2f}',
            'Temperature Difference (°C)': '{:.2f}',
            'Heat Loss (W)': '{:.2f}'
        }))
        
        # Create bar chart
        fig = px.bar(
            components_df,
            x='Component',
            y='Heat Loss (W)',
            title="Heat Loss by Building Component",
            color='Component',
            color_discrete_sequence=px.colors.sequential.Viridis,
            text='Heat Loss (W)'
        )
        
        # Improve layout and formatting
        fig.update_traces(
            texttemplate='%{text:.1f} W',
            textposition='outside',
            hovertemplate='<b>%{x}</b><br>Heat Loss: %{y:.1f} W<extra></extra>'
        )
        
        # Improve layout
        fig.update_layout(
            xaxis_title="Building Component",
            yaxis_title="Heat Loss (W)",
            font=dict(size=14),
            title_font=dict(size=18),
            title_x=0.5,  # Center the title
            margin=dict(t=50, b=50, l=50, r=50),
            xaxis={'categoryorder':'total descending'}  # Sort by highest heat loss
        )
        
        st.plotly_chart(fig)
    
    with tabs[1]:
        st.subheader("Ventilation & Infiltration Heat Losses")
        
        # Get ventilation data
        ventilation_data = st.session_state.heating_form_data['ventilation']
        
        # Create dataframe
        ventilation_df = pd.DataFrame([
            {
                'Source': 'Infiltration',
                'Air Changes per Hour': ventilation_data['infiltration']['air_changes'],
                'Volume (m³)': ventilation_data['infiltration']['volume'],
                'Temperature Difference (°C)': ventilation_data['infiltration']['temp_diff'],
                'Heat Loss (W)': ventilation_data['infiltration']['heat_loss']
            },
            {
                'Source': 'Ventilation',
                'Air Changes per Hour': ventilation_data['ventilation']['air_changes'],
                'Volume (m³)': ventilation_data['ventilation']['volume'],
                'Temperature Difference (°C)': ventilation_data['ventilation']['temp_diff'],
                'Heat Loss (W)': ventilation_data['ventilation']['heat_loss']
            }
        ])
        
        # Display table
        st.dataframe(ventilation_df.style.format({
            'Air Changes per Hour': '{:.2f}',
            'Volume (m³)': '{:.2f}',
            'Temperature Difference (°C)': '{:.2f}',
            'Heat Loss (W)': '{:.2f}'
        }))
        
        # Create bar chart
        fig = px.bar(
            ventilation_df,
            x='Source',
            y='Heat Loss (W)',
            title="Ventilation & Infiltration Heat Losses",
            color='Source',
            color_discrete_sequence=px.colors.sequential.Plasma,
            text='Heat Loss (W)'
        )
        
        # Improve layout and formatting
        fig.update_traces(
            texttemplate='%{text:.1f} W',
            textposition='outside',
            hovertemplate='<b>%{x}</b><br>Heat Loss: %{y:.1f} W<br>Air Changes: %{customdata[0]:.2f} ACH<extra></extra>'
        )
        
        # Add custom data for hover
        fig.update_traces(customdata=ventilation_df[['Air Changes per Hour']])
        
        # Improve layout
        fig.update_layout(
            xaxis_title="Ventilation Source",
            yaxis_title="Heat Loss (W)",
            font=dict(size=14),
            title_font=dict(size=18),
            title_x=0.5,  # Center the title
            margin=dict(t=50, b=50, l=50, r=50)
        )
        
        st.plotly_chart(fig)
    
    with tabs[2]:
        st.subheader("Annual Heating Energy")
        
        # Get occupancy data
        occupancy_data = st.session_state.heating_form_data['occupancy']
        
        # Create dataframe
        annual_data = pd.DataFrame([
            {
                'Parameter': 'Heating Degree Days',
                'Value': results['heating_degree_days'],
                'Unit': 'HDD'
            },
            {
                'Parameter': 'Base Temperature',
                'Value': occupancy_data['base_temp'],
                'Unit': '°C'
            },
            {
                'Parameter': 'Occupancy Type',
                'Value': occupancy_data['occupancy_type'].capitalize(),
                'Unit': ''
            },
            {
                'Parameter': 'Correction Factor',
                'Value': results['correction_factor'],
                'Unit': ''
            },
            {
                'Parameter': 'Annual Heating Energy',
                'Value': results['annual_energy_kwh'],
                'Unit': 'kWh'
            },
            {
                'Parameter': 'Annual Heating Energy',
                'Value': results['annual_energy_mj'],
                'Unit': 'MJ'
            }
        ])
        
        # Display table
        st.dataframe(annual_data.style.format({
            'Value': lambda x: f"{x:.2f}" if isinstance(x, (int, float)) else str(x)
        }))
        
        # Create bar chart for monthly distribution (estimated)
        # This is a simplified distribution based on heating degree days
        months = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec']
        
        # Get location
        location = st.session_state.heating_form_data['building_info'].get('location', 'sydney')
        
        # Simplified monthly distribution factors based on hemisphere
        # Southern hemisphere: winter is June-August
        # Northern hemisphere: winter is December-February
        southern_hemisphere = ['sydney', 'melbourne', 'brisbane', 'perth', 'adelaide', 'hobart', 'darwin', 'canberra', 'mildura']
        
        if location.lower() in southern_hemisphere:
            # Southern hemisphere distribution
            monthly_factors = [0.02, 0.01, 0.03, 0.08, 0.12, 0.16, 0.18, 0.16, 0.12, 0.08, 0.03, 0.01]
        else:
            # Northern hemisphere distribution
            monthly_factors = [0.18, 0.16, 0.12, 0.08, 0.03, 0.01, 0.01, 0.01, 0.03, 0.08, 0.12, 0.17]
        
        # Calculate monthly energy
        monthly_energy = [results['annual_energy_kwh'] * factor for factor in monthly_factors]
        
        # Create dataframe
        monthly_df = pd.DataFrame({
            'Month': months,
            'Energy (kWh)': monthly_energy
        })
        
        # Create bar chart
        fig = px.bar(
            monthly_df,
            x='Month',
            y='Energy (kWh)',
            title="Estimated Monthly Heating Energy Distribution",
            color_discrete_sequence=['indianred']
        )
        
        st.plotly_chart(fig)
    
    # Export options
    st.write("### Export Options")
    
    col1, col2 = st.columns(2)
    
    with col1:
        if st.button("Export Results as CSV", key="export_csv_heating"):
            # Create a CSV file with results
            csv_data = export_data(st.session_state.heating_form_data, st.session_state.heating_results, format='csv')
            
            # Provide download link
            st.download_button(
                label="Download CSV",
                data=csv_data,
                file_name=f"heating_load_results_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv",
                mime="text/csv"
            )
    
    with col2:
        if st.button("Export Results as JSON", key="export_json_heating"):
            # Create a JSON file with results
            json_data = export_data(st.session_state.heating_form_data, st.session_state.heating_results, format='json')
            
            # Provide download link
            st.download_button(
                label="Download JSON",
                data=json_data,
                file_name=f"heating_load_results_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
                mime="application/json"
            )
    
    # Navigation buttons
    col1, col2 = st.columns([1, 1])
    
    with col1:
        prev_button = st.button("← Back: Occupancy", key="heating_results_prev")
        if prev_button:
            st.session_state.heating_active_tab = "occupancy"
            st.experimental_rerun()
    
    with col2:
        recalculate_button = st.button("Recalculate", key="heating_results_recalculate")
        if recalculate_button:
            # Recalculate heating load
            calculate_heating_load()
            st.experimental_rerun()


def heating_calculator():
    """Main function for the heating load calculator page."""
    st.title("Heating Load Calculator")
    
    # Initialize reference data
    ref_data = ReferenceData()
    
    # Initialize session state
    load_session_state()
    
    # Initialize active tab if not already set
    if 'heating_active_tab' not in st.session_state:
        st.session_state.heating_active_tab = "building_info"
    
    # Create tabs for different steps
    tabs = st.tabs([
        "1. Building Information", 
        "2. Building Envelope", 
        "3. Windows & Doors", 
        "4. Ventilation", 
        "5. Occupancy", 
        "6. Results"
    ])
    
    # Add direct navigation buttons at the top
    st.write("### Navigation")
    st.write("Click on any button below to navigate directly to that section:")
    
    col1, col2, col3 = st.columns(3)
    with col1:
        if st.button("1. Building Information", key="direct_nav_heating_info"):
            st.session_state.heating_active_tab = "building_info"
            st.experimental_rerun()
        
        if st.button("2. Building Envelope", key="direct_nav_heating_envelope"):
            st.session_state.heating_active_tab = "building_envelope"
            st.experimental_rerun()
    
    with col2:
        if st.button("3. Windows & Doors", key="direct_nav_heating_windows"):
            st.session_state.heating_active_tab = "windows"
            st.experimental_rerun()
        
        if st.button("4. Ventilation", key="direct_nav_heating_ventilation"):
            st.session_state.heating_active_tab = "ventilation"
            st.experimental_rerun()
    
    with col3:
        if st.button("5. Occupancy", key="direct_nav_heating_occupancy"):
            st.session_state.heating_active_tab = "occupancy"
            st.experimental_rerun()
        
        if st.button("6. Results", key="direct_nav_heating_results"):
            # Only enable if all previous steps are completed
            if all(st.session_state.heating_completed.values()):
                st.session_state.heating_active_tab = "results"
                st.experimental_rerun()
            else:
                st.warning("Please complete all previous steps before viewing results.")
    
    # Display the active tab
    with tabs[0]:
        if st.session_state.heating_active_tab == "building_info":
            building_info_form(ref_data)
    
    with tabs[1]:
        if st.session_state.heating_active_tab == "building_envelope":
            building_envelope_form(ref_data)
    
    with tabs[2]:
        if st.session_state.heating_active_tab == "windows":
            windows_form(ref_data)
    
    with tabs[3]:
        if st.session_state.heating_active_tab == "ventilation":
            ventilation_form(ref_data)
    
    with tabs[4]:
        if st.session_state.heating_active_tab == "occupancy":
            occupancy_form(ref_data)
    
    with tabs[5]:
        if st.session_state.heating_active_tab == "results":
            results_page()


if __name__ == "__main__":
    heating_calculator()