data/climate_data.py

"""
ASHRAE 169 climate data module for HVAC Load Calculator.
This module provides access to climate data for various locations based on ASHRAE 169 standard.

Author: Dr Majed Abuseif
Date: March 2025
Version: 1.0.0
"""

from typing import Dict, List, Any, Optional
import pandas as pd
import numpy as np
import os
import json
from dataclasses import dataclass
import streamlit as st
import plotly.graph_objects as go
from io import StringIO

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

@dataclass
class ClimateLocation:
    """Class representing a climate location with ASHRAE 169 data."""
    
    id: str
    country: str
    state_province: str
    city: str
    latitude: float
    longitude: float
    elevation: float  # meters
    climate_zone: str
    heating_degree_days: float  # base 18°C
    cooling_degree_days: float  # base 18°C
    winter_design_temp: float  # 99.6% heating design temperature (°C)
    summer_design_temp_db: float  # 0.4% cooling design dry-bulb temperature (°C)
    summer_design_temp_wb: float  # 0.4% cooling design wet-bulb temperature (°C)
    summer_daily_range: float  # Mean daily temperature range in summer (°C)
    monthly_temps: Dict[str, float]  # Average monthly temperatures (°C)
    monthly_humidity: Dict[str, float]  # Average monthly relative humidity (%)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert the climate location to a dictionary."""
        return {
            "id": self.id,
            "country": self.country,
            "state_province": self.state_province,
            "city": self.city,
            "latitude": self.latitude,
            "longitude": self.longitude,
            "elevation": self.elevation,
            "climate_zone": self.climate_zone,
            "heating_degree_days": self.heating_degree_days,
            "cooling_degree_days": self.cooling_degree_days,
            "winter_design_temp": self.winter_design_temp,
            "summer_design_temp_db": self.summer_design_temp_db,
            "summer_design_temp_wb": self.summer_design_temp_wb,
            "summer_daily_range": self.summer_daily_range,
            "monthly_temps": self.monthly_temps,
            "monthly_humidity": self.monthly_humidity
        }

class ClimateData:
    """Class for managing ASHRAE 169 climate data."""
    
    def __init__(self):
        """Initialize climate data."""
        self.locations = {}
        self.countries = []
        self.country_states = {}
    
    def _group_locations_by_country_state(self) -> Dict[str, Dict[str, List[str]]]:
        """Group locations by country and state/province."""
        result = {}
        for loc in self.locations.values():
            if loc.country not in result:
                result[loc.country] = {}
            if loc.state_province not in result[loc.country]:
                result[loc.country][loc.state_province] = []
            result[loc.country][loc.state_province].append(loc.city)
        for country in result:
            for state in result[country]:
                result[country][state] = sorted(result[country][state])
        return result
    
    def add_location(self, location: ClimateLocation):
        """Add a new location to the dictionary."""
        self.locations[location.id] = location
        self.countries = sorted(list(set(loc.country for loc in self.locations.values())))
        self.country_states = self._group_locations_by_country_state()

    def get_location_by_id(self, location_id: str, session_state: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """Retrieve climate data by ID from session state or locations."""
        if "climate_data" in session_state and session_state["climate_data"].get("id") == location_id:
            return session_state["climate_data"]
        if location_id in self.locations:
            return self.locations[location_id].to_dict()
        return None

    @staticmethod
    def validate_climate_data(data: Dict[str, Any]) -> bool:
        """Validate climate data for required fields and ranges."""
        required_fields = [
            "id", "country", "city", "latitude", "longitude", "elevation",
            "climate_zone", "heating_degree_days", "cooling_degree_days",
            "winter_design_temp", "summer_design_temp_db", "summer_design_temp_wb",
            "summer_daily_range", "monthly_temps", "monthly_humidity"
        ]
        month_names = ["Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"]
        
        for field in required_fields:
            if field not in data:
                return False
        
        if not (-90 <= data["latitude"] <= 90 and -180 <= data["longitude"] <= 180):
            return False
        if data["elevation"] < 0:
            return False
        if data["climate_zone"] not in ["0A", "0B", "1A", "1B", "2A", "2B", "3A", "3B", "3C", "4A", "4B", "4C", "5A", "5B", "5C", "6A", "6B", "7", "8"]:
            return False
        if not (data["heating_degree_days"] >= 0 and data["cooling_degree_days"] >= 0):
            return False
        if not (-50 <= data["winter_design_temp"] <= 20):
            return False
        if not (0 <= data["summer_design_temp_db"] <= 50 and 0 <= data["summer_design_temp_wb"] <= 40):
            return False
        if data["summer_daily_range"] < 0:
            return False
        
        for month in month_names:
            if month not in data["monthly_temps"] or month not in data["monthly_humidity"]:
                return False
            if not (-50 <= data["monthly_temps"][month] <= 50):
                return False
            if not (0 <= data["monthly_humidity"][month] <= 100):
                return False
        
        return True

    @staticmethod
    def calculate_wet_bulb(dry_bulb: np.ndarray, relative_humidity: np.ndarray) -> np.ndarray:
        """Calculate Wet Bulb Temperature using Stull (2011) approximation."""
        db = np.array(dry_bulb, dtype=float)
        rh = np.array(relative_humidity, dtype=float)
        
        term1 = db * np.arctan(0.151977 * (rh + 8.313659)**0.5)
        term2 = np.arctan(db + rh)
        term3 = np.arctan(rh - 1.676331)
        term4 = 0.00391838 * rh**1.5 * np.arctan(0.023101 * rh)
        term5 = -4.686035
        
        wet_bulb = term1 + term2 - term3 + term4 + term5
        
        invalid_mask = (rh < 5) | (rh > 99) | (db < -20) | (db > 50) | np.isnan(db) | np.isnan(rh)
        wet_bulb[invalid_mask] = np.nan
        
        return wet_bulb

    def display_climate_input(self, session_state: Dict[str, Any]):
        """Display form for manual input or EPW upload in Streamlit."""
        st.title("Climate Data")
        
        if not session_state.building_info.get("country") or not session_state.building_info.get("city"):
            st.warning("Please enter country and city in Building Information first.")
            st.button("Go to Building Information", on_click=lambda: setattr(session_state, "page", "Building Information"))
            return
        
        st.subheader(f"Location: {session_state.building_info['country']}, {session_state.building_info['city']}")
        tab1, tab2 = st.tabs(["Manual Input", "Upload EPW File"])
        
        # Manual Input Tab
        with tab1:
            with st.form("manual_climate_form"):
                col1, col2 = st.columns(2)
                with col1:
                    latitude = st.number_input(
                        "Latitude",
                        min_value=-90.0,
                        max_value=90.0,
                        value=0.0,
                        step=0.1,
                        help="Enter the latitude of the location in degrees (e.g., 64.1 for Reykjavik)"
                    )
                    longitude = st.number_input(
                        "Longitude",
                        min_value=-180.0,
                        max_value=180.0,
                        value=0.0,
                        step=0.1,
                        help="Enter the longitude of the location in degrees (e.g., -21.9 for Reykjavik)"
                    )
                    elevation = st.number_input(
                        "Elevation (m)",
                        min_value=0.0,
                        value=0.0,
                        step=10.0,
                        help="Enter the elevation of the location above sea level in meters"
                    )
                    climate_zone = st.selectbox(
                        "Climate Zone",
                        ["0A", "0B", "1A", "1B", "2A", "2B", "3A", "3B", "3C", "4A", "4B", "4C", "5A", "5B", "5C", "6A", "6B", "7", "8"],
                        help="Select the ASHRAE climate zone for the location (e.g., 6A for cold, humid climates)"
                    )
                
                with col2:
                    hdd = st.number_input(
                        "Heating Degree Days (base 18°C)",
                        min_value=0.0,
                        value=0.0,
                        step=100.0,
                        help="Enter the annual heating degree days using an 18°C base temperature"
                    )
                    cdd = st.number_input(
                        "Cooling Degree Days (base 18°C)",
                        min_value=0.0,
                        value=0.0,
                        step=100.0,
                        help="Enter the annual cooling degree days using an 18°C base temperature"
                    )
                    winter_design_temp = st.number_input(
                        "Winter Design Temp (99.6%) (°C)",
                        min_value=-50.0,
                        max_value=20.0,
                        value=0.0,
                        step=0.5,
                        help="Enter the 99.6% winter design temperature in °C (extreme cold condition)"
                    )
                    summer_design_temp_db = st.number_input(
                        "Summer Design Temp DB (0.4%) (°C)",
                        min_value=0.0,
                        max_value=50.0,
                        value=35.0,
                        step=0.5,
                        help="Enter the 0.4% summer design dry-bulb temperature in °C (extreme hot condition)"
                    )
                    summer_design_temp_wb = st.number_input(
                        "Summer Design Temp WB (0.4%) (°C)",
                        min_value=0.0,
                        max_value=40.0,
                        value=25.0,
                        step=0.5,
                        help="Enter the 0.4% summer design wet-bulb temperature in °C (for humidity consideration)"
                    )
                    summer_daily_range = st.number_input(
                        "Summer Daily Range (°C)",
                        min_value=0.0,
                        value=5.0,
                        step=0.5,
                        help="Enter the average daily temperature range in summer in °C"
                    )
                
                # Monthly Data with clear titles (no help added here)
                monthly_temps = {}
                monthly_humidity = {}
                month_names = ["Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"]
                
                st.subheader("Monthly Temperatures")
                col1, col2 = st.columns(2)
                with col1:
                    for month in month_names[:6]:
                        monthly_temps[month] = st.number_input(f"{month} Temp (°C)", min_value=-50.0, max_value=50.0, value=20.0, step=0.5, key=f"temp_{month}")
                with col2:
                    for month in month_names[6:]:
                        monthly_temps[month] = st.number_input(f"{month} Temp (°C)", min_value=-50.0, max_value=50.0, value=20.0, step=0.5, key=f"temp_{month}")
                
                st.subheader("Monthly Humidity")
                col1, col2 = st.columns(2)
                with col1:
                    for month in month_names[:6]:
                        monthly_humidity[month] = st.number_input(f"{month} Humidity (%)", min_value=0.0, max_value=100.0, value=50.0, step=5.0, key=f"hum_{month}")
                with col2:
                    for month in month_names[6:]:
                        monthly_humidity[month] = st.number_input(f"{month} Humidity (%)", min_value=0.0, max_value=100.0, value=50.0, step=5.0, key=f"hum_{month}")
                
                if st.form_submit_button("Save Climate Data"):
                    try:
                        # Generate ID internally using country and city from session_state
                        generated_id = f"{session_state.building_info['country'][:2].upper()}-{session_state.building_info['city'][:3].upper()}"
                        location = ClimateLocation(
                            id=generated_id,
                            country=session_state.building_info["country"],
                            state_province="N/A",  # Default since input removed
                            city=session_state.building_info["city"],
                            latitude=latitude,
                            longitude=longitude,
                            elevation=elevation,
                            climate_zone=climate_zone,
                            heating_degree_days=hdd,
                            cooling_degree_days=cdd,
                            winter_design_temp=winter_design_temp,
                            summer_design_temp_db=summer_design_temp_db,
                            summer_design_temp_wb=summer_design_temp_wb,
                            summer_daily_range=summer_daily_range,
                            monthly_temps=monthly_temps,
                            monthly_humidity=monthly_humidity
                        )
                        self.add_location(location)
                        climate_data_dict = location.to_dict()
                        if not self.validate_climate_data(climate_data_dict):
                            raise ValueError("Invalid climate data. Please check all inputs.")
                        session_state["climate_data"] = climate_data_dict  # Save to session state
                        st.success("Climate data saved manually!")
                        st.write(f"Debug: Saved climate data for {location.city} (ID: {location.id}): {climate_data_dict}")  # Debug
                        self.display_design_conditions(location)
                        self.visualize_data(location, epw_data=None)
                    except Exception as e:
                        st.error(f"Error saving climate data: {str(e)}. Please check inputs and try again.")

        # EPW Upload Tab
        with tab2:
            uploaded_file = st.file_uploader("Upload EPW File", type=["epw"])
            if uploaded_file:
                try:
                    epw_content = uploaded_file.read().decode("utf-8")
                    epw_lines = epw_content.splitlines()
                    header = next(line for line in epw_lines if line.startswith("LOCATION"))
                    header_parts = header.split(",")
                    latitude = float(header_parts[6])
                    longitude = float(header_parts[7])
                    elevation = float(header_parts[8])
                    
                    data_start_idx = next(i for i, line in enumerate(epw_lines) if line.startswith("DATA PERIODS")) + 1
                    epw_data = pd.read_csv(StringIO("\n".join(epw_lines[data_start_idx:])), header=None, dtype=str)
                    if len(epw_data) != 8760:
                        raise ValueError(f"EPW file has {len(epw_data)} records, expected 8760.")
                    
                    for col in epw_data.columns:
                        epw_data[col] = pd.to_numeric(epw_data[col], errors='coerce')
                    
                    months = epw_data[1].values  # Month
                    dry_bulb = epw_data[6].values  # Dry-bulb temperature (°C)
                    humidity = epw_data[8].values  # Relative humidity (%)
                    pressure = epw_data[9].values  # Atmospheric pressure (Pa)
                    
                    wet_bulb = self.calculate_wet_bulb(dry_bulb, humidity)
                    
                    if np.all(np.isnan(dry_bulb)) or np.all(np.isnan(humidity)) or np.all(np.isnan(wet_bulb)):
                        raise ValueError("Dry bulb, humidity, or calculated wet bulb data is entirely NaN.")
                    
                    daily_temps = np.nanmean(dry_bulb.reshape(-1, 24), axis=1)
                    hdd = round(np.nansum(np.maximum(18 - daily_temps, 0)))
                    cdd = round(np.nansum(np.maximum(daily_temps - 18, 0)))
                    
                    winter_design_temp = round(np.nanpercentile(dry_bulb, 0.4), 1)
                    summer_design_temp_db = round(np.nanpercentile(dry_bulb, 99.6), 1)
                    summer_design_temp_wb = round(np.nanpercentile(wet_bulb, 99.6), 1)
                    summer_mask = (months >= 6) & (months <= 8)
                    summer_temps = dry_bulb[summer_mask].reshape(-1, 24)
                    summer_daily_range = round(np.nanmean(np.nanmax(summer_temps, axis=1) - np.nanmin(summer_temps, axis=1)), 1)
                    
                    monthly_temps = {}
                    monthly_humidity = {}
                    month_names = ["Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"]
                    for i in range(1, 13):
                        month_mask = (months == i)
                        monthly_temps[month_names[i-1]] = round(np.nanmean(dry_bulb[month_mask]), 1)
                        monthly_humidity[month_names[i-1]] = round(np.nanmean(humidity[month_mask]), 1)
                    
                    avg_humidity = np.nanmean(humidity)
                    climate_zone = self.assign_climate_zone(hdd, cdd, avg_humidity)
                    
                    location = ClimateLocation(
                        id=f"{session_state.building_info['country'][:2].upper()}-{session_state.building_info['city'][:3].upper()}",
                        country=session_state.building_info["country"],
                        state_province="N/A",
                        city=session_state.building_info["city"],
                        latitude=latitude,
                        longitude=longitude,
                        elevation=elevation,
                        climate_zone=climate_zone,
                        heating_degree_days=hdd,
                        cooling_degree_days=cdd,
                        winter_design_temp=winter_design_temp,
                        summer_design_temp_db=summer_design_temp_db,
                        summer_design_temp_wb=summer_design_temp_wb,
                        summer_daily_range=summer_daily_range,
                        monthly_temps=monthly_temps,
                        monthly_humidity=monthly_humidity
                    )
                    self.add_location(location)
                    climate_data_dict = location.to_dict()
                    if not self.validate_climate_data(climate_data_dict):
                        raise ValueError("Invalid climate data extracted from EPW file.")
                    session_state["climate_data"] = climate_data_dict  # Save to session state
                    st.success("Climate data extracted from EPW file with calculated Wet Bulb Temperature!")
                    st.write(f"Debug: Saved climate data for {location.city} (ID: {location.id}): {climate_data_dict}")  # Debug
                    self.display_design_conditions(location)
                    self.visualize_data(location, epw_data=epw_data)
                except Exception as e:
                    st.error(f"Error processing EPW file: {str(e)}. Ensure it has 8760 hourly records and correct format.")

        col1, col2 = st.columns(2)
        with col1:
            st.button("Back to Building Information", on_click=lambda: setattr(session_state, "page", "Building Information"))
        with col2:
            if self.locations:
                st.button("Continue to Building Components", on_click=lambda: setattr(session_state, "page", "Building Components"))
            else:
                st.button("Continue to Building Components", disabled=True)

        # Display saved session state data (if any)
        if "climate_data" in session_state and session_state["climate_data"]:
            st.subheader("Saved Climate Data")
            st.json(session_state["climate_data"])  # Display as JSON for clarity

    def display_design_conditions(self, location: ClimateLocation):
        """Display a table of design conditions including additional parameters for HVAC calculations."""
        st.subheader("Design Conditions for HVAC Calculations")
        
        design_data = pd.DataFrame({
            "Parameter": [
                "Latitude",
                "Longitude",
                "Elevation (m)",
                "Climate Zone",
                "Heating Degree Days (base 18°C)",
                "Cooling Degree Days (base 18°C)",
                "Winter Design Temperature (99.6%)",
                "Summer Design Dry-Bulb Temp (0.4%)",
                "Summer Design Wet-Bulb Temp (0.4%)",
                "Summer Daily Temperature Range"
            ],
            "Value": [
                f"{location.latitude}°",
                f"{location.longitude}°",
                f"{location.elevation} m",
                location.climate_zone,
                f"{location.heating_degree_days} HDD",
                f"{location.cooling_degree_days} CDD",
                f"{location.winter_design_temp} °C",
                f"{location.summer_design_temp_db} °C",
                f"{location.summer_design_temp_wb} °C",
                f"{location.summer_daily_range} °C"
            ]
        })
        
        month_names = ["Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"]
        monthly_temp_data = pd.DataFrame({
            "Parameter": [f"{month} Avg Temp" for month in month_names],
            "Value": [f"{location.monthly_temps[month]} °C" for month in month_names]
        })
        
        monthly_humidity_data = pd.DataFrame({
            "Parameter": [f"{month} Avg Humidity" for month in month_names],
            "Value": [f"{location.monthly_humidity[month]} %" for month in month_names]
        })
        
        full_design_data = pd.concat([design_data, monthly_temp_data, monthly_humidity_data], ignore_index=True)
        st.table(full_design_data)

    @staticmethod
    def assign_climate_zone(hdd: float, cdd: float, avg_humidity: float) -> str:
        """Assign ASHRAE 169 climate zone based on HDD, CDD, and humidity."""
        if cdd > 10000:
            return "0A" if avg_humidity > 60 else "0B"
        elif cdd > 5000:
            return "1A" if avg_humidity > 60 else "1B"
        elif cdd > 2500:
            return "2A" if avg_humidity > 60 else "2B"
        elif hdd < 2000 and cdd > 1000:
            return "3A" if avg_humidity > 60 else "3B" if avg_humidity < 40 else "3C"
        elif hdd < 3000:
            return "4A" if avg_humidity > 60 else "4B" if avg_humidity < 40 else "4C"
        elif hdd < 4000:
            return "5A" if avg_humidity > 60 else "5B" if avg_humidity < 40 else "5C"
        elif hdd < 5000:
            return "6A" if avg_humidity > 60 else "6B"
        elif hdd < 7000:
            return "7"
        else:
            return "8"

    @staticmethod
    def visualize_data(location: ClimateLocation, epw_data: Optional[pd.DataFrame] = None):
        """Visualize monthly temperature and humidity data."""
        st.subheader("Monthly Climate Data Visualization")
        
        months = list(range(1, 13))
        month_names = ["Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"]
        temps_avg = [location.monthly_temps[m] for m in month_names]
        humidity_avg = [location.monthly_humidity[m] for m in month_names]
        
        fig_temp = go.Figure()
        fig_temp.add_trace(go.Scatter(
            x=months,
            y=temps_avg,
            mode='lines+markers',
            name='Avg Temperature (°C)',
            line=dict(color='red'),
            marker=dict(size=8)
        ))
        
        if epw_data is not None:
            dry_bulb = epw_data[6].values
            month_col = epw_data[1].values
            temps_min = []
            temps_max = []
            for i in range(1, 13):
                month_mask = (month_col == i)
                temps_min.append(round(np.nanmin(dry_bulb[month_mask]), 1))
                temps_max.append(round(np.nanmax(dry_bulb[month_mask]), 1))
            fig_temp.add_trace(go.Scatter(
                x=months,
                y=temps_max,
                mode='lines',
                name='Max Temperature (°C)',
                line=dict(color='red', dash='dash'),
                opacity=0.5
            ))
            fig_temp.add_trace(go.Scatter(
                x=months,
                y=temps_min,
                mode='lines',
                name='Min Temperature (°C)',
                line=dict(color='red', dash='dash'),
                opacity=0.5,
                fill='tonexty',
                fillcolor='rgba(255, 0, 0, 0.1)'
            ))
        
        fig_temp.update_layout(
            title='Monthly Temperatures',
            xaxis_title='Month',
            yaxis_title='Temperature (°C)',
            xaxis=dict(tickmode='array', tickvals=months, ticktext=month_names),
            legend=dict(yanchor="top", y=0.99, xanchor="left", x=0.01)
        )
        st.plotly_chart(fig_temp, use_container_width=True)
        
        fig_hum = go.Figure()
        fig_hum.add_trace(go.Scatter(
            x=months,
            y=humidity_avg,
            mode='lines+markers',
            name='Avg Humidity (%)',
            line=dict(color='blue'),
            marker=dict(size=8)
        ))
        
        if epw_data is not None:
            humidity = epw_data[8].values
            month_col = epw_data[1].values
            humidity_min = []
            humidity_max = []
            for i in range(1, 13):
                month_mask = (month_col == i)
                humidity_min.append(round(np.nanmin(humidity[month_mask]), 1))
                humidity_max.append(round(np.nanmax(humidity[month_mask]), 1))
            fig_hum.add_trace(go.Scatter(
                x=months,
                y=humidity_max,
                mode='lines',
                name='Max Humidity (%)',
                line=dict(color='blue', dash='dash'),
                opacity=0.5
            ))
            fig_hum.add_trace(go.Scatter(
                x=months,
                y=humidity_min,
                mode='lines',
                name='Min Humidity (%)',
                line=dict(color='blue', dash='dash'),
                opacity=0.5,
                fill='tonexty',
                fillcolor='rgba(0, 0, 255, 0.1)'
            ))
        
        fig_hum.update_layout(
            title='Monthly Relative Humidity',
            xaxis_title='Month',
            yaxis_title='Relative Humidity (%)',
            xaxis=dict(tickmode='array', tickvals=months, ticktext=month_names),
            legend=dict(yanchor="top", y=0.99, xanchor="left", x=0.01)
        )
        st.plotly_chart(fig_hum, use_container_width=True)

    def export_to_json(self, file_path: str) -> None:
        """Export all climate data to a JSON file."""
        data = {loc_id: loc.to_dict() for loc_id, loc in self.locations.items()}
        with open(file_path, 'w') as f:
            json.dump(data, f, indent=4)

    @classmethod
    def from_json(cls, file_path: str) -> 'ClimateData':
        """Load climate data from a JSON file."""
        with open(file_path, 'r') as f:
            data = json.load(f)
        climate_data = cls()
        for loc_id, loc_dict in data.items():
            location = ClimateLocation(**loc_dict)
            climate_data.add_location(location)
        return climate_data

if __name__ == "__main__":
    climate_data = ClimateData()
    session_state = {"building_info": {"country": "Iceland", "city": "Reyugalvik"}, "page": "Climate Data"}
    climate_data.display_climate_input(session_state)