"""
Psychrometric module for HVAC Load Calculator.
This module implements psychrometric calculations for air properties.
Reference: ASHRAE Handbook—Fundamentals (2017), Chapter 1.
"""

from typing import Dict, List, Any, Optional, Tuple
import math
import numpy as np

# Constants
ATMOSPHERIC_PRESSURE = 101325  # Standard atmospheric pressure in Pa
WATER_MOLECULAR_WEIGHT = 18.01534  # kg/kmol
DRY_AIR_MOLECULAR_WEIGHT = 28.9645  # kg/kmol
UNIVERSAL_GAS_CONSTANT = 8314.462618  # J/(kmol·K)
GAS_CONSTANT_DRY_AIR = UNIVERSAL_GAS_CONSTANT / DRY_AIR_MOLECULAR_WEIGHT  # J/(kg·K)
GAS_CONSTANT_WATER_VAPOR = UNIVERSAL_GAS_CONSTANT / WATER_MOLECULAR_WEIGHT  # J/(kg·K)


class Psychrometrics:
    """Class for psychrometric calculations."""
    
    @staticmethod
    def validate_inputs(t_db: float, rh: Optional[float] = None, p_atm: Optional[float] = None) -> None:
        """
        Validate input parameters for psychrometric calculations.
        
        Args:
            t_db: Dry-bulb temperature in °C
            rh: Relative humidity in % (0-100), optional
            p_atm: Atmospheric pressure in Pa, optional
            
        Raises:
            ValueError: If inputs are invalid
        """
        if not -50 <= t_db <= 60:
            raise ValueError(f"Temperature {t_db}°C must be between -50°C and 60°C")
        if rh is not None and not 0 <= rh <= 100:
            raise ValueError(f"Relative humidity {rh}% must be between 0 and 100%")
        if p_atm is not None and p_atm <= 0:
            raise ValueError(f"Atmospheric pressure {p_atm} Pa must be positive")

    @staticmethod
    def saturation_pressure(t_db: float) -> float:
        """
        Calculate saturation pressure of water vapor.
        Reference: ASHRAE Handbook—Fundamentals (2017), Chapter 1, Equations 5 and 6.
        
        Args:
            t_db: Dry-bulb temperature in °C
            
        Returns:
            Saturation pressure in Pa
        """
        Psychrometrics.validate_inputs(t_db)
        
        # Convert temperature to Kelvin
        t_k = t_db + 273.15
        
        # ASHRAE Fundamentals 2017 Chapter 1, Equation 5 & 6
        if t_db >= 0:
            # Equation 5 for temperatures above freezing
            c1 = -5.8002206e3
            c2 = 1.3914993
            c3 = -4.8640239e-2
            c4 = 4.1764768e-5
            c5 = -1.4452093e-8
            c6 = 6.5459673
        else:
            # Equation 6 for temperatures below freezing
            c1 = -5.6745359e3
            c2 = 6.3925247
            c3 = -9.6778430e-3
            c4 = 6.2215701e-7
            c5 = 2.0747825e-9
            c6 = -9.4840240e-13
            c7 = 4.1635019
        
        # Calculate natural log of saturation pressure in Pa
        if t_db >= 0:
            ln_p_ws = c1 / t_k + c2 + c3 * t_k + c4 * t_k**2 + c5 * t_k**3 + c6 * math.log(t_k)
        else:
            ln_p_ws = c1 / t_k + c2 + c3 * t_k + c4 * t_k**2 + c5 * t_k**3 + c6 * t_k**4 + c7 * math.log(t_k)
        
        # Convert from natural log to actual pressure in Pa
        p_ws = math.exp(ln_p_ws)
        
        return p_ws
    
    @staticmethod
    def humidity_ratio(t_db: float, rh: float, p_atm: float = ATMOSPHERIC_PRESSURE) -> float:
        """
        Calculate humidity ratio (mass of water vapor per unit mass of dry air).
        Reference: ASHRAE Handbook—Fundamentals (2017), Chapter 1, Equation 20.
        
        Args:
            t_db: Dry-bulb temperature in °C
            rh: Relative humidity (0-100)
            p_atm: Atmospheric pressure in Pa (default: standard atmospheric pressure)
            
        Returns:
            Humidity ratio in kg water vapor / kg dry air
        """
        Psychrometrics.validate_inputs(t_db, rh, p_atm)
        
        # Convert relative humidity to decimal
        rh_decimal = rh / 100.0
        
        # Calculate saturation pressure
        p_ws = Psychrometrics.saturation_pressure(t_db)
        
        # Calculate partial pressure of water vapor
        p_w = rh_decimal * p_ws
        
        if p_w >= p_atm:
            raise ValueError("Partial pressure of water vapor exceeds atmospheric pressure")
        
        # Calculate humidity ratio
        w = 0.621945 * p_w / (p_atm - p_w)
        
        return w
    
    @staticmethod
    def relative_humidity(t_db: float, w: float, p_atm: float = ATMOSPHERIC_PRESSURE) -> float:
        """
        Calculate relative humidity from humidity ratio.
        Reference: ASHRAE Handbook—Fundamentals (2017), Chapter 1, Equation 20 (rearranged).
        
        Args:
            t_db: Dry-bulb temperature in °C
            w: Humidity ratio in kg water vapor / kg dry air
            p_atm: Atmospheric pressure in Pa (default: standard atmospheric pressure)
            
        Returns:
            Relative humidity (0-100)
        """
        Psychrometrics.validate_inputs(t_db, p_atm=p_atm)
        if w < 0:
            raise ValueError("Humidity ratio cannot be negative")
        
        # Calculate saturation pressure
        p_ws = Psychrometrics.saturation_pressure(t_db)
        
        # Calculate partial pressure of water vapor
        p_w = p_atm * w / (0.621945 + w)
        
        # Calculate relative humidity
        rh = 100.0 * p_w / p_ws
        
        return rh
    
    @staticmethod
    def wet_bulb_temperature(t_db: float, rh: float, p_atm: float = ATMOSPHERIC_PRESSURE) -> float:
        """
        Calculate wet-bulb temperature using iterative method.
        Reference: ASHRAE Handbook—Fundamentals (2017), Chapter 1, Equation 35.
        
        Args:
            t_db: Dry-bulb temperature in °C
            rh: Relative humidity (0-100)
            p_atm: Atmospheric pressure in Pa (default: standard atmospheric pressure)
            
        Returns:
            Wet-bulb temperature in °C
        """
        Psychrometrics.validate_inputs(t_db, rh, p_atm)
        
        # Calculate humidity ratio at given conditions
        w = Psychrometrics.humidity_ratio(t_db, rh, p_atm)
        
        # Initial guess for wet-bulb temperature
        t_wb = t_db
        
        # Iterative solution
        max_iterations = 100
        tolerance = 0.001  # °C
        
        for i in range(max_iterations):
            # Validate wet-bulb temperature
            Psychrometrics.validate_inputs(t_wb)
            
            # Calculate saturation pressure at wet-bulb temperature
            p_ws_wb = Psychrometrics.saturation_pressure(t_wb)
            
            # Calculate saturation humidity ratio at wet-bulb temperature
            w_s_wb = 0.621945 * p_ws_wb / (p_atm - p_ws_wb)
            
            # Calculate humidity ratio from wet-bulb temperature
            h_fg = 2501000 + 1840 * t_wb  # Latent heat of vaporization at t_wb in J/kg
            c_pa = 1006  # Specific heat of dry air in J/(kg·K)
            c_pw = 1860  # Specific heat of water vapor in J/(kg·K)
            
            w_calc = ((h_fg - c_pw * (t_db - t_wb)) * w_s_wb - c_pa * (t_db - t_wb)) / (h_fg + c_pw * t_db - c_pw * t_wb)
            
            # Check convergence
            if abs(w - w_calc) < tolerance:
                break
            
            # Adjust wet-bulb temperature
            if w_calc > w:
                t_wb -= 0.1
            else:
                t_wb += 0.1
        
        return t_wb
    
    @staticmethod
    def dew_point_temperature(t_db: float, rh: float) -> float:
        """
        Calculate dew point temperature.
        Reference: ASHRAE Handbook—Fundamentals (2017), Chapter 1, Equations 39 and 40.
        
        Args:
            t_db: Dry-bulb temperature in °C
            rh: Relative humidity (0-100)
            
        Returns:
            Dew point temperature in °C
        """
        Psychrometrics.validate_inputs(t_db, rh)
        
        # Convert relative humidity to decimal
        rh_decimal = rh / 100.0
        
        # Calculate saturation pressure
        p_ws = Psychrometrics.saturation_pressure(t_db)
        
        # Calculate partial pressure of water vapor
        p_w = rh_decimal * p_ws
        
        # Calculate dew point temperature
        alpha = math.log(p_w / 1000.0)  # Convert to kPa for the formula
        
        if t_db >= 0:
            # For temperatures above freezing
            c14 = 6.54
            c15 = 14.526
            c16 = 0.7389
            c17 = 0.09486
            c18 = 0.4569
            
            t_dp = c14 + c15 * alpha + c16 * alpha**2 + c17 * alpha**3 + c18 * p_w**(0.1984)
        else:
            # For temperatures below freezing
            c14 = 6.09
            c15 = 12.608
            c16 = 0.4959
            
            t_dp = c14 + c15 * alpha + c16 * alpha**2
        
        return t_dp
    
    @staticmethod
    def enthalpy(t_db: float, w: float) -> float:
        """
        Calculate specific enthalpy of moist air.
        Reference: ASHRAE Handbook—Fundamentals (2017), Chapter 1, Equation 30.
        
        Args:
            t_db: Dry-bulb temperature in °C
            w: Humidity ratio in kg water vapor / kg dry air
            
        Returns:
            Specific enthalpy in J/kg dry air
        """
        Psychrometrics.validate_inputs(t_db)
        if w < 0:
            raise ValueError("Humidity ratio cannot be negative")
        
        c_pa = 1006  # Specific heat of dry air in J/(kg·K)
        h_fg = 2501000  # Latent heat of vaporization at 0°C in J/kg
        c_pw = 1860  # Specific heat of water vapor in J/(kg·K)
        
        h = c_pa * t_db + w * (h_fg + c_pw * t_db)
        
        return h
    
    @staticmethod
    def specific_volume(t_db: float, w: float, p_atm: float = ATMOSPHERIC_PRESSURE) -> float:
        """
        Calculate specific volume of moist air.
        Reference: ASHRAE Handbook—Fundamentals (2017), Chapter 1, Equation 28.
        
        Args:
            t_db: Dry-bulb temperature in °C
            w: Humidity ratio in kg water vapor / kg dry air
            p_atm: Atmospheric pressure in Pa (default: standard atmospheric pressure)
            
        Returns:
            Specific volume in m³/kg dry air
        """
        Psychrometrics.validate_inputs(t_db, p_atm=p_atm)
        if w < 0:
            raise ValueError("Humidity ratio cannot be negative")
        
        # Convert temperature to Kelvin
        t_k = t_db + 273.15
        
        r_da = GAS_CONSTANT_DRY_AIR  # Gas constant for dry air in J/(kg·K)
        
        v = r_da * t_k * (1 + 1.607858 * w) / p_atm
        
        return v
    
    @staticmethod
    def density(t_db: float, w: float, p_atm: float = ATMOSPHERIC_PRESSURE) -> float:
        """
        Calculate density of moist air.
        Reference: ASHRAE Handbook—Fundamentals (2017), Chapter 1, derived from Equation 28.
        
        Args:
            t_db: Dry-bulb temperature in °C
            w: Humidity ratio in kg water vapor / kg dry air
            p_atm: Atmospheric pressure in Pa (default: standard atmospheric pressure)
            
        Returns:
            Density in kg/m³
        """
        Psychrometrics.validate_inputs(t_db, p_atm=p_atm)
        if w < 0:
            raise ValueError("Humidity ratio cannot be negative")
        
        # Calculate specific volume
        v = Psychrometrics.specific_volume(t_db, w, p_atm)
        
        # Density is the reciprocal of specific volume
        rho = (1 + w) / v
        
        return rho
    
    @staticmethod
    def moist_air_properties(t_db: float, rh: float, p_atm: float = ATMOSPHERIC_PRESSURE) -> Dict[str, float]:
        """
        Calculate all psychrometric properties of moist air.
        Reference: ASHRAE Handbook—Fundamentals (2017), Chapter 1.
        
        Args:
            t_db: Dry-bulb temperature in °C
            rh: Relative humidity (0-100)
            p_atm: Atmospheric pressure in Pa (default: standard atmospheric pressure)
            
        Returns:
            Dictionary with all psychrometric properties
        """
        Psychrometrics.validate_inputs(t_db, rh, p_atm)
        
        # Calculate humidity ratio
        w = Psychrometrics.humidity_ratio(t_db, rh, p_atm)
        
        # Calculate wet-bulb temperature
        t_wb = Psychrometrics.wet_bulb_temperature(t_db, rh, p_atm)
        
        # Calculate dew point temperature
        t_dp = Psychrometrics.dew_point_temperature(t_db, rh)
        
        # Calculate enthalpy
        h = Psychrometrics.enthalpy(t_db, w)
        
        # Calculate specific volume
        v = Psychrometrics.specific_volume(t_db, w, p_atm)
        
        # Calculate density
        rho = Psychrometrics.density(t_db, w, p_atm)
        
        # Calculate saturation pressure
        p_ws = Psychrometrics.saturation_pressure(t_db)
        
        # Calculate partial pressure of water vapor
        p_w = rh / 100.0 * p_ws
        
        # Return all properties
        return {
            "dry_bulb_temperature": t_db,
            "wet_bulb_temperature": t_wb,
            "dew_point_temperature": t_dp,
            "relative_humidity": rh,
            "humidity_ratio": w,
            "enthalpy": h,
            "specific_volume": v,
            "density": rho,
            "saturation_pressure": p_ws,
            "partial_pressure": p_w,
            "atmospheric_pressure": p_atm
        }
    
    @staticmethod
    def find_humidity_ratio_for_enthalpy(t_db: float, h: float) -> float:
        """
        Find humidity ratio for a given dry-bulb temperature and enthalpy.
        Reference: ASHRAE Handbook—Fundamentals (2017), Chapter 1, Equation 30 (rearranged).
        
        Args:
            t_db: Dry-bulb temperature in °C
            h: Specific enthalpy in J/kg dry air
            
        Returns:
            Humidity ratio in kg water vapor / kg dry air
        """
        Psychrometrics.validate_inputs(t_db)
        if h < 0:
            raise ValueError("Enthalpy cannot be negative")
        
        c_pa = 1006  # Specific heat of dry air in J/(kg·K)
        h_fg = 2501000  # Latent heat of vaporization at 0°C in J/kg
        c_pw = 1860  # Specific heat of water vapor in J/(kg·K)
        
        w = (h - c_pa * t_db) / (h_fg + c_pw * t_db)
        
        return max(0, w)
    
    @staticmethod
    def find_temperature_for_enthalpy(w: float, h: float) -> float:
        """
        Find dry-bulb temperature for a given humidity ratio and enthalpy.
        Reference: ASHRAE Handbook—Fundamentals (2017), Chapter 1, Equation 30 (rearranged).
        
        Args:
            w: Humidity ratio in kg water vapor / kg dry air
            h: Specific enthalpy in J/kg dry air
            
        Returns:
            Dry-bulb temperature in °C
        """
        if w < 0:
            raise ValueError("Humidity ratio cannot be negative")
        if h < 0:
            raise ValueError("Enthalpy cannot be negative")
        
        c_pa = 1006  # Specific heat of dry air in J/(kg·K)
        h_fg = 2501000  # Latent heat of vaporization at 0°C in J/kg
        c_pw = 1860  # Specific heat of water vapor in J/(kg·K)
        
        t_db = (h - w * h_fg) / (c_pa + w * c_pw)
        
        Psychrometrics.validate_inputs(t_db)
        return t_db
    
    @staticmethod
    def sensible_heat_ratio(q_sensible: float, q_total: float) -> float:
        """
        Calculate sensible heat ratio.
        Reference: ASHRAE Handbook—Fundamentals (2017), Chapter 1, Section 1.5.
        
        Args:
            q_sensible: Sensible heat load in W
            q_total: Total heat load in W
            
        Returns:
            Sensible heat ratio (0-1)
        """
        if q_total == 0:
            return 1.0
        if q_sensible < 0 or q_total < 0:
            raise ValueError("Heat loads cannot be negative")
        
        return q_sensible / q_total
    
    @staticmethod
    def air_flow_rate_for_load(q_sensible: float, t_supply: float, t_return: float, 
                              rh_return: float = 50.0, p_atm: float = ATMOSPHERIC_PRESSURE) -> Dict[str, float]:
        """
        Calculate required air flow rate for a given sensible load.
        Reference: ASHRAE Handbook—Fundamentals (2017), Chapter 1, Section 1.6.
        
        Args:
            q_sensible: Sensible heat load in W
            t_supply: Supply air temperature in °C
            t_return: Return air temperature in °C
            rh_return: Return air relative humidity in % (default: 50%)
            p_atm: Atmospheric pressure in Pa (default: standard atmospheric pressure)
            
        Returns:
            Dictionary with air flow rate in different units
        """
        Psychrometrics.validate_inputs(t_return, rh_return, p_atm)
        Psychrometrics.validate_inputs(t_supply)
        
        # Calculate return air properties
        w_return = Psychrometrics.humidity_ratio(t_return, rh_return, p_atm)
        rho_return = Psychrometrics.density(t_return, w_return, p_atm)
        
        # Calculate specific heat of moist air
        c_pa = 1006  # Specific heat of dry air in J/(kg·K)
        c_pw = 1860  # Specific heat of water vapor in J/(kg·K)
        c_p_moist = c_pa + w_return * c_pw
        
        # Calculate mass flow rate
        delta_t = t_return - t_supply
        if delta_t == 0:
            raise ValueError("Supply and return temperatures cannot be equal")
        
        m_dot = q_sensible / (c_p_moist * delta_t)
        
        # Calculate volumetric flow rate
        v_dot = m_dot / rho_return
        
        # Convert to different units
        v_dot_m3_s = v_dot
        v_dot_m3_h = v_dot * 3600
        v_dot_cfm = v_dot * 2118.88
        v_dot_l_s = v_dot * 1000
        
        return {
            "mass_flow_rate_kg_s": m_dot,
            "volumetric_flow_rate_m3_s": v_dot_m3_s,
            "volumetric_flow_rate_m3_h": v_dot_m3_h,
            "volumetric_flow_rate_cfm": v_dot_cfm,
            "volumetric_flow_rate_l_s": v_dot_l_s
        }
    
    @staticmethod
    def mixing_air_properties(m1: float, t_db1: float, rh1: float, 
                             m2: float, t_db2: float, rh2: float, 
                             p_atm: float = ATMOSPHERIC_PRESSURE) -> Dict[str, float]:
        """
        Calculate properties of mixed airstreams.
        Reference: ASHRAE Handbook—Fundamentals (2017), Chapter 1, Section 1.7.
        
        Args:
            m1: Mass flow rate of airstream 1 in kg/s
            t_db1: Dry-bulb temperature of airstream 1 in °C
            rh1: Relative humidity of airstream 1 in %
            m2: Mass flow rate of airstream 2 in kg/s
            t_db2: Dry-bulb temperature of airstream 2 in °C
            rh2: Relative humidity of airstream 2 in %
            p_atm: Atmospheric pressure in Pa (default: standard atmospheric pressure)
            
        Returns:
            Dictionary with mixed air properties
        """
        Psychrometrics.validate_inputs(t_db1, rh1, p_atm)
        Psychrometrics.validate_inputs(t_db2, rh2, p_atm)
        if m1 < 0 or m2 < 0:
            raise ValueError("Mass flow rates cannot be negative")
        
        # Calculate humidity ratios
        w1 = Psychrometrics.humidity_ratio(t_db1, rh1, p_atm)
        w2 = Psychrometrics.humidity_ratio(t_db2, rh2, p_atm)
        
        # Calculate enthalpies
        h1 = Psychrometrics.enthalpy(t_db1, w1)
        h2 = Psychrometrics.enthalpy(t_db2, w2)
        
        # Calculate mixed air properties
        m_total = m1 + m2
        
        if m_total == 0:
            raise ValueError("Total mass flow rate cannot be zero")
        
        w_mix = (m1 * w1 + m2 * w2) / m_total
        h_mix = (m1 * h1 + m2 * h2) / m_total
        
        # Find dry-bulb temperature for the mixed air
        t_db_mix = Psychrometrics.find_temperature_for_enthalpy(w_mix, h_mix)
        
        # Calculate relative humidity for the mixed air
        rh_mix = Psychrometrics.relative_humidity(t_db_mix, w_mix, p_atm)
        
        # Return mixed air properties
        return Psychrometrics.moist_air_properties(t_db_mix, rh_mix, p_atm)


# Create a singleton instance
psychrometrics = Psychrometrics()

# Example usage
if __name__ == "__main__":
    # Calculate properties of air at 25°C and 50% RH
    properties = psychrometrics.moist_air_properties(25, 50)
    
    print("Air Properties at 25°C and 50% RH:")
    print(f"Dry-bulb temperature: {properties['dry_bulb_temperature']:.2f} °C")
    print(f"Wet-bulb temperature: {properties['wet_bulb_temperature']:.2f} °C")
    print(f"Dew point temperature: {properties['dew_point_temperature']:.2f} °C")
    print(f"Relative humidity: {properties['relative_humidity']:.2f} %")
    print(f"Humidity ratio: {properties['humidity_ratio']:.6f} kg/kg")
    print(f"Enthalpy: {properties['enthalpy']/1000:.2f} kJ/kg")
    print(f"Specific volume: {properties['specific_volume']:.4f} m³/kg")
    print(f"Density: {properties['density']:.4f} kg/m³")
    print(f"Saturation pressure: {properties['saturation_pressure']/1000:.2f} kPa")
    print(f"Partial pressure: {properties['partial_pressure']/1000:.2f} kPa")