app/data_persistence.py

"""
Data persistence module for HVAC Load Calculator.
This module provides functionality for saving and loading project data.
"""

import streamlit as st
import pandas as pd
import numpy as np
from typing import Dict, List, Any, Optional, Tuple
import json
import os
import base64
import io
import pickle
from datetime import datetime

# Import data models
from data.building_components import Wall, Roof, Floor, Window, Door, Orientation, ComponentType, Skylight


class DataPersistence:
    """Class for data persistence functionality."""
    
    @staticmethod
    def save_project_to_json(session_state: Dict[str, Any], file_path: str = None) -> Optional[str]:
        """
        Save project data to a JSON file.
        
        Args:
            session_state: Streamlit session state containing project data
            file_path: Optional path to save the JSON file
            
        Returns:
            JSON string if file_path is None, otherwise None
        """
        try:
            # Create project data dictionary
            project_data = {
                "building_info": session_state.get("building_info", {}),
                "components": DataPersistence._serialize_components(session_state.get("components", {})),
                "internal_loads": session_state.get("internal_loads", {}),
                "calculation_settings": session_state.get("calculation_settings", {}),
                "saved_scenarios": DataPersistence._serialize_scenarios(session_state.get("saved_scenarios", {})),
                "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            }
            
            # Convert to JSON
            json_data = json.dumps(project_data, indent=4)
            
            # Save to file if path provided
            if file_path:
                with open(file_path, "w") as f:
                    f.write(json_data)
                return None
            
            # Return JSON string if no path provided
            return json_data
        
        except Exception as e:
            st.error(f"Error saving project data: {e}")
            return None
    
    @staticmethod
    def load_project_from_json(json_data: str = None, file_path: str = None) -> Optional[Dict[str, Any]]:
        """
        Load project data from a JSON file or string.
        
        Args:
            json_data: Optional JSON string containing project data
            file_path: Optional path to the JSON file
            
        Returns:
            Dictionary with project data if successful, None otherwise
        """
        try:
            # Load from file if path provided
            if file_path and not json_data:
                with open(file_path, "r") as f:
                    json_data = f.read()
            
            # Parse JSON data
            if json_data:
                project_data = json.loads(json_data)
                
                # Deserialize components
                if "components" in project_data:
                    project_data["components"] = DataPersistence._deserialize_components(project_data["components"])
                
                # Deserialize scenarios
                if "saved_scenarios" in project_data:
                    project_data["saved_scenarios"] = DataPersistence._deserialize_scenarios(project_data["saved_scenarios"])
                
                return project_data
            
            return None
        
        except Exception as e:
            st.error(f"Error loading project data: {e}")
            return None
    
    @staticmethod
    def _serialize_components(components: Dict[str, List[Any]]) -> Dict[str, List[Dict[str, Any]]]:
        """
        Serialize components for JSON storage.
        
        Args:
            components: Dictionary with building components
            
        Returns:
            Dictionary with serialized components
        """
        serialized_components = {
            "walls": [],
            "roofs": [],
            "floors": [],
            "windows": [],
            "doors": [],
            "skylights": []
        }
        
        # Serialize walls
        for wall in components.get("walls", []):
            serialized_wall = wall.__dict__.copy()
            
            # Convert enums to strings
            if hasattr(serialized_wall["orientation"], "name"):
                serialized_wall["orientation"] = serialized_wall["orientation"].name
            
            if hasattr(serialized_wall["component_type"], "name"):
                serialized_wall["component_type"] = serialized_wall["component_type"].name
            
            serialized_components["walls"].append(serialized_wall)
        
        # Serialize roofs
        for roof in components.get("roofs", []):
            serialized_roof = roof.__dict__.copy()
            
            # Convert enums to strings
            if hasattr(serialized_roof["orientation"], "name"):
                serialized_roof["orientation"] = serialized_roof["orientation"].name
            
            if hasattr(serialized_roof["component_type"], "name"):
                serialized_roof["component_type"] = serialized_roof["component_type"].name
            
            serialized_components["roofs"].append(serialized_roof)
        
        # Serialize floors
        for floor in components.get("floors", []):
            serialized_floor = floor.__dict__.copy()
            
            # Convert enums to strings
            if hasattr(serialized_floor["component_type"], "name"):
                serialized_floor["component_type"] = serialized_floor["component_type"].name
            
            serialized_components["floors"].append(serialized_floor)
        
        # Serialize windows
        for window in components.get("windows", []):
            serialized_window = window.__dict__.copy()
            
            # Convert enums to strings
            if hasattr(serialized_window["orientation"], "name"):
                serialized_window["orientation"] = serialized_window["orientation"].name
            
            if hasattr(serialized_window["component_type"], "name"):
                serialized_window["component_type"] = serialized_window["component_type"].name
            
            serialized_components["windows"].append(serialized_window)
        
        # Serialize doors
        for door in components.get("doors", []):
            serialized_door = door.__dict__.copy()
            
            # Convert enums to strings
            if hasattr(serialized_door["orientation"], "name"):
                serialized_door["orientation"] = serialized_door["orientation"].name
            
            if hasattr(serialized_door["component_type"], "name"):
                serialized_door["component_type"] = serialized_door["component_type"].name
            
            serialized_components["doors"].append(serialized_door)
        
        # Serialize skylights
        for skylight in components.get("skylights", []):
            serialized_skylight = skylight.__dict__.copy()
            
            # Convert enums to strings
            if hasattr(serialized_skylight["orientation"], "name"):
                serialized_skylight["orientation"] = serialized_skylight["orientation"].name
            
            if hasattr(serialized_skylight["component_type"], "name"):
                serialized_skylight["component_type"] = serialized_skylight["component_type"].name
            
            serialized_components["skylights"].append(serialized_skylight)
        
        return serialized_components
    
    @staticmethod
    def _deserialize_components(serialized_components: Dict[str, List[Dict[str, Any]]]) -> Dict[str, List[Any]]:
        """
        Deserialize components from JSON storage.
        
        Args:
            serialized_components: Dictionary with serialized components
            
        Returns:
            Dictionary with deserialized components
        """
        components = {
            "walls": [],
            "roofs": [],
            "floors": [],
            "windows": [],
            "doors": [],
            "skylights": []
        }
        
        # Deserialize walls
        for wall_dict in serialized_components.get("walls", []):
            wall = Wall(
                id=wall_dict.get("id", ""),
                name=wall_dict.get("name", ""),
                component_type=ComponentType[wall_dict.get("component_type", "WALL")],
                u_value=wall_dict.get("u_value", 0.0),
                area=wall_dict.get("area", 0.0),
                orientation=Orientation[wall_dict.get("orientation", "NORTH")],
                wall_type=wall_dict.get("wall_type", ""),
                wall_group=wall_dict.get("wall_group", ""),
                solar_absorptivity=wall_dict.get("solar_absorptivity", 0.6)
            )
            components["walls"].append(wall)
        
        # Deserialize roofs
        for roof_dict in serialized_components.get("roofs", []):
            roof = Roof(
                id=roof_dict.get("id", ""),
                name=roof_dict.get("name", ""),
                component_type=ComponentType[roof_dict.get("component_type", "ROOF")],
                u_value=roof_dict.get("u_value", 0.0),
                area=roof_dict.get("area", 0.0),
                orientation=Orientation[roof_dict.get("orientation", "HORIZONTAL")],
                roof_type=roof_dict.get("roof_type", ""),
                roof_group=roof_dict.get("roof_group", ""),
                solar_absorptivity=roof_dict.get("solar_absorptivity", 0.6)
            )
            components["roofs"].append(roof)
        
        # Deserialize floors
        for floor_dict in serialized_components.get("floors", []):
            floor = Floor(
                id=floor_dict.get("id", ""),
                name=floor_dict.get("name", ""),
                component_type=ComponentType[floor_dict.get("component_type", "FLOOR")],
                u_value=floor_dict.get("u_value", 0.0),
                area=floor_dict.get("area", 0.0),
                floor_type=floor_dict.get("floor_type", ""),
                solar_absorptivity=floor_dict.get("solar_absorptivity", 0.6)
            )
            components["floors"].append(floor)
        
        # Deserialize windows
        for window_dict in serialized_components.get("windows", []):
            window = Window(
                id=window_dict.get("id", ""),
                name=window_dict.get("name", ""),
                component_type=ComponentType[window_dict.get("component_type", "WINDOW")],
                u_value=window_dict.get("u_value", 0.0),
                area=window_dict.get("area", 0.0),
                orientation=Orientation[window_dict.get("orientation", "NORTH")],
                shgc=window_dict.get("shgc", 0.0),
                vt=window_dict.get("vt", 0.0),
                window_type=window_dict.get("window_type", ""),
                glazing_layers=window_dict.get("glazing_layers", 1),
                gas_fill=window_dict.get("gas_fill", ""),
                low_e_coating=window_dict.get("low_e_coating", False),
                solar_absorptivity=window_dict.get("solar_absorptivity", 0.6)
            )
            components["windows"].append(window)
        
        # Deserialize doors
        for door_dict in serialized_components.get("doors", []):
            door = Door(
                id=door_dict.get("id", ""),
                name=door_dict.get("name", ""),
                component_type=ComponentType[door_dict.get("component_type", "DOOR")],
                u_value=door_dict.get("u_value", 0.0),
                area=door_dict.get("area", 0.0),
                orientation=Orientation[door_dict.get("orientation", "NORTH")],
                door_type=door_dict.get("door_type", ""),
                solar_absorptivity=door_dict.get("solar_absorptivity", 0.6)
            )
            components["doors"].append(door)
        
        # Deserialize skylights
        for skylight_dict in serialized_components.get("skylights", []):
            skylight = Skylight(
                id=skylight_dict.get("id", ""),
                name=skylight_dict.get("name", ""),
                component_type=ComponentType[skylight_dict.get("component_type", "SKYLIGHT")],
                u_value=skylight_dict.get("u_value", 0.0),
                area=skylight_dict.get("area", 0.0),
                orientation=Orientation[skylight_dict.get("orientation", "HORIZONTAL")],
                shgc=skylight_dict.get("shgc", 0.0),
                vt=skylight_dict.get("vt", 0.0),
                skylight_type=skylight_dict.get("skylight_type", ""),
                frame_type=skylight_dict.get("frame_type", ""),
                drapery_openness=skylight_dict.get("drapery_openness", "Open"),
                solar_absorptivity=skylight_dict.get("solar_absorptivity", 0.6)
            )
            components["skylights"].append(skylight)
        
        return components
    
    @staticmethod
    def _serialize_scenarios(scenarios: Dict[str, Dict[str, Any]]) -> Dict[str, Dict[str, Any]]:
        """
        Serialize scenarios for JSON storage.
        
        Args:
            scenarios: Dictionary with saved scenarios
            
        Returns:
            Dictionary with serialized scenarios
        """
        serialized_scenarios = {}
        
        for scenario_name, scenario_data in scenarios.items():
            serialized_scenario = {
                "results": scenario_data.get("results", {}),
                "building_info": scenario_data.get("building_info", {}),
                "components": DataPersistence._serialize_components(scenario_data.get("components", {})),
                "timestamp": scenario_data.get("timestamp", datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
            }
            
            serialized_scenarios[scenario_name] = serialized_scenario
        
        return serialized_scenarios
    
    @staticmethod
    def _deserialize_scenarios(serialized_scenarios: Dict[str, Dict[str, Any]]) -> Dict[str, Dict[str, Any]]:
        """
        Deserialize scenarios from JSON storage.
        
        Args:
            serialized_scenarios: Dictionary with serialized scenarios
            
        Returns:
            Dictionary with deserialized scenarios
        """
        scenarios = {}
        
        for scenario_name, serialized_scenario in serialized_scenarios.items():
            scenario = {
                "results": serialized_scenario.get("results", {}),
                "building_info": serialized_scenario.get("building_info", {}),
                "components": DataPersistence._deserialize_components(serialized_scenario.get("components", {})),
                "timestamp": serialized_scenario.get("timestamp", datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
            }
            
            scenarios[scenario_name] = scenario
        
        return scenarios
    
    @staticmethod
    def get_download_link(data: str, filename: str, text: str) -> str:
        """
        Generate a download link for data.
        
        Args:
            data: Data to download
            filename: Name of the file to download
            text: Text to display for the download link
            
        Returns:
            HTML string with download link
        """
        b64 = base64.b64encode(data.encode()).decode()
        href = f'<a href="data:file/txt;base64,{b64}" download="{filename}">{text}</a>'
        return href
    
    @staticmethod
    def display_project_management(session_state: Dict[str, Any]) -> None:
        """
        Display project management interface in Streamlit.
        
        Args:
            session_state: Streamlit session state containing project data
        """
        st.header("Project Management")
        
        # Create tabs for different project management functions
        tab1, tab2, tab3 = st.tabs(["Save Project", "Load Project", "Project History"])
        
        with tab1:
            DataPersistence._display_save_project(session_state)
        
        with tab2:
            DataPersistence._display_load_project(session_state)
        
        with tab3:
            DataPersistence._display_project_history(session_state)
    
    @staticmethod
    def _display_save_project(session_state: Dict[str, Any]) -> None:
        """
        Display save project interface.
        
        Args:
            session_state: Streamlit session state containing project data
        """
        st.subheader("Save Project")
        
        # Get project name
        project_name = st.text_input(
            "Project Name",
            value=session_state.get("building_info", {}).get("project_name", "HVAC_Project"),
            key="save_project_name"
        )
        
        # Add description
        project_description = st.text_area(
            "Project Description",
            value=session_state.get("project_description", ""),
            key="save_project_description"
        )
        
        # Save project description
        session_state["project_description"] = project_description
        
        # Add save button
        if st.button("Save Project"):
            # Validate project data
            if "building_info" not in session_state or not session_state["building_info"]:
                st.error("No building information found. Please enter building information before saving.")
                return
            
            if "components" not in session_state or not any(session_state["components"].values()):
                st.warning("No building components found. It's recommended to add components before saving.")
            
            # Save project data to JSON
            json_data = DataPersistence.save_project_to_json(session_state)
            
            if json_data:
                # Generate download link
                filename = f"{project_name.replace(' ', '_')}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.hvac"
                download_link = DataPersistence.get_download_link(json_data, filename, "Download Project File")
                
                # Display download link
                st.success("Project saved successfully!")
                st.markdown(download_link, unsafe_allow_html=True)
                
                # Save to project history
                if "project_history" not in session_state:
                    session_state["project_history"] = []
                
                session_state["project_history"].append({
                    "name": project_name,
                    "description": project_description,
                    "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                    "data": json_data
                })
            else:
                st.error("Error saving project data.")
    
    @staticmethod
    def _display_load_project(session_state: Dict[str, Any]) -> None:
        """
        Display load project interface.
        
        Args:
            session_state: Streamlit session state containing project data
        """
        st.subheader("Load Project")
        
        # Add file uploader
        uploaded_file = st.file_uploader("Upload Project File", type=["hvac", "json"])
        
        if uploaded_file is not None:
            # Read file content
            json_data = uploaded_file.read().decode("utf-8")
            
            # Load project data
            project_data = DataPersistence.load_project_from_json(json_data)
            
            if project_data:
                # Add load button
                if st.button("Load Project Data"):
                    # Update session state with project data
                    for key, value in project_data.items():
                        session_state[key] = value
                    
                    st.success("Project loaded successfully!")
                    st.experimental_rerun()
            else:
                st.error("Error loading project data. Invalid file format.")
    
    @staticmethod
    def _display_project_history(session_state: Dict[str, Any]) -> None:
        """
        Display project history interface.
        
        Args:
            session_state: Streamlit session state containing project data
        """
        st.subheader("Project History")
        
        # Check if project history exists
        if "project_history" not in session_state or not session_state["project_history"]:
            st.info("No project history found. Save a project to see it in the history.")
            return
        
        # Display project history
        for i, project in enumerate(reversed(session_state["project_history"])):
            with st.expander(f"{project['name']} - {project['timestamp']}"):
                st.write(f"**Description:** {project['description']}")
                
                # Add load button
                if st.button(f"Load Project", key=f"load_history_{i}"):
                    # Load project data
                    project_data = DataPersistence.load_project_from_json(project["data"])
                    
                    if project_data:
                        # Update session state with project data
                        for key, value in project_data.items():
                            session_state[key] = value
                        
                        st.success("Project loaded successfully!")
                        st.experimental_rerun()
                
                # Add download button
                filename = f"{project['name'].replace(' ', '_')}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.hvac"
                download_link = DataPersistence.get_download_link(project["data"], filename, "Download Project File")
                st.markdown(download_link, unsafe_allow_html=True)
                
                # Add delete button
                if st.button(f"Delete from History", key=f"delete_history_{i}"):
                    # Remove project from history
                    session_state["project_history"].remove(project)
                    
                    st.success("Project removed from history.")
                    st.experimental_rerun()


# Create a singleton instance
data_persistence = DataPersistence()

# Example usage
if __name__ == "__main__":
    import streamlit as st
    
    # Initialize session state with dummy data for testing
    if "building_info" not in st.session_state:
        st.session_state["building_info"] = {
            "project_name": "Test Project",
            "building_name": "Test Building",
            "location": "New York",
            "climate_zone": "4A",
            "building_type": "Office",
            "floor_area": 1000.0,
            "num_floors": 2,
            "floor_height": 3.0,
            "orientation": "NORTH",
            "occupancy": 50,
            "operating_hours": "8:00-18:00",
            "design_conditions": {
                "summer_outdoor_db": 35.0,
                "summer_outdoor_wb": 25.0,
                "summer_indoor_db": 24.0,
                "summer_indoor_rh": 50.0,
                "winter_outdoor_db": -5.0,
                "winter_outdoor_rh": 80.0,
                "winter_indoor_db": 21.0,
                "winter_indoor_rh": 40.0
            }
        }
    
    # Display project management interface
    data_persistence.display_project_management(st.session_state)