"""
Heating load calculation module for HVAC Load Calculator.
This module implements steady-state methods for calculating heating loads.
"""

from typing import Dict, List, Any, Optional, Tuple
import math
import numpy as np
import pandas as pd
import os
from datetime import datetime, timedelta
from enum import Enum

# Import data models and utilities
from data.building_components import Wall, Roof, Floor, Window, Door, Orientation, ComponentType
from utils.psychrometrics import Psychrometrics
from utils.heat_transfer import HeatTransfer

# Define paths
DATA_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


class HeatingLoad:
    """Class for calculating heating loads using steady-state methods."""
    
    def __init__(self):
        """Initialize heating load calculator."""
        self.heat_transfer = HeatTransfer()
        self.psychrometrics = Psychrometrics()
    
    def calculate_wall_heating_load(self, wall: Wall, outdoor_temp: float, indoor_temp: float) -> float:
        """
        Calculate heating load through a wall using steady-state conduction.
        
        Args:
            wall: Wall object
            outdoor_temp: Outdoor temperature in °C
            indoor_temp: Indoor temperature in °C
            
        Returns:
            Heating load in W
        """
        # Get wall properties
        u_value = wall.u_value
        area = wall.area
        
        # Calculate heating load
        delta_t = indoor_temp - outdoor_temp
        heating_load = u_value * area * delta_t
        
        return heating_load
    
    def calculate_roof_heating_load(self, roof: Roof, outdoor_temp: float, indoor_temp: float) -> float:
        """
        Calculate heating load through a roof using steady-state conduction.
        
        Args:
            roof: Roof object
            outdoor_temp: Outdoor temperature in °C
            indoor_temp: Indoor temperature in °C
            
        Returns:
            Heating load in W
        """
        # Get roof properties
        u_value = roof.u_value
        area = roof.area
        
        # Calculate heating load
        delta_t = indoor_temp - outdoor_temp
        heating_load = u_value * area * delta_t
        
        return heating_load
    
    def calculate_floor_heating_load(self, floor: Floor, ground_temp: float, indoor_temp: float) -> float:
        """
        Calculate heating load through a floor.
        
        Args:
            floor: Floor object
            ground_temp: Ground or adjacent space temperature in °C
            indoor_temp: Indoor temperature in °C
            
        Returns:
            Heating load in W
        """
        # Get floor properties
        u_value = floor.u_value
        area = floor.area
        
        # Calculate heating load
        delta_t = indoor_temp - ground_temp
        heating_load = u_value * area * delta_t
        
        return heating_load
    
    def calculate_window_heating_load(self, window: Window, outdoor_temp: float, indoor_temp: float) -> float:
        """
        Calculate heating load through a window using steady-state conduction.
        
        Args:
            window: Window object
            outdoor_temp: Outdoor temperature in °C
            indoor_temp: Indoor temperature in °C
            
        Returns:
            Heating load in W
        """
        # Get window properties
        u_value = window.u_value
        area = window.area
        
        # Calculate heating load
        delta_t = indoor_temp - outdoor_temp
        heating_load = u_value * area * delta_t
        
        return heating_load
    
    def calculate_door_heating_load(self, door: Door, outdoor_temp: float, indoor_temp: float) -> float:
        """
        Calculate heating load through a door using steady-state conduction.
        
        Args:
            door: Door object
            outdoor_temp: Outdoor temperature in °C
            indoor_temp: Indoor temperature in °C
            
        Returns:
            Heating load in W
        """
        # Get door properties
        u_value = door.u_value
        area = door.area
        
        # Calculate heating load
        delta_t = indoor_temp - outdoor_temp
        heating_load = u_value * area * delta_t
        
        return heating_load
    
    def calculate_infiltration_heating_load(self, flow_rate: float, outdoor_temp: float, indoor_temp: float,
                                          outdoor_rh: float, indoor_rh: float) -> Dict[str, float]:
        """
        Calculate sensible and latent heating loads due to infiltration.
        
        Args:
            flow_rate: Infiltration flow rate in m³/s
            outdoor_temp: Outdoor temperature in °C
            indoor_temp: Indoor temperature in °C
            outdoor_rh: Outdoor relative humidity in %
            indoor_rh: Indoor relative humidity in %
            
        Returns:
            Dictionary with sensible, latent, and total heating loads in W
        """
        # Calculate sensible heating load
        sensible_load = self.heat_transfer.infiltration_heat_transfer(
            flow_rate=flow_rate,
            delta_t=indoor_temp - outdoor_temp
        )
        
        # Calculate humidity ratios
        w_outdoor = self.psychrometrics.humidity_ratio(outdoor_temp, outdoor_rh)
        w_indoor = self.psychrometrics.humidity_ratio(indoor_temp, indoor_rh)
        
        # Calculate latent heating load (only if indoor humidity is higher than outdoor)
        delta_w = w_indoor - w_outdoor
        if delta_w > 0:
            latent_load = self.heat_transfer.infiltration_latent_heat_transfer(
                flow_rate=flow_rate,
                delta_w=delta_w
            )
        else:
            latent_load = 0
        
        # Calculate total heating load
        total_load = sensible_load + latent_load
        
        return {
            "sensible": sensible_load,
            "latent": latent_load,
            "total": total_load
        }
    
    def calculate_ventilation_heating_load(self, flow_rate: float, outdoor_temp: float, indoor_temp: float,
                                         outdoor_rh: float, indoor_rh: float) -> Dict[str, float]:
        """
        Calculate sensible and latent heating loads due to ventilation.
        
        Args:
            flow_rate: Ventilation flow rate in m³/s
            outdoor_temp: Outdoor temperature in °C
            indoor_temp: Indoor temperature in °C
            outdoor_rh: Outdoor relative humidity in %
            indoor_rh: Indoor relative humidity in %
            
        Returns:
            Dictionary with sensible, latent, and total heating loads in W
        """
        # Ventilation load calculation is the same as infiltration
        return self.calculate_infiltration_heating_load(
            flow_rate=flow_rate,
            outdoor_temp=outdoor_temp,
            indoor_temp=indoor_temp,
            outdoor_rh=outdoor_rh,
            indoor_rh=indoor_rh
        )
    
    def calculate_internal_gains_offset(self, people_load: float, lights_load: float, 
                                      equipment_load: float, usage_factor: float = 0.7) -> float:
        """
        Calculate internal gains offset for heating load.
        
        Args:
            people_load: Heat gain from people in W
            lights_load: Heat gain from lights in W
            equipment_load: Heat gain from equipment in W
            usage_factor: Usage factor for internal gains (0-1)
            
        Returns:
            Internal gains offset in W
        """
        # Calculate total internal gains
        total_gains = people_load + lights_load + equipment_load
        
        # Apply usage factor
        offset = total_gains * usage_factor
        
        return offset
    
    def calculate_design_heating_load(self, building_components: Dict[str, List[Any]], 
                                    outdoor_conditions: Dict[str, Any],
                                    indoor_conditions: Dict[str, Any],
                                    internal_loads: Dict[str, Any],
                                    safety_factor: float = 1.15) -> Dict[str, float]:
        """
        Calculate design heating load for a building.
        
        Args:
            building_components: Dictionary with lists of building components
            outdoor_conditions: Dictionary with outdoor conditions
            indoor_conditions: Dictionary with indoor conditions
            internal_loads: Dictionary with internal loads
            safety_factor: Safety factor for heating load (default: 1.15)
            
        Returns:
            Dictionary with heating load results
        """
        # Initialize results dictionary
        results = {
            "walls": 0,
            "roofs": 0,
            "floors": 0,
            "windows": 0,
            "doors": 0,
            "infiltration_sensible": 0,
            "infiltration_latent": 0,
            "ventilation_sensible": 0,
            "ventilation_latent": 0,
            "internal_gains_offset": 0,
            "detailed_loads": {
                "walls": [],
                "roofs": [],
                "floors": [],
                "windows": [],
                "doors": []
            }
        }
        
        # Get design conditions
        outdoor_temp = outdoor_conditions.get("design_heating_temp", -10)
        indoor_temp = indoor_conditions.get("design_heating_temp", 21)
        outdoor_rh = outdoor_conditions.get("design_heating_rh", 90)
        indoor_rh = indoor_conditions.get("design_heating_rh", 30)
        
        # Calculate wall heating loads
        for wall in building_components.get("walls", []):
            wall_load = self.calculate_wall_heating_load(
                wall=wall,
                outdoor_temp=outdoor_temp,
                indoor_temp=indoor_temp
            )
            
            results["walls"] += wall_load
            
            # Add detailed load
            results["detailed_loads"]["walls"].append({
                "name": wall.name,
                "orientation": wall.orientation.value,
                "area": wall.area,
                "u_value": wall.u_value,
                "delta_t": indoor_temp - outdoor_temp,
                "load": wall_load / 1000  # Convert to kW
            })
        
        # Calculate roof heating loads
        for roof in building_components.get("roofs", []):
            roof_load = self.calculate_roof_heating_load(
                roof=roof,
                outdoor_temp=outdoor_temp,
                indoor_temp=indoor_temp
            )
            
            results["roofs"] += roof_load
            
            # Add detailed load
            results["detailed_loads"]["roofs"].append({
                "name": roof.name,
                "area": roof.area,
                "u_value": roof.u_value,
                "delta_t": indoor_temp - outdoor_temp,
                "load": roof_load / 1000  # Convert to kW
            })
        
        # Calculate floor heating loads
        for floor in building_components.get("floors", []):
            ground_temp = outdoor_conditions.get("ground_temp", 5)
            floor_load = self.calculate_floor_heating_load(
                floor=floor,
                ground_temp=ground_temp,
                indoor_temp=indoor_temp
            )
            
            results["floors"] += floor_load
            
            # Add detailed load
            results["detailed_loads"]["floors"].append({
                "name": floor.name,
                "area": floor.area,
                "u_value": floor.u_value,
                "delta_t": indoor_temp - ground_temp,
                "load": floor_load / 1000  # Convert to kW
            })
        
        # Calculate window heating loads
        for window in building_components.get("windows", []):
            window_load = self.calculate_window_heating_load(
                window=window,
                outdoor_temp=outdoor_temp,
                indoor_temp=indoor_temp
            )
            
            results["windows"] += window_load
            
            # Add detailed load
            results["detailed_loads"]["windows"].append({
                "name": window.name,
                "orientation": window.orientation.value,
                "area": window.area,
                "u_value": window.u_value,
                "delta_t": indoor_temp - outdoor_temp,
                "load": window_load / 1000  # Convert to kW
            })
        
        # Calculate door heating loads
        for door in building_components.get("doors", []):
            door_load = self.calculate_door_heating_load(
                door=door,
                outdoor_temp=outdoor_temp,
                indoor_temp=indoor_temp
            )
            
            results["doors"] += door_load
            
            # Add detailed load
            results["detailed_loads"]["doors"].append({
                "name": door.name,
                "orientation": door.orientation.value,
                "area": door.area,
                "u_value": door.u_value,
                "delta_t": indoor_temp - outdoor_temp,
                "load": door_load / 1000  # Convert to kW
            })
        
        # Calculate infiltration heating load
        if "infiltration" in indoor_conditions:
            infiltration_info = indoor_conditions["infiltration"]
            
            # Convert ACH to flow rate
            if "ach" in infiltration_info:
                volume = indoor_conditions.get("volume", indoor_conditions.get("floor_area", 100) * 3)  # Assume 3m ceiling height if not specified
                flow_rate = self.heat_transfer.air_exchange_rate_to_flow_rate(
                    ach=infiltration_info["ach"],
                    volume=volume
                )
            else:
                flow_rate = infiltration_info.get("flow_rate", 0.05)  # Default to 0.05 m³/s
            
            infiltration_load = self.calculate_infiltration_heating_load(
                flow_rate=flow_rate,
                outdoor_temp=outdoor_temp,
                indoor_temp=indoor_temp,
                outdoor_rh=outdoor_rh,
                indoor_rh=indoor_rh
            )
            
            results["infiltration_sensible"] += infiltration_load["sensible"]
            results["infiltration_latent"] += infiltration_load["latent"]
        
        # Calculate ventilation heating load
        if "ventilation" in indoor_conditions:
            ventilation_info = indoor_conditions["ventilation"]
            
            # Convert ACH to flow rate
            if "ach" in ventilation_info:
                volume = indoor_conditions.get("volume", indoor_conditions.get("floor_area", 100) * 3)  # Assume 3m ceiling height if not specified
                flow_rate = self.heat_transfer.air_exchange_rate_to_flow_rate(
                    ach=ventilation_info["ach"],
                    volume=volume
                )
            else:
                flow_rate = ventilation_info.get("flow_rate", 0.1)  # Default to 0.1 m³/s
            
            ventilation_load = self.calculate_ventilation_heating_load(
                flow_rate=flow_rate,
                outdoor_temp=outdoor_temp,
                indoor_temp=indoor_temp,
                outdoor_rh=outdoor_rh,
                indoor_rh=indoor_rh
            )
            
            results["ventilation_sensible"] += ventilation_load["sensible"]
            results["ventilation_latent"] += ventilation_load["latent"]
        
        # Calculate internal gains offset
        people_load = internal_loads.get("people", {}).get("total", 0)
        lights_load = internal_loads.get("lights", {}).get("total", 0)
        equipment_load = internal_loads.get("equipment", {}).get("total", 0)
        
        internal_gains_offset = self.calculate_internal_gains_offset(
            people_load=people_load,
            lights_load=lights_load,
            equipment_load=equipment_load,
            usage_factor=internal_loads.get("usage_factor", 0.7)
        )
        
        results["internal_gains_offset"] = internal_gains_offset
        
        # Calculate subtotals
        envelope_load = (
            results["walls"] + 
            results["roofs"] + 
            results["floors"] + 
            results["windows"] + 
            results["doors"]
        )
        
        infiltration_load = results["infiltration_sensible"] + results["infiltration_latent"]
        ventilation_load = results["ventilation_sensible"] + results["ventilation_latent"]
        
        # Calculate total load
        total_load = (
            envelope_load + 
            infiltration_load + 
            ventilation_load - 
            internal_gains_offset
        )
        
        # Apply safety factor
        design_load = total_load * safety_factor
        
        # Add subtotals and totals to results
        results.update({
            "envelope": envelope_load,
            "infiltration": infiltration_load,
            "ventilation": ventilation_load,
            "total": total_load,
            "design_load": design_load,
            "safety_factor_percent": (safety_factor - 1) * 100
        })
        
        # Add per area metrics
        floor_area = indoor_conditions.get("floor_area", 100)  # Default to 100 m²
        results["load_per_area"] = design_load / floor_area if floor_area > 0 else 0
        
        return results
    
    def calculate_heating_load_profile(self, building_components: Dict[str, List[Any]], 
                                     outdoor_conditions: Dict[str, Any],
                                     indoor_conditions: Dict[str, Any],
                                     internal_loads: Dict[str, Any],
                                     temperature_profile: List[float]) -> List[Dict[str, float]]:
        """
        Calculate heating load profile for a range of outdoor temperatures.
        
        Args:
            building_components: Dictionary with lists of building components
            outdoor_conditions: Dictionary with outdoor conditions
            indoor_conditions: Dictionary with indoor conditions
            internal_loads: Dictionary with internal loads
            temperature_profile: List of outdoor temperatures in °C
            
        Returns:
            List of dictionaries with heating load results for each temperature
        """
        # Initialize results list
        results = []
        
        # Calculate heating load for each temperature
        for temp in temperature_profile:
            # Update outdoor temperature
            outdoor_conditions_temp = outdoor_conditions.copy()
            outdoor_conditions_temp["design_heating_temp"] = temp
            
            # Calculate heating load
            load = self.calculate_design_heating_load(
                building_components=building_components,
                outdoor_conditions=outdoor_conditions_temp,
                indoor_conditions=indoor_conditions,
                internal_loads=internal_loads
            )
            
            # Add temperature to results
            load["outdoor_temp"] = temp
            
            # Add to results list
            results.append(load)
        
        return results
    
    def calculate_balance_point(self, building_components: Dict[str, List[Any]], 
                              outdoor_conditions: Dict[str, Any],
                              indoor_conditions: Dict[str, Any],
                              internal_loads: Dict[str, Any],
                              temperature_range: List[float] = None) -> float:
        """
        Calculate balance point temperature for a building.
        
        Args:
            building_components: Dictionary with lists of building components
            outdoor_conditions: Dictionary with outdoor conditions
            indoor_conditions: Dictionary with indoor conditions
            internal_loads: Dictionary with internal loads
            temperature_range: List of outdoor temperatures to check (default: range from -20 to indoor temp)
            
        Returns:
            Balance point temperature in °C
        """
        # Get indoor temperature
        indoor_temp = indoor_conditions.get("design_heating_temp", 21)
        
        # Create temperature range if not provided
        if temperature_range is None:
            temperature_range = list(range(-20, int(indoor_temp) + 1))
        
        # Calculate heating load profile
        load_profile = self.calculate_heating_load_profile(
            building_components=building_components,
            outdoor_conditions=outdoor_conditions,
            indoor_conditions=indoor_conditions,
            internal_loads=internal_loads,
            temperature_profile=temperature_range
        )
        
        # Find balance point (where total load is approximately zero)
        for load in load_profile:
            if load["total"] <= 0:
                return load["outdoor_temp"]
        
        # If no balance point found, return indoor temperature
        return indoor_temp
    
    def calculate_bin_method_energy(self, building_components: Dict[str, List[Any]], 
                                  outdoor_conditions: Dict[str, Any],
                                  indoor_conditions: Dict[str, Any],
                                  internal_loads: Dict[str, Any],
                                  bin_data: Dict[str, List[int]]) -> Dict[str, float]:
        """
        Calculate annual heating energy using the bin method.
        
        Args:
            building_components: Dictionary with lists of building components
            outdoor_conditions: Dictionary with outdoor conditions
            indoor_conditions: Dictionary with indoor conditions
            internal_loads: Dictionary with internal loads
            bin_data: Dictionary with temperature bins and hours
            
        Returns:
            Dictionary with annual heating energy results
        """
        # Initialize results
        results = {
            "annual_energy": 0,
            "peak_load": 0,
            "bin_results": []
        }
        
        # Get indoor temperature
        indoor_temp = indoor_conditions.get("design_heating_temp", 21)
        
        # Calculate balance point
        balance_point = self.calculate_balance_point(
            building_components=building_components,
            outdoor_conditions=outdoor_conditions,
            indoor_conditions=indoor_conditions,
            internal_loads=internal_loads
        )
        
        # Calculate heating load for each bin
        for temp, hours in zip(bin_data["temperatures"], bin_data["hours"]):
            # Skip bins above balance point
            if temp >= balance_point:
                continue
            
            # Update outdoor temperature
            outdoor_conditions_temp = outdoor_conditions.copy()
            outdoor_conditions_temp["design_heating_temp"] = temp
            
            # Calculate heating load
            load = self.calculate_design_heating_load(
                building_components=building_components,
                outdoor_conditions=outdoor_conditions_temp,
                indoor_conditions=indoor_conditions,
                internal_loads=internal_loads
            )
            
            # Calculate energy for this bin
            energy = load["total"] * hours / 1000  # kWh
            
            # Add to results
            results["annual_energy"] += energy
            results["peak_load"] = max(results["peak_load"], load["total"])
            
            # Add bin result
            results["bin_results"].append({
                "temperature": temp,
                "hours": hours,
                "load": load["total"],
                "energy": energy
            })
        
        return results