utils/psychrometric_visualization.py

"""
Psychrometric visualization module for HVAC Load Calculator.
This module provides visualization tools for psychrometric processes.
"""

import streamlit as st
import pandas as pd
import numpy as np
import plotly.graph_objects as go
import plotly.express as px
from typing import Dict, List, Any, Optional, Tuple
import math

# Import psychrometrics module
from utils.psychrometrics import Psychrometrics


class PsychrometricVisualization:
    """Class for psychrometric visualization."""
    
    def __init__(self):
        """Initialize psychrometric visualization."""
        self.psychrometrics = Psychrometrics()
        
        # Define temperature and humidity ratio ranges for chart
        self.temp_min = -10
        self.temp_max = 50
        self.w_min = 0
        self.w_max = 0.030
        
        # Define standard atmospheric pressure
        self.pressure = 101325  # Pa
    
    def create_psychrometric_chart(self, points: Optional[List[Dict[str, Any]]] = None,
                                  processes: Optional[List[Dict[str, Any]]] = None,
                                  comfort_zone: Optional[Dict[str, Any]] = None) -> go.Figure:
        """
        Create an interactive psychrometric chart.
        
        Args:
            points: List of points to plot on the chart
            processes: List of processes to plot on the chart
            comfort_zone: Dictionary with comfort zone parameters
            
        Returns:
            Plotly figure with psychrometric chart
        """
        # Create figure
        fig = go.Figure()
        
        # Generate temperature and humidity ratio grids
        temp_range = np.linspace(self.temp_min, self.temp_max, 100)
        w_range = np.linspace(self.w_min, self.w_max, 100)
        
        # Generate saturation curve
        sat_temps = np.linspace(self.temp_min, self.temp_max, 100)
        sat_w = [self.psychrometrics.humidity_ratio(t, 100, self.pressure) for t in sat_temps]
        
        # Plot saturation curve
        fig.add_trace(go.Scatter(
            x=sat_temps,
            y=sat_w,
            mode="lines",
            line=dict(color="blue", width=2),
            name="Saturation Curve"
        ))
        
        # Generate constant RH curves
        rh_values = [10, 20, 30, 40, 50, 60, 70, 80, 90]
        
        for rh in rh_values:
            rh_temps = np.linspace(self.temp_min, self.temp_max, 50)
            rh_w = [self.psychrometrics.humidity_ratio(t, rh, self.pressure) for t in rh_temps]
            
            # Filter out values above saturation
            valid_points = [(t, w) for t, w in zip(rh_temps, rh_w) if w <= self.psychrometrics.humidity_ratio(t, 100, self.pressure)]
            
            if valid_points:
                valid_temps, valid_w = zip(*valid_points)
                
                fig.add_trace(go.Scatter(
                    x=valid_temps,
                    y=valid_w,
                    mode="lines",
                    line=dict(color="rgba(0, 0, 255, 0.3)", width=1, dash="dot"),
                    name=f"{rh}% RH",
                    hoverinfo="name"
                ))
        
        # Generate constant wet-bulb temperature lines
        wb_values = np.arange(0, 35, 5)
        
        for wb in wb_values:
            wb_temps = np.linspace(wb, self.temp_max, 50)
            wb_points = []
            
            for t in wb_temps:
                # Binary search to find humidity ratio for this wet-bulb temperature
                w_low = 0
                w_high = self.psychrometrics.humidity_ratio(t, 100, self.pressure)
                
                for _ in range(10):  # 10 iterations should be enough for good precision
                    w_mid = (w_low + w_high) / 2
                    rh = self.psychrometrics.relative_humidity(t, w_mid, self.pressure)
                    t_wb_calc = self.psychrometrics.wet_bulb_temperature(t, rh, self.pressure)
                    
                    if abs(t_wb_calc - wb) < 0.1:
                        wb_points.append((t, w_mid))
                        break
                    elif t_wb_calc < wb:
                        w_low = w_mid
                    else:
                        w_high = w_mid
            
            if wb_points:
                wb_temps, wb_w = zip(*wb_points)
                
                fig.add_trace(go.Scatter(
                    x=wb_temps,
                    y=wb_w,
                    mode="lines",
                    line=dict(color="rgba(0, 128, 0, 0.3)", width=1, dash="dash"),
                    name=f"{wb}°C WB",
                    hoverinfo="name"
                ))
        
        # Generate constant enthalpy lines
        h_values = np.arange(0, 100, 10) * 1000  # kJ/kg to J/kg
        
        for h in h_values:
            h_temps = np.linspace(self.temp_min, self.temp_max, 50)
            h_points = []
            
            for t in h_temps:
                # Calculate humidity ratio for this enthalpy
                w = self.psychrometrics.find_humidity_ratio_for_enthalpy(t, h)
                
                if 0 <= w <= self.psychrometrics.humidity_ratio(t, 100, self.pressure):
                    h_points.append((t, w))
            
            if h_points:
                h_temps, h_w = zip(*h_points)
                
                fig.add_trace(go.Scatter(
                    x=h_temps,
                    y=h_w,
                    mode="lines",
                    line=dict(color="rgba(255, 0, 0, 0.3)", width=1, dash="dashdot"),
                    name=f"{h/1000:.0f} kJ/kg",
                    hoverinfo="name"
                ))
        
        # Generate constant specific volume lines
        v_values = [0.8, 0.85, 0.9, 0.95, 1.0, 1.05]
        
        for v in v_values:
            v_temps = np.linspace(self.temp_min, self.temp_max, 50)
            v_points = []
            
            for t in h_temps:
                # Binary search to find humidity ratio for this specific volume
                w_low = 0
                w_high = self.psychrometrics.humidity_ratio(t, 100, self.pressure)
                
                for _ in range(10):  # 10 iterations should be enough for good precision
                    w_mid = (w_low + w_high) / 2
                    v_calc = self.psychrometrics.specific_volume(t, w_mid, self.pressure)
                    
                    if abs(v_calc - v) < 0.01:
                        v_points.append((t, w_mid))
                        break
                    elif v_calc < v:
                        w_low = w_mid
                    else:
                        w_high = w_mid
            
            if v_points:
                v_temps, v_w = zip(*v_points)
                
                fig.add_trace(go.Scatter(
                    x=v_temps,
                    y=v_w,
                    mode="lines",
                    line=dict(color="rgba(128, 0, 128, 0.3)", width=1, dash="longdash"),
                    name=f"{v:.2f} m³/kg",
                    hoverinfo="name"
                ))
        
        # Add comfort zone if specified
        if comfort_zone:
            temp_min = comfort_zone.get("temp_min", 20)
            temp_max = comfort_zone.get("temp_max", 26)
            rh_min = comfort_zone.get("rh_min", 30)
            rh_max = comfort_zone.get("rh_max", 60)
            
            # Calculate humidity ratios at corners
            w_bottom_left = self.psychrometrics.humidity_ratio(temp_min, rh_min, self.pressure)
            w_bottom_right = self.psychrometrics.humidity_ratio(temp_max, rh_min, self.pressure)
            w_top_right = self.psychrometrics.humidity_ratio(temp_max, rh_max, self.pressure)
            w_top_left = self.psychrometrics.humidity_ratio(temp_min, rh_max, self.pressure)
            
            # Add comfort zone as a filled polygon
            fig.add_trace(go.Scatter(
                x=[temp_min, temp_max, temp_max, temp_min, temp_min],
                y=[w_bottom_left, w_bottom_right, w_top_right, w_top_left, w_bottom_left],
                fill="toself",
                fillcolor="rgba(0, 255, 0, 0.2)",
                line=dict(color="green", width=2),
                name="Comfort Zone"
            ))
        
        # Add points if specified
        if points:
            for i, point in enumerate(points):
                temp = point.get("temp", 0)
                rh = point.get("rh", 0)
                w = point.get("w", self.psychrometrics.humidity_ratio(temp, rh, self.pressure))
                name = point.get("name", f"Point {i+1}")
                color = point.get("color", "blue")
                
                fig.add_trace(go.Scatter(
                    x=[temp],
                    y=[w],
                    mode="markers+text",
                    marker=dict(size=10, color=color),
                    text=[name],
                    textposition="top center",
                    name=name
                ))
        
        # Add processes if specified
        if processes:
            for i, process in enumerate(processes):
                start_point = process.get("start", {})
                end_point = process.get("end", {})
                
                start_temp = start_point.get("temp", 0)
                start_rh = start_point.get("rh", 0)
                start_w = start_point.get("w", self.psychrometrics.humidity_ratio(start_temp, start_rh, self.pressure))
                
                end_temp = end_point.get("temp", 0)
                end_rh = end_point.get("rh", 0)
                end_w = end_point.get("w", self.psychrometrics.humidity_ratio(end_temp, end_rh, self.pressure))
                
                name = process.get("name", f"Process {i+1}")
                color = process.get("color", "red")
                
                fig.add_trace(go.Scatter(
                    x=[start_temp, end_temp],
                    y=[start_w, end_w],
                    mode="lines+markers",
                    line=dict(color=color, width=2, dash="solid"),
                    marker=dict(size=8),
                    name=name
                ))
        
        # Update layout
        fig.update_layout(
            title="Psychrometric Chart",
            xaxis_title="Dry-Bulb Temperature (°C)",
            yaxis_title="Humidity Ratio (kg/kg)",
            legend_title="Legend",
            height=700,
            margin=dict(l=50, r=50, t=50, b=50),
            plot_bgcolor="white",
            paper_bgcolor="white",
            font=dict(size=12)
        )
        
        # Set axis ranges
        fig.update_xaxes(range=[self.temp_min, self.temp_max], gridcolor="lightgray")
        fig.update_yaxes(range=[self.w_min, self.w_max], gridcolor="lightgray")
        
        return fig

    def display_psychrometric_chart(self, calculation_results: Dict[str, Any], design_conditions: Dict[str, Any]) -> None:
        """
        Display psychrometric chart with calculation results.
        
        Args:
            calculation_results: Dictionary containing calculation results
            design_conditions: Dictionary containing design conditions
        """
        # Extract design conditions
        summer_outdoor_db = design_conditions.get("summer_outdoor_db", 35)
        summer_outdoor_wb = design_conditions.get("summer_outdoor_wb", 25)
        summer_indoor_db = design_conditions.get("summer_indoor_db", 24)
        summer_indoor_rh = design_conditions.get("summer_indoor_rh", 50)
        
        winter_outdoor_db = design_conditions.get("winter_outdoor_db", 0)
        winter_outdoor_rh = design_conditions.get("winter_outdoor_rh", 80)
        winter_indoor_db = design_conditions.get("winter_indoor_db", 22)
        winter_indoor_rh = design_conditions.get("winter_indoor_rh", 40)
        
        # Calculate humidity ratios
        summer_outdoor_w = self.psychrometrics.humidity_ratio_from_wb(summer_outdoor_db, summer_outdoor_wb, self.pressure)
        summer_indoor_w = self.psychrometrics.humidity_ratio(summer_indoor_db, summer_indoor_rh, self.pressure)
        winter_outdoor_w = self.psychrometrics.humidity_ratio(winter_outdoor_db, winter_outdoor_rh, self.pressure)
        winter_indoor_w = self.psychrometrics.humidity_ratio(winter_indoor_db, winter_indoor_rh, self.pressure)
        
        # Create points for psychrometric chart
        points = [
            {
                "temp": summer_outdoor_db,
                "w": summer_outdoor_w,
                "name": "Summer Outdoor",
                "color": "red"
            },
            {
                "temp": summer_indoor_db,
                "w": summer_indoor_w,
                "name": "Summer Indoor",
                "color": "blue"
            },
            {
                "temp": winter_outdoor_db,
                "w": winter_outdoor_w,
                "name": "Winter Outdoor",
                "color": "purple"
            },
            {
                "temp": winter_indoor_db,
                "w": winter_indoor_w,
                "name": "Winter Indoor",
                "color": "green"
            }
        ]
        
        # Create processes for psychrometric chart
        processes = [
            {
                "start": {"temp": summer_outdoor_db, "w": summer_outdoor_w},
                "end": {"temp": summer_indoor_db, "w": summer_indoor_w},
                "name": "Cooling Process",
                "color": "blue"
            },
            {
                "start": {"temp": winter_outdoor_db, "w": winter_outdoor_w},
                "end": {"temp": winter_indoor_db, "w": winter_indoor_w},
                "name": "Heating Process",
                "color": "red"
            }
        ]
        
        # Create comfort zone
        comfort_zone = {
            "temp_min": 20,
            "temp_max": 26,
            "rh_min": 30,
            "rh_max": 60
        }
        
        # Create psychrometric chart
        fig = self.create_psychrometric_chart(points, processes, comfort_zone)
        
        # Display chart in Streamlit
        st.plotly_chart(fig, use_container_width=True)
        
        # Display psychrometric properties
        st.subheader("Psychrometric Properties")
        
        # Create dataframe for properties
        properties = []
        
        # Summer outdoor properties
        summer_outdoor_rh = self.psychrometrics.relative_humidity_from_wb(summer_outdoor_db, summer_outdoor_wb, self.pressure)
        summer_outdoor_dp = self.psychrometrics.dew_point(summer_outdoor_db, summer_outdoor_rh, self.pressure)
        summer_outdoor_h = self.psychrometrics.enthalpy(summer_outdoor_db, summer_outdoor_w)
        summer_outdoor_v = self.psychrometrics.specific_volume(summer_outdoor_db, summer_outdoor_w, self.pressure)
        
        properties.append({
            "Point": "Summer Outdoor",
            "Dry-Bulb (°C)": f"{summer_outdoor_db:.1f}",
            "Wet-Bulb (°C)": f"{summer_outdoor_wb:.1f}",
            "Relative Humidity (%)": f"{summer_outdoor_rh:.1f}",
            "Humidity Ratio (g/kg)": f"{summer_outdoor_w*1000:.1f}",
            "Dew Point (°C)": f"{summer_outdoor_dp:.1f}",
            "Enthalpy (kJ/kg)": f"{summer_outdoor_h/1000:.1f}",
            "Specific Volume (m³/kg)": f"{summer_outdoor_v:.3f}"
        })
        
        # Summer indoor properties
        summer_indoor_wb = self.psychrometrics.wet_bulb_temperature(summer_indoor_db, summer_indoor_rh, self.pressure)
        summer_indoor_dp = self.psychrometrics.dew_point(summer_indoor_db, summer_indoor_rh, self.pressure)
        summer_indoor_h = self.psychrometrics.enthalpy(summer_indoor_db, summer_indoor_w)
        summer_indoor_v = self.psychrometrics.specific_volume(summer_indoor_db, summer_indoor_w, self.pressure)
        
        properties.append({
            "Point": "Summer Indoor",
            "Dry-Bulb (°C)": f"{summer_indoor_db:.1f}",
            "Wet-Bulb (°C)": f"{summer_indoor_wb:.1f}",
            "Relative Humidity (%)": f"{summer_indoor_rh:.1f}",
            "Humidity Ratio (g/kg)": f"{summer_indoor_w*1000:.1f}",
            "Dew Point (°C)": f"{summer_indoor_dp:.1f}",
            "Enthalpy (kJ/kg)": f"{summer_indoor_h/1000:.1f}",
            "Specific Volume (m³/kg)": f"{summer_indoor_v:.3f}"
        })
        
        # Winter outdoor properties
        winter_outdoor_wb = self.psychrometrics.wet_bulb_temperature(winter_outdoor_db, winter_outdoor_rh, self.pressure)
        winter_outdoor_dp = self.psychrometrics.dew_point(winter_outdoor_db, winter_outdoor_rh, self.pressure)
        winter_outdoor_h = self.psychrometrics.enthalpy(winter_outdoor_db, winter_outdoor_w)
        winter_outdoor_v = self.psychrometrics.specific_volume(winter_outdoor_db, winter_outdoor_w, self.pressure)
        
        properties.append({
            "Point": "Winter Outdoor",
            "Dry-Bulb (°C)": f"{winter_outdoor_db:.1f}",
            "Wet-Bulb (°C)": f"{winter_outdoor_wb:.1f}",
            "Relative Humidity (%)": f"{winter_outdoor_rh:.1f}",
            "Humidity Ratio (g/kg)": f"{winter_outdoor_w*1000:.1f}",
            "Dew Point (°C)": f"{winter_outdoor_dp:.1f}",
            "Enthalpy (kJ/kg)": f"{winter_outdoor_h/1000:.1f}",
            "Specific Volume (m³/kg)": f"{winter_outdoor_v:.3f}"
        })
        
        # Winter indoor properties
        winter_indoor_wb = self.psychrometrics.wet_bulb_temperature(winter_indoor_db, winter_indoor_rh, self.pressure)
        winter_indoor_dp = self.psychrometrics.dew_point(winter_indoor_db, winter_indoor_rh, self.pressure)
        winter_indoor_h = self.psychrometrics.enthalpy(winter_indoor_db, winter_indoor_w)
        winter_indoor_v = self.psychrometrics.specific_volume(winter_indoor_db, winter_indoor_w, self.pressure)
        
        properties.append({
            "Point": "Winter Indoor",
            "Dry-Bulb (°C)": f"{winter_indoor_db:.1f}",
            "Wet-Bulb (°C)": f"{winter_indoor_wb:.1f}",
            "Relative Humidity (%)": f"{winter_indoor_rh:.1f}",
            "Humidity Ratio (g/kg)": f"{winter_indoor_w*1000:.1f}",
            "Dew Point (°C)": f"{winter_indoor_dp:.1f}",
            "Enthalpy (kJ/kg)": f"{winter_indoor_h/1000:.1f}",
            "Specific Volume (m³/kg)": f"{winter_indoor_v:.3f}"
        })
        
        # Display properties table
        st.dataframe(pd.DataFrame(properties), use_container_width=True)