"""
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,
                    hovertemplate=(
                        f"<b>{name}</b><br>" +
                        "Temperature: %{x:.1f}°C<br>" +
                        "Humidity Ratio: %{y:.5f} kg/kg<br>" +
                        f"Relative Humidity: {rh:.1f}%<br>"
                    )
                ))
        
        # 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, color=color),
                    name=name
                ))
                
                # Add arrow to indicate direction
                fig.add_annotation(
                    x=end_temp,
                    y=end_w,
                    ax=start_temp,
                    ay=start_w,
                    xref="x",
                    yref="y",
                    axref="x",
                    ayref="y",
                    showarrow=True,
                    arrowhead=2,
                    arrowsize=1,
                    arrowwidth=2,
                    arrowcolor=color
                )
        
        # Update layout
        fig.update_layout(
            title="Psychrometric Chart",
            xaxis_title="Dry-Bulb Temperature (°C)",
            yaxis_title="Humidity Ratio (kg/kg)",
            xaxis=dict(
                range=[self.temp_min, self.temp_max],
                gridcolor="rgba(0, 0, 0, 0.1)",
                showgrid=True
            ),
            yaxis=dict(
                range=[self.w_min, self.w_max],
                gridcolor="rgba(0, 0, 0, 0.1)",
                showgrid=True
            ),
            height=700,
            margin=dict(l=50, r=50, b=50, t=50),
            legend=dict(
                orientation="h",
                yanchor="bottom",
                y=1.02,
                xanchor="right",
                x=1
            ),
            hovermode="closest"
        )
        
        return fig
    
    def create_process_visualization(self, process: Dict[str, Any]) -> go.Figure:
        """
        Create a visualization of a psychrometric process.
        
        Args:
            process: Dictionary with process parameters
            
        Returns:
            Plotly figure with process visualization
        """
        # Extract process parameters
        start_point = process.get("start", {})
        end_point = process.get("end", {})
        
        start_temp = start_point.get("temp", 0)
        start_rh = start_point.get("rh", 0)
        
        end_temp = end_point.get("temp", 0)
        end_rh = end_point.get("rh", 0)
        
        # Calculate psychrometric properties
        start_props = self.psychrometrics.moist_air_properties(start_temp, start_rh, self.pressure)
        end_props = self.psychrometrics.moist_air_properties(end_temp, end_rh, self.pressure)
        
        # Calculate process changes
        delta_t = end_temp - start_temp
        delta_w = end_props["humidity_ratio"] - start_props["humidity_ratio"]
        delta_h = end_props["enthalpy"] - start_props["enthalpy"]
        
        # Determine process type
        process_type = "Unknown"
        if abs(delta_w) < 0.0001:  # Sensible heating/cooling
            if delta_t > 0:
                process_type = "Sensible Heating"
            else:
                process_type = "Sensible Cooling"
        elif abs(delta_t) < 0.1:  # Humidification/Dehumidification
            if delta_w > 0:
                process_type = "Humidification"
            else:
                process_type = "Dehumidification"
        elif delta_t > 0 and delta_w > 0:
            process_type = "Heating and Humidification"
        elif delta_t < 0 and delta_w < 0:
            process_type = "Cooling and Dehumidification"
        elif delta_t > 0 and delta_w < 0:
            process_type = "Heating and Dehumidification"
        elif delta_t < 0 and delta_w > 0:
            process_type = "Cooling and Humidification"
        
        # Create figure
        fig = go.Figure()
        
        # Add process to psychrometric chart
        chart_fig = self.create_psychrometric_chart(
            points=[
                {"temp": start_temp, "rh": start_rh, "name": "Start", "color": "blue"},
                {"temp": end_temp, "rh": end_rh, "name": "End", "color": "red"}
            ],
            processes=[
                {"start": {"temp": start_temp, "rh": start_rh}, 
                 "end": {"temp": end_temp, "rh": end_rh}, 
                 "name": process_type, 
                 "color": "green"}
            ]
        )
        
        # Create process diagram
        # Create data for process parameters
        params = [
            "Dry-Bulb Temperature (°C)",
            "Relative Humidity (%)",
            "Humidity Ratio (g/kg)",
            "Enthalpy (kJ/kg)",
            "Wet-Bulb Temperature (°C)",
            "Dew Point Temperature (°C)",
            "Specific Volume (m³/kg)"
        ]
        
        start_values = [
            start_props["dry_bulb_temperature"],
            start_props["relative_humidity"],
            start_props["humidity_ratio"] * 1000,  # Convert to g/kg
            start_props["enthalpy"] / 1000,  # Convert to kJ/kg
            start_props["wet_bulb_temperature"],
            start_props["dew_point_temperature"],
            start_props["specific_volume"]
        ]
        
        end_values = [
            end_props["dry_bulb_temperature"],
            end_props["relative_humidity"],
            end_props["humidity_ratio"] * 1000,  # Convert to g/kg
            end_props["enthalpy"] / 1000,  # Convert to kJ/kg
            end_props["wet_bulb_temperature"],
            end_props["dew_point_temperature"],
            end_props["specific_volume"]
        ]
        
        delta_values = [end - start for start, end in zip(start_values, end_values)]
        
        # Create table
        table_fig = go.Figure(data=[go.Table(
            header=dict(
                values=["Parameter", "Start", "End", "Change"],
                fill_color="paleturquoise",
                align="left",
                font=dict(size=12)
            ),
            cells=dict(
                values=[
                    params,
                    [f"{val:.2f}" for val in start_values],
                    [f"{val:.2f}" for val in end_values],
                    [f"{val:.2f}" for val in delta_values]
                ],
                fill_color="lavender",
                align="left",
                font=dict(size=11)
            )
        )])
        
        table_fig.update_layout(
            title=f"Process Parameters: {process_type}",
            height=300,
            margin=dict(l=0, r=0, b=0, t=30)
        )
        
        return chart_fig, table_fig
    
    def display_psychrometric_visualization(self) -> None:
        """
        Display psychrometric visualization in Streamlit.
        """
        st.header("Psychrometric Visualization")
        
        # Create tabs for different visualizations
        tab1, tab2, tab3 = st.tabs([
            "Interactive Psychrometric Chart", 
            "Process Visualization", 
            "Comfort Zone Analysis"
        ])
        
        with tab1:
            st.subheader("Interactive Psychrometric Chart")
            
            # Add controls for points
            st.write("Add points to the chart:")
            
            col1, col2, col3 = st.columns(3)
            
            with col1:
                point1_temp = st.number_input("Point 1 Temperature (°C)", -10.0, 50.0, 20.0, key="point1_temp")
                point1_rh = st.number_input("Point 1 RH (%)", 0.0, 100.0, 50.0, key="point1_rh")
            
            with col2:
                point2_temp = st.number_input("Point 2 Temperature (°C)", -10.0, 50.0, 30.0, key="point2_temp")
                point2_rh = st.number_input("Point 2 RH (%)", 0.0, 100.0, 40.0, key="point2_rh")
            
            with col3:
                show_process = st.checkbox("Show Process Line", True, key="show_process")
                process_name = st.text_input("Process Name", "Cooling Process", key="process_name")
            
            # Create points
            points = [
                {"temp": point1_temp, "rh": point1_rh, "name": "Point 1", "color": "blue"},
                {"temp": point2_temp, "rh": point2_rh, "name": "Point 2", "color": "red"}
            ]
            
            # Create process if enabled
            processes = []
            if show_process:
                processes.append({
                    "start": {"temp": point1_temp, "rh": point1_rh},
                    "end": {"temp": point2_temp, "rh": point2_rh},
                    "name": process_name,
                    "color": "green"
                })
            
            # Create and display chart
            fig = self.create_psychrometric_chart(points=points, processes=processes)
            st.plotly_chart(fig, use_container_width=True)
            
            # Display point properties
            col1, col2 = st.columns(2)
            
            with col1:
                st.subheader("Point 1 Properties")
                props1 = self.psychrometrics.moist_air_properties(point1_temp, point1_rh, self.pressure)
                st.write(f"Dry-Bulb Temperature: {props1['dry_bulb_temperature']:.2f} °C")
                st.write(f"Relative Humidity: {props1['relative_humidity']:.2f} %")
                st.write(f"Humidity Ratio: {props1['humidity_ratio']*1000:.2f} g/kg")
                st.write(f"Enthalpy: {props1['enthalpy']/1000:.2f} kJ/kg")
                st.write(f"Wet-Bulb Temperature: {props1['wet_bulb_temperature']:.2f} °C")
                st.write(f"Dew Point Temperature: {props1['dew_point_temperature']:.2f} °C")
            
            with col2:
                st.subheader("Point 2 Properties")
                props2 = self.psychrometrics.moist_air_properties(point2_temp, point2_rh, self.pressure)
                st.write(f"Dry-Bulb Temperature: {props2['dry_bulb_temperature']:.2f} °C")
                st.write(f"Relative Humidity: {props2['relative_humidity']:.2f} %")
                st.write(f"Humidity Ratio: {props2['humidity_ratio']*1000:.2f} g/kg")
                st.write(f"Enthalpy: {props2['enthalpy']/1000:.2f} kJ/kg")
                st.write(f"Wet-Bulb Temperature: {props2['wet_bulb_temperature']:.2f} °C")
                st.write(f"Dew Point Temperature: {props2['dew_point_temperature']:.2f} °C")
        
        with tab2:
            st.subheader("Process Visualization")
            
            # Add controls for process
            st.write("Define a psychrometric process:")
            
            col1, col2 = st.columns(2)
            
            with col1:
                st.write("Starting Point")
                start_temp = st.number_input("Temperature (°C)", -10.0, 50.0, 24.0, key="start_temp")
                start_rh = st.number_input("RH (%)", 0.0, 100.0, 50.0, key="start_rh")
            
            with col2:
                st.write("Ending Point")
                end_temp = st.number_input("Temperature (°C)", -10.0, 50.0, 14.0, key="end_temp")
                end_rh = st.number_input("RH (%)", 0.0, 100.0, 90.0, key="end_rh")
            
            # Create process
            process = {
                "start": {"temp": start_temp, "rh": start_rh},
                "end": {"temp": end_temp, "rh": end_rh}
            }
            
            # Create and display process visualization
            chart_fig, table_fig = self.create_process_visualization(process)
            
            st.plotly_chart(chart_fig, use_container_width=True)
            st.plotly_chart(table_fig, use_container_width=True)
            
            # Calculate process energy requirements
            start_props = self.psychrometrics.moist_air_properties(start_temp, start_rh, self.pressure)
            end_props = self.psychrometrics.moist_air_properties(end_temp, end_rh, self.pressure)
            
            delta_h = end_props["enthalpy"] - start_props["enthalpy"]  # J/kg
            
            st.subheader("Energy Calculations")
            
            air_flow = st.number_input("Air Flow Rate (m³/s)", 0.1, 100.0, 1.0, key="air_flow")
            
            # Calculate mass flow rate
            density = start_props["density"]  # kg/m³
            mass_flow = air_flow * density  # kg/s
            
            # Calculate energy rate
            energy_rate = mass_flow * delta_h  # W
            
            st.write(f"Air Density: {density:.2f} kg/m³")
            st.write(f"Mass Flow Rate: {mass_flow:.2f} kg/s")
            st.write(f"Enthalpy Change: {delta_h/1000:.2f} kJ/kg")
            st.write(f"Energy Rate: {energy_rate/1000:.2f} kW")
        
        with tab3:
            st.subheader("Comfort Zone Analysis")
            
            # Add controls for comfort zone
            st.write("Define comfort zone parameters:")
            
            col1, col2 = st.columns(2)
            
            with col1:
                temp_min = st.number_input("Minimum Temperature (°C)", 10.0, 30.0, 20.0, key="temp_min")
                temp_max = st.number_input("Maximum Temperature (°C)", 10.0, 30.0, 26.0, key="temp_max")
            
            with col2:
                rh_min = st.number_input("Minimum RH (%)", 0.0, 100.0, 30.0, key="rh_min")
                rh_max = st.number_input("Maximum RH (%)", 0.0, 100.0, 60.0, key="rh_max")
            
            # Create comfort zone
            comfort_zone = {
                "temp_min": temp_min,
                "temp_max": temp_max,
                "rh_min": rh_min,
                "rh_max": rh_max
            }
            
            # Add point to check if it's in comfort zone
            st.write("Check if a point is within the comfort zone:")
            
            col1, col2 = st.columns(2)
            
            with col1:
                check_temp = st.number_input("Temperature (°C)", -10.0, 50.0, 22.0, key="check_temp")
                check_rh = st.number_input("RH (%)", 0.0, 100.0, 45.0, key="check_rh")
            
            # Check if point is in comfort zone
            in_comfort_zone = (
                temp_min <= check_temp <= temp_max and
                rh_min <= check_rh <= rh_max
            )
            
            with col2:
                if in_comfort_zone:
                    st.success("✅ Point is within the comfort zone")
                else:
                    st.error("❌ Point is outside the comfort zone")
                
                # Calculate properties
                check_props = self.psychrometrics.moist_air_properties(check_temp, check_rh, self.pressure)
                st.write(f"Humidity Ratio: {check_props['humidity_ratio']*1000:.2f} g/kg")
                st.write(f"Enthalpy: {check_props['enthalpy']/1000:.2f} kJ/kg")
                st.write(f"Wet-Bulb Temperature: {check_props['wet_bulb_temperature']:.2f} °C")
            
            # Create and display chart with comfort zone
            fig = self.create_psychrometric_chart(
                points=[{"temp": check_temp, "rh": check_rh, "name": "Test Point", "color": "purple"}],
                comfort_zone=comfort_zone
            )
            
            st.plotly_chart(fig, use_container_width=True)


# Create a singleton instance
psychrometric_visualization = PsychrometricVisualization()

# Example usage
if __name__ == "__main__":
    import streamlit as st
    
    # Display psychrometric visualization
    psychrometric_visualization.display_psychrometric_visualization()