utils/component_visualization.py

"""
Hierarchical component visualization module for HVAC Load Calculator.
This module provides visualization tools for building components.
"""

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 data models
from data.building_components import Wall, Roof, Floor, Window, Door, Orientation, ComponentType


class ComponentVisualization:
    """Class for hierarchical component visualization."""
    
    @staticmethod
    def create_component_summary_table(components: Dict[str, List[Any]]) -> pd.DataFrame:
        """
        Create a summary table of building components.
        
        Args:
            components: Dictionary with lists of building components
            
        Returns:
            DataFrame with component summary
        """
        # Initialize data
        data = []
        
        # Process walls
        for wall in components.get("walls", []):
            data.append({
                "Component Type": "Wall",
                "Name": wall.name,
                "Orientation": wall.orientation.name,
                "Area (m²)": wall.area,
                "U-Value (W/m²·K)": wall.u_value,
                "Heat Transfer (W/K)": wall.area * wall.u_value
            })
        
        # Process roofs
        for roof in components.get("roofs", []):
            data.append({
                "Component Type": "Roof",
                "Name": roof.name,
                "Orientation": roof.orientation.name,
                "Area (m²)": roof.area,
                "U-Value (W/m²·K)": roof.u_value,
                "Heat Transfer (W/K)": roof.area * roof.u_value
            })
        
        # Process floors
        for floor in components.get("floors", []):
            data.append({
                "Component Type": "Floor",
                "Name": floor.name,
                "Orientation": "Horizontal",
                "Area (m²)": floor.area,
                "U-Value (W/m²·K)": floor.u_value,
                "Heat Transfer (W/K)": floor.area * floor.u_value
            })
        
        # Process windows
        for window in components.get("windows", []):
            data.append({
                "Component Type": "Window",
                "Name": window.name,
                "Orientation": window.orientation.name,
                "Area (m²)": window.area,
                "U-Value (W/m²·K)": window.u_value,
                "Heat Transfer (W/K)": window.area * window.u_value,
                "SHGC": window.shgc if hasattr(window, "shgc") else None
            })
        
        # Process doors
        for door in components.get("doors", []):
            data.append({
                "Component Type": "Door",
                "Name": door.name,
                "Orientation": door.orientation.name,
                "Area (m²)": door.area,
                "U-Value (W/m²·K)": door.u_value,
                "Heat Transfer (W/K)": door.area * door.u_value
            })
        
        # Create DataFrame
        df = pd.DataFrame(data)
        
        return df
    
    @staticmethod
    def create_component_area_chart(components: Dict[str, List[Any]]) -> go.Figure:
        """
        Create a pie chart of component areas.
        
        Args:
            components: Dictionary with lists of building components
            
        Returns:
            Plotly figure with component area breakdown
        """
        # Calculate total areas by component type
        areas = {
            "Walls": sum(wall.area for wall in components.get("walls", [])),
            "Roofs": sum(roof.area for roof in components.get("roofs", [])),
            "Floors": sum(floor.area for floor in components.get("floors", [])),
            "Windows": sum(window.area for window in components.get("windows", [])),
            "Doors": sum(door.area for door in components.get("doors", []))
        }
        
        # Create labels and values
        labels = list(areas.keys())
        values = list(areas.values())
        
        # Create pie chart
        fig = go.Figure(data=[go.Pie(
            labels=labels,
            values=values,
            hole=0.3,
            textinfo="label+percent",
            insidetextorientation="radial"
        )])
        
        # Update layout
        fig.update_layout(
            title="Building Component Areas",
            height=500,
            legend=dict(
                orientation="h",
                yanchor="bottom",
                y=1.02,
                xanchor="right",
                x=1
            )
        )
        
        return fig
    
    @staticmethod
    def create_orientation_area_chart(components: Dict[str, List[Any]]) -> go.Figure:
        """
        Create a bar chart of areas by orientation.
        
        Args:
            components: Dictionary with lists of building components
            
        Returns:
            Plotly figure with area breakdown by orientation
        """
        # Initialize areas by orientation
        orientation_areas = {
            "NORTH": 0,
            "NORTHEAST": 0,
            "EAST": 0,
            "SOUTHEAST": 0,
            "SOUTH": 0,
            "SOUTHWEST": 0,
            "WEST": 0,
            "NORTHWEST": 0,
            "HORIZONTAL": 0
        }
        
        # Calculate areas by orientation for walls
        for wall in components.get("walls", []):
            orientation_areas[wall.orientation.name] += wall.area
        
        # Calculate areas by orientation for windows
        for window in components.get("windows", []):
            orientation_areas[window.orientation.name] += window.area
        
        # Calculate areas by orientation for doors
        for door in components.get("doors", []):
            orientation_areas[door.orientation.name] += door.area
        
        # Add roofs and floors to horizontal
        for roof in components.get("roofs", []):
            if roof.orientation.name == "HORIZONTAL":
                orientation_areas["HORIZONTAL"] += roof.area
            else:
                orientation_areas[roof.orientation.name] += roof.area
        
        for floor in components.get("floors", []):
            orientation_areas["HORIZONTAL"] += floor.area
        
        # Create labels and values
        orientations = []
        areas = []
        
        for orientation, area in orientation_areas.items():
            if area > 0:
                orientations.append(orientation)
                areas.append(area)
        
        # Create bar chart
        fig = go.Figure(data=[go.Bar(
            x=orientations,
            y=areas,
            text=areas,
            texttemplate="%{y:.1f} m²",
            textposition="auto"
        )])
        
        # Update layout
        fig.update_layout(
            title="Building Component Areas by Orientation",
            xaxis_title="Orientation",
            yaxis_title="Area (m²)",
            height=500
        )
        
        return fig
    
    @staticmethod
    def create_heat_transfer_chart(components: Dict[str, List[Any]]) -> go.Figure:
        """
        Create a bar chart of heat transfer coefficients by component type.
        
        Args:
            components: Dictionary with lists of building components
            
        Returns:
            Plotly figure with heat transfer breakdown
        """
        # Calculate heat transfer by component type
        heat_transfer = {
            "Walls": sum(wall.area * wall.u_value for wall in components.get("walls", [])),
            "Roofs": sum(roof.area * roof.u_value for roof in components.get("roofs", [])),
            "Floors": sum(floor.area * floor.u_value for floor in components.get("floors", [])),
            "Windows": sum(window.area * window.u_value for window in components.get("windows", [])),
            "Doors": sum(door.area * door.u_value for door in components.get("doors", []))
        }
        
        # Create labels and values
        labels = list(heat_transfer.keys())
        values = list(heat_transfer.values())
        
        # Create bar chart
        fig = go.Figure(data=[go.Bar(
            x=labels,
            y=values,
            text=values,
            texttemplate="%{y:.1f} W/K",
            textposition="auto"
        )])
        
        # Update layout
        fig.update_layout(
            title="Heat Transfer Coefficients by Component Type",
            xaxis_title="Component Type",
            yaxis_title="Heat Transfer Coefficient (W/K)",
            height=500
        )
        
        return fig
        
    def display_component_breakdown(self, calculation_results: Dict[str, Any]) -> None:
        """
        Display component breakdown of calculation results.
        
        Args:
            calculation_results: Dictionary containing calculation results
        """
        # Ensure calculation_results has the required structure
        if not calculation_results or "components" not in calculation_results:
            st.warning("No component data available for visualization.")
            return
            
        # Get components from calculation results
        components = calculation_results.get("components", {})
        
        # Create tabs for different visualizations
        tab1, tab2, tab3, tab4 = st.tabs([
            "Component Summary", 
            "Area Breakdown", 
            "Orientation Analysis",
            "Heat Transfer Analysis"
        ])
        
        with tab1:
            # Display component summary table
            st.subheader("Building Component Summary")
            df = self.create_component_summary_table(components)
            if not df.empty:
                st.dataframe(df, use_container_width=True)
            else:
                st.info("No components added yet.")
        
        with tab2:
            # Display component area chart
            st.subheader("Component Area Breakdown")
            fig = self.create_component_area_chart(components)
            st.plotly_chart(fig, use_container_width=True)
            
            # Display total area
            total_area = sum(
                sum(wall.area for wall in components.get("walls", [])) +
                sum(roof.area for roof in components.get("roofs", [])) +
                sum(floor.area for floor in components.get("floors", [])) +
                sum(window.area for window in components.get("windows", [])) +
                sum(door.area for door in components.get("doors", []))
            )
            st.metric("Total Building Envelope Area", f"{total_area:.2f} m²")
        
        with tab3:
            # Display orientation area chart
            st.subheader("Component Orientation Analysis")
            fig = self.create_orientation_area_chart(components)
            st.plotly_chart(fig, use_container_width=True)
        
        with tab4:
            # Display heat transfer chart
            st.subheader("Heat Transfer Analysis")
            fig = self.create_heat_transfer_chart(components)
            st.plotly_chart(fig, use_container_width=True)
            
            # Display total heat transfer coefficient
            total_heat_transfer = sum(
                sum(wall.area * wall.u_value for wall in components.get("walls", [])) +
                sum(roof.area * roof.u_value for roof in components.get("roofs", [])) +
                sum(floor.area * floor.u_value for floor in components.get("floors", [])) +
                sum(window.area * window.u_value for window in components.get("windows", [])) +
                sum(door.area * door.u_value for door in components.get("doors", []))
            )
            st.metric("Total Heat Transfer Coefficient", f"{total_heat_transfer:.2f} W/K")
        
        return fig
    
    @staticmethod
    def create_3d_building_model(components: Dict[str, List[Any]]) -> go.Figure:
        """
        Create a 3D visualization of the building components.
        
        Args:
            components: Dictionary with lists of building components
            
        Returns:
            Plotly figure with 3D building model
        """
        # Initialize figure
        fig = go.Figure()
        
        # Define colors
        colors = {
            "Wall": "lightblue",
            "Roof": "red",
            "Floor": "brown",
            "Window": "skyblue",
            "Door": "orange"
        }
        
        # Define orientation vectors
        orientation_vectors = {
            "NORTH": (0, 1, 0),
            "NORTHEAST": (0.7071, 0.7071, 0),
            "EAST": (1, 0, 0),
            "SOUTHEAST": (0.7071, -0.7071, 0),
            "SOUTH": (0, -1, 0),
            "SOUTHWEST": (-0.7071, -0.7071, 0),
            "WEST": (-1, 0, 0),
            "NORTHWEST": (-0.7071, 0.7071, 0),
            "HORIZONTAL": (0, 0, 1)
        }
        
        # Define building dimensions (simplified model)
        building_width = 10
        building_depth = 10
        building_height = 3
        
        # Create walls
        for i, wall in enumerate(components.get("walls", [])):
            orientation = wall.orientation.name
            vector = orientation_vectors[orientation]
            
            # Determine wall position and dimensions
            if orientation in ["NORTH", "SOUTH"]:
                width = building_width
                height = building_height
                depth = 0.3
                
                if orientation == "NORTH":
                    x = 0
                    y = building_depth / 2
                else:  # SOUTH
                    x = 0
                    y = -building_depth / 2
                
                z = building_height / 2
                
            elif orientation in ["EAST", "WEST"]:
                width = 0.3
                height = building_height
                depth = building_depth
                
                if orientation == "EAST":
                    x = building_width / 2
                    y = 0
                else:  # WEST
                    x = -building_width / 2
                    y = 0
                
                z = building_height / 2
            
            else:  # Diagonal orientations
                width = building_width / 2
                height = building_height
                depth = 0.3
                
                if orientation == "NORTHEAST":
                    x = building_width / 4
                    y = building_depth / 4
                elif orientation == "SOUTHEAST":
                    x = building_width / 4
                    y = -building_depth / 4
                elif orientation == "SOUTHWEST":
                    x = -building_width / 4
                    y = -building_depth / 4
                else:  # NORTHWEST
                    x = -building_width / 4
                    y = building_depth / 4
                
                z = building_height / 2
            
            # Add wall to figure
            fig.add_trace(go.Mesh3d(
                x=[x - width/2, x + width/2, x + width/2, x - width/2, x - width/2, x + width/2, x + width/2, x - width/2],
                y=[y - depth/2, y - depth/2, y + depth/2, y + depth/2, y - depth/2, y - depth/2, y + depth/2, y + depth/2],
                z=[z - height/2, z - height/2, z - height/2, z - height/2, z + height/2, z + height/2, z + height/2, z + height/2],
                i=[0, 0, 0, 1, 4, 4],
                j=[1, 2, 4, 2, 5, 6],
                k=[2, 3, 7, 3, 6, 7],
                color=colors["Wall"],
                opacity=0.7,
                name=f"Wall: {wall.name}"
            ))
        
        # Create roof
        for i, roof in enumerate(components.get("roofs", [])):
            # Add roof to figure
            fig.add_trace(go.Mesh3d(
                x=[-building_width/2, building_width/2, building_width/2, -building_width/2],
                y=[-building_depth/2, -building_depth/2, building_depth/2, building_depth/2],
                z=[building_height, building_height, building_height, building_height],
                i=[0],
                j=[1],
                k=[2],
                color=colors["Roof"],
                opacity=0.7,
                name=f"Roof: {roof.name}"
            ))
            
            fig.add_trace(go.Mesh3d(
                x=[-building_width/2, -building_width/2, building_width/2],
                y=[building_depth/2, -building_depth/2, -building_depth/2],
                z=[building_height, building_height, building_height],
                i=[0],
                j=[1],
                k=[2],
                color=colors["Roof"],
                opacity=0.7,
                name=f"Roof: {roof.name}"
            ))
        
        # Create floor
        for i, floor in enumerate(components.get("floors", [])):
            # Add floor to figure
            fig.add_trace(go.Mesh3d(
                x=[-building_width/2, building_width/2, building_width/2, -building_width/2],
                y=[-building_depth/2, -building_depth/2, building_depth/2, building_depth/2],
                z=[0, 0, 0, 0],
                i=[0],
                j=[1],
                k=[2],
                color=colors["Floor"],
                opacity=0.7,
                name=f"Floor: {floor.name}"
            ))
            
            fig.add_trace(go.Mesh3d(
                x=[-building_width/2, -building_width/2, building_width/2],
                y=[building_depth/2, -building_depth/2, -building_depth/2],
                z=[0, 0, 0],
                i=[0],
                j=[1],
                k=[2],
                color=colors["Floor"],
                opacity=0.7,
                name=f"Floor: {floor.name}"
            ))
        
        # Create windows
        for i, window in enumerate(components.get("windows", [])):
            orientation = window.orientation.name
            vector = orientation_vectors[orientation]
            
            # Determine window position and dimensions
            window_width = 1.5
            window_height = 1.2
            window_depth = 0.1
            
            if orientation == "NORTH":
                x = i * 3 - building_width/4
                y = building_depth / 2
                z = building_height / 2
            elif orientation == "SOUTH":
                x = i * 3 - building_width/4
                y = -building_depth / 2
                z = building_height / 2
            elif orientation == "EAST":
                x = building_width / 2
                y = i * 3 - building_depth/4
                z = building_height / 2
            elif orientation == "WEST":
                x = -building_width / 2
                y = i * 3 - building_depth/4
                z = building_height / 2
            else:
                # Skip diagonal orientations for simplicity
                continue
            
            # Add window to figure
            fig.add_trace(go.Mesh3d(
                x=[x - window_width/2, x + window_width/2, x + window_width/2, x - window_width/2, x - window_width/2, x + window_width/2, x + window_width/2, x - window_width/2],
                y=[y - window_depth/2, y - window_depth/2, y + window_depth/2, y + window_depth/2, y - window_depth/2, y - window_depth/2, y + window_depth/2, y + window_depth/2],
                z=[z - window_height/2, z - window_height/2, z - window_height/2, z - window_height/2, z + window_height/2, z + window_height/2, z + window_height/2, z + window_height/2],
                i=[0, 0, 0, 1, 4, 4],
                j=[1, 2, 4, 2, 5, 6],
                k=[2, 3, 7, 3, 6, 7],
                color=colors["Window"],
                opacity=0.5,
                name=f"Window: {window.name}"
            ))
        
        # Create doors
        for i, door in enumerate(components.get("doors", [])):
            orientation = door.orientation.name
            vector = orientation_vectors[orientation]
            
            # Determine door position and dimensions
            door_width = 1.0
            door_height = 2.0
            door_depth = 0.1
            
            if orientation == "NORTH":
                x = i * 3
                y = building_depth / 2
                z = door_height / 2
            elif orientation == "SOUTH":
                x = i * 3
                y = -building_depth / 2
                z = door_height / 2
            elif orientation == "EAST":
                x = building_width / 2
                y = i * 3
                z = door_height / 2
            elif orientation == "WEST":
                x = -building_width / 2
                y = i * 3
                z = door_height / 2
            else:
                # Skip diagonal orientations for simplicity
                continue
            
            # Add door to figure
            fig.add_trace(go.Mesh3d(
                x=[x - door_width/2, x + door_width/2, x + door_width/2, x - door_width/2, x - door_width/2, x + door_width/2, x + door_width/2, x - door_width/2],
                y=[y - door_depth/2, y - door_depth/2, y + door_depth/2, y + door_depth/2, y - door_depth/2, y - door_depth/2, y + door_depth/2, y + door_depth/2],
                z=[z - door_height/2, z - door_height/2, z - door_height/2, z - door_height/2, z + door_height/2, z + door_height/2, z + door_height/2, z + door_height/2],
                i=[0, 0, 0, 1, 4, 4],
                j=[1, 2, 4, 2, 5, 6],
                k=[2, 3, 7, 3, 6, 7],
                color=colors["Door"],
                opacity=0.7,
                name=f"Door: {door.name}"
            ))
        
        # Update layout
        fig.update_layout(
            title="3D Building Model",
            scene=dict(
                xaxis_title="X",
                yaxis_title="Y",
                zaxis_title="Z",
                aspectmode="data"
            ),
            height=700,
            margin=dict(l=0, r=0, b=0, t=30)
        )
        
        return fig
    
    @staticmethod
    def display_component_visualization(components: Dict[str, List[Any]]) -> None:
        """
        Display component visualization in Streamlit.
        
        Args:
            components: Dictionary with lists of building components
        """
        st.header("Building Component Visualization")
        
        # Create tabs for different visualizations
        tab1, tab2, tab3, tab4, tab5 = st.tabs([
            "Component Summary", 
            "Area Breakdown", 
            "Orientation Analysis", 
            "Heat Transfer Analysis",
            "3D Building Model"
        ])
        
        with tab1:
            st.subheader("Component Summary")
            df = ComponentVisualization.create_component_summary_table(components)
            st.dataframe(df, use_container_width=True)
            
            # Add download button for CSV
            csv = df.to_csv(index=False).encode('utf-8')
            st.download_button(
                label="Download Component Summary as CSV",
                data=csv,
                file_name="component_summary.csv",
                mime="text/csv"
            )
        
        with tab2:
            st.subheader("Area Breakdown")
            fig = ComponentVisualization.create_component_area_chart(components)
            st.plotly_chart(fig, use_container_width=True)
        
        with tab3:
            st.subheader("Orientation Analysis")
            fig = ComponentVisualization.create_orientation_area_chart(components)
            st.plotly_chart(fig, use_container_width=True)
        
        with tab4:
            st.subheader("Heat Transfer Analysis")
            fig = ComponentVisualization.create_heat_transfer_chart(components)
            st.plotly_chart(fig, use_container_width=True)
        
        with tab5:
            st.subheader("3D Building Model")
            fig = ComponentVisualization.create_3d_building_model(components)
            st.plotly_chart(fig, use_container_width=True)


# Create a singleton instance
component_visualization = ComponentVisualization()

# Example usage
if __name__ == "__main__":
    import streamlit as st
    from data.building_components import Wall, Roof, Floor, Window, Door, Orientation, ComponentType
    
    # Create sample building components
    walls = [
        Wall(
            id="wall1",
            name="North Wall",
            component_type=ComponentType.WALL,
            u_value=0.5,
            area=20.0,
            orientation=Orientation.NORTH,
            wall_type="Brick",
            wall_group="B"
        ),
        Wall(
            id="wall2",
            name="South Wall",
            component_type=ComponentType.WALL,
            u_value=0.5,
            area=20.0,
            orientation=Orientation.SOUTH,
            wall_type="Brick",
            wall_group="B"
        ),
        Wall(
            id="wall3",
            name="East Wall",
            component_type=ComponentType.WALL,
            u_value=0.5,
            area=15.0,
            orientation=Orientation.EAST,
            wall_type="Brick",
            wall_group="B"
        ),
        Wall(
            id="wall4",
            name="West Wall",
            component_type=ComponentType.WALL,
            u_value=0.5,
            area=15.0,
            orientation=Orientation.WEST,
            wall_type="Brick",
            wall_group="B"
        )
    ]
    
    roofs = [
        Roof(
            id="roof1",
            name="Flat Roof",
            component_type=ComponentType.ROOF,
            u_value=0.3,
            area=100.0,
            orientation=Orientation.HORIZONTAL,
            roof_type="Concrete",
            roof_group="C"
        )
    ]
    
    floors = [
        Floor(
            id="floor1",
            name="Ground Floor",
            component_type=ComponentType.FLOOR,
            u_value=0.4,
            area=100.0,
            floor_type="Concrete"
        )
    ]
    
    windows = [
        Window(
            id="window1",
            name="North Window 1",
            component_type=ComponentType.WINDOW,
            u_value=2.8,
            area=4.0,
            orientation=Orientation.NORTH,
            shgc=0.7,
            vt=0.8,
            window_type="Double Glazed",
            glazing_layers=2,
            gas_fill="Air",
            low_e_coating=False
        ),
        Window(
            id="window2",
            name="South Window 1",
            component_type=ComponentType.WINDOW,
            u_value=2.8,
            area=6.0,
            orientation=Orientation.SOUTH,
            shgc=0.7,
            vt=0.8,
            window_type="Double Glazed",
            glazing_layers=2,
            gas_fill="Air",
            low_e_coating=False
        ),
        Window(
            id="window3",
            name="East Window 1",
            component_type=ComponentType.WINDOW,
            u_value=2.8,
            area=3.0,
            orientation=Orientation.EAST,
            shgc=0.7,
            vt=0.8,
            window_type="Double Glazed",
            glazing_layers=2,
            gas_fill="Air",
            low_e_coating=False
        )
    ]
    
    doors = [
        Door(
            id="door1",
            name="Front Door",
            component_type=ComponentType.DOOR,
            u_value=2.0,
            area=2.0,
            orientation=Orientation.SOUTH,
            door_type="Solid Wood"
        )
    ]
    
    # Create components dictionary
    components = {
        "walls": walls,
        "roofs": roofs,
        "floors": floors,
        "windows": windows,
        "doors": doors
    }
    
    # Display component visualization
    component_visualization.display_component_visualization(components)