"""
U-Value calculator module for HVAC Load Calculator.
This module implements the layer-by-layer assembly builder and U-value calculation functions.
"""

from typing import Dict, List, Any, Optional, Tuple
import pandas as pd
import numpy as np
import os
import json
from dataclasses import dataclass, field

# Import data models
from data.building_components import MaterialLayer
from data.reference_data import reference_data

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


@dataclass
class MaterialAssembly:
    """Class representing a material assembly for U-value calculation."""
    
    name: str
    description: str = ""
    layers: List[MaterialLayer] = field(default_factory=list)
    
    # Surface resistances (m²·K/W)
    r_si: float = 0.13  # Interior surface resistance
    r_se: float = 0.04  # Exterior surface resistance
    
    def add_layer(self, layer: MaterialLayer) -> None:
        """
        Add a material layer to the assembly.
        
        Args:
            layer: MaterialLayer object
        """
        self.layers.append(layer)
    
    def remove_layer(self, index: int) -> bool:
        """
        Remove a material layer from the assembly.
        
        Args:
            index: Index of the layer to remove
            
        Returns:
            True if the layer was removed, False otherwise
        """
        if index < 0 or index >= len(self.layers):
            return False
        
        self.layers.pop(index)
        return True
    
    def move_layer(self, from_index: int, to_index: int) -> bool:
        """
        Move a material layer within the assembly.
        
        Args:
            from_index: Current index of the layer
            to_index: New index for the layer
            
        Returns:
            True if the layer was moved, False otherwise
        """
        if (from_index < 0 or from_index >= len(self.layers) or
            to_index < 0 or to_index >= len(self.layers)):
            return False
        
        layer = self.layers.pop(from_index)
        self.layers.insert(to_index, layer)
        return True
    
    @property
    def total_thickness(self) -> float:
        """Calculate the total thickness of the assembly in meters."""
        return sum(layer.thickness for layer in self.layers)
    
    @property
    def r_value_layers(self) -> float:
        """Calculate the total thermal resistance of all layers in m²·K/W."""
        return sum(layer.r_value for layer in self.layers)
    
    @property
    def r_value_total(self) -> float:
        """Calculate the total thermal resistance including surface resistances in m²·K/W."""
        return self.r_si + self.r_value_layers + self.r_se
    
    @property
    def u_value(self) -> float:
        """Calculate the U-value of the assembly in W/(m²·K)."""
        if self.r_value_total == 0:
            return float('inf')
        return 1 / self.r_value_total
    
    @property
    def thermal_mass(self) -> Optional[float]:
        """Calculate the total thermal mass of the assembly in J/(m²·K)."""
        masses = [layer.thermal_mass for layer in self.layers]
        if None in masses:
            return None
        return sum(masses)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert the material assembly to a dictionary."""
        return {
            "name": self.name,
            "description": self.description,
            "layers": [layer.to_dict() for layer in self.layers],
            "r_si": self.r_si,
            "r_se": self.r_se,
            "total_thickness": self.total_thickness,
            "r_value_layers": self.r_value_layers,
            "r_value_total": self.r_value_total,
            "u_value": self.u_value,
            "thermal_mass": self.thermal_mass
        }


class UValueCalculator:
    """Class for calculating U-values of material assemblies."""
    
    def __init__(self):
        """Initialize U-value calculator."""
        self.assemblies = {}
        self.load_preset_assemblies()
    
    def load_preset_assemblies(self) -> None:
        """Load preset material assemblies."""
        # Create preset assemblies from reference data
        
        # Wall assemblies
        for wall_id, wall_data in reference_data.wall_types.items():
            # Create material layers
            layers = []
            for layer_data in wall_data.get("layers", []):
                material_id = layer_data.get("material")
                thickness = layer_data.get("thickness")
                
                material = reference_data.get_material(material_id)
                if material:
                    layer = MaterialLayer(
                        name=material["name"],
                        thickness=thickness,
                        conductivity=material["conductivity"],
                        density=material.get("density"),
                        specific_heat=material.get("specific_heat")
                    )
                    layers.append(layer)
            
            # Create assembly
            assembly_id = f"preset_wall_{wall_id}"
            assembly = MaterialAssembly(
                name=wall_data["name"],
                description=wall_data["description"],
                layers=layers
            )
            
            self.assemblies[assembly_id] = assembly
        
        # Roof assemblies
        for roof_id, roof_data in reference_data.roof_types.items():
            # Create material layers
            layers = []
            for layer_data in roof_data.get("layers", []):
                material_id = layer_data.get("material")
                thickness = layer_data.get("thickness")
                
                material = reference_data.get_material(material_id)
                if material:
                    layer = MaterialLayer(
                        name=material["name"],
                        thickness=thickness,
                        conductivity=material["conductivity"],
                        density=material.get("density"),
                        specific_heat=material.get("specific_heat")
                    )
                    layers.append(layer)
            
            # Create assembly
            assembly_id = f"preset_roof_{roof_id}"
            assembly = MaterialAssembly(
                name=roof_data["name"],
                description=roof_data["description"],
                layers=layers
            )
            
            self.assemblies[assembly_id] = assembly
        
        # Floor assemblies
        for floor_id, floor_data in reference_data.floor_types.items():
            # Create material layers
            layers = []
            for layer_data in floor_data.get("layers", []):
                material_id = layer_data.get("material")
                thickness = layer_data.get("thickness")
                
                material = reference_data.get_material(material_id)
                if material:
                    layer = MaterialLayer(
                        name=material["name"],
                        thickness=thickness,
                        conductivity=material["conductivity"],
                        density=material.get("density"),
                        specific_heat=material.get("specific_heat")
                    )
                    layers.append(layer)
            
            # Create assembly
            assembly_id = f"preset_floor_{floor_id}"
            assembly = MaterialAssembly(
                name=floor_data["name"],
                description=floor_data["description"],
                layers=layers
            )
            
            self.assemblies[assembly_id] = assembly
    
    def get_assembly(self, assembly_id: str) -> Optional[MaterialAssembly]:
        """
        Get a material assembly by ID.
        
        Args:
            assembly_id: Assembly identifier
            
        Returns:
            MaterialAssembly object or None if not found
        """
        return self.assemblies.get(assembly_id)
    
    def get_preset_assemblies(self) -> Dict[str, MaterialAssembly]:
        """
        Get all preset material assemblies.
        
        Returns:
            Dictionary of preset MaterialAssembly objects
        """
        return {assembly_id: assembly for assembly_id, assembly in self.assemblies.items()
                if assembly_id.startswith("preset_")}
    
    def get_custom_assemblies(self) -> Dict[str, MaterialAssembly]:
        """
        Get all custom material assemblies.
        
        Returns:
            Dictionary of custom MaterialAssembly objects
        """
        return {assembly_id: assembly for assembly_id, assembly in self.assemblies.items()
                if assembly_id.startswith("custom_")}
    
    def create_assembly(self, name: str, description: str = "") -> str:
        """
        Create a new material assembly.
        
        Args:
            name: Assembly name
            description: Assembly description
            
        Returns:
            Assembly ID
        """
        import uuid
        
        assembly_id = f"custom_assembly_{str(uuid.uuid4())[:8]}"
        assembly = MaterialAssembly(name=name, description=description)
        
        self.assemblies[assembly_id] = assembly
        return assembly_id
    
    def add_layer_to_assembly(self, assembly_id: str, material_id: str, thickness: float) -> bool:
        """
        Add a material layer to an assembly.
        
        Args:
            assembly_id: Assembly identifier
            material_id: Material identifier
            thickness: Layer thickness in meters
            
        Returns:
            True if the layer was added, False otherwise
        """
        if assembly_id not in self.assemblies:
            return False
        
        material = reference_data.get_material(material_id)
        if not material:
            return False
        
        layer = MaterialLayer(
            name=material["name"],
            thickness=thickness,
            conductivity=material["conductivity"],
            density=material.get("density"),
            specific_heat=material.get("specific_heat")
        )
        
        self.assemblies[assembly_id].add_layer(layer)
        return True
    
    def add_custom_layer_to_assembly(self, assembly_id: str, name: str, thickness: float,
                                    conductivity: float, density: float = None,
                                    specific_heat: float = None) -> bool:
        """
        Add a custom material layer to an assembly.
        
        Args:
            assembly_id: Assembly identifier
            name: Layer name
            thickness: Layer thickness in meters
            conductivity: Thermal conductivity in W/(m·K)
            density: Density in kg/m³ (optional)
            specific_heat: Specific heat capacity in J/(kg·K) (optional)
            
        Returns:
            True if the layer was added, False otherwise
        """
        if assembly_id not in self.assemblies:
            return False
        
        layer = MaterialLayer(
            name=name,
            thickness=thickness,
            conductivity=conductivity,
            density=density,
            specific_heat=specific_heat
        )
        
        self.assemblies[assembly_id].add_layer(layer)
        return True
    
    def remove_layer_from_assembly(self, assembly_id: str, layer_index: int) -> bool:
        """
        Remove a material layer from an assembly.
        
        Args:
            assembly_id: Assembly identifier
            layer_index: Index of the layer to remove
            
        Returns:
            True if the layer was removed, False otherwise
        """
        if assembly_id not in self.assemblies:
            return False
        
        return self.assemblies[assembly_id].remove_layer(layer_index)
    
    def move_layer_in_assembly(self, assembly_id: str, from_index: int, to_index: int) -> bool:
        """
        Move a material layer within an assembly.
        
        Args:
            assembly_id: Assembly identifier
            from_index: Current index of the layer
            to_index: New index for the layer
            
        Returns:
            True if the layer was moved, False otherwise
        """
        if assembly_id not in self.assemblies:
            return False
        
        return self.assemblies[assembly_id].move_layer(from_index, to_index)
    
    def calculate_u_value(self, assembly_id: str) -> Optional[float]:
        """
        Calculate the U-value of an assembly.
        
        Args:
            assembly_id: Assembly identifier
            
        Returns:
            U-value in W/(m²·K) or None if the assembly was not found
        """
        if assembly_id not in self.assemblies:
            return None
        
        return self.assemblies[assembly_id].u_value
    
    def calculate_r_value(self, assembly_id: str) -> Optional[float]:
        """
        Calculate the R-value of an assembly.
        
        Args:
            assembly_id: Assembly identifier
            
        Returns:
            R-value in m²·K/W or None if the assembly was not found
        """
        if assembly_id not in self.assemblies:
            return None
        
        return self.assemblies[assembly_id].r_value_total
    
    def export_to_json(self, file_path: str) -> None:
        """
        Export all assemblies to a JSON file.
        
        Args:
            file_path: Path to the output JSON file
        """
        data = {assembly_id: assembly.to_dict() for assembly_id, assembly in self.assemblies.items()}
        
        with open(file_path, 'w') as f:
            json.dump(data, f, indent=4)
    
    def import_from_json(self, file_path: str) -> int:
        """
        Import assemblies from a JSON file.
        
        Args:
            file_path: Path to the input JSON file
            
        Returns:
            Number of assemblies imported
        """
        with open(file_path, 'r') as f:
            data = json.load(f)
        
        count = 0
        for assembly_id, assembly_data in data.items():
            try:
                # Create assembly
                assembly = MaterialAssembly(
                    name=assembly_data["name"],
                    description=assembly_data.get("description", ""),
                    r_si=assembly_data.get("r_si", 0.13),
                    r_se=assembly_data.get("r_se", 0.04)
                )
                
                # Add layers
                for layer_data in assembly_data.get("layers", []):
                    layer = MaterialLayer(
                        name=layer_data["name"],
                        thickness=layer_data["thickness"],
                        conductivity=layer_data["conductivity"],
                        density=layer_data.get("density"),
                        specific_heat=layer_data.get("specific_heat")
                    )
                    assembly.add_layer(layer)
                
                self.assemblies[assembly_id] = assembly
                count += 1
            except Exception as e:
                print(f"Error importing assembly {assembly_id}: {e}")
        
        return count


# Create a singleton instance
u_value_calculator = UValueCalculator()

# Export U-value calculator to JSON if needed
if __name__ == "__main__":
    u_value_calculator.export_to_json(os.path.join(DATA_DIR, "data", "u_value_calculator.json"))