"""
Time-based visualization module for HVAC Load Calculator.
This module provides visualization tools for time-based load analysis.
"""

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 calendar
from datetime import datetime, timedelta


class TimeBasedVisualization:
    """Class for time-based visualization."""
    
    @staticmethod
    def create_hourly_load_profile(hourly_loads: Dict[str, List[float]], 
                                  date: str = "Jul 15") -> go.Figure:
        """
        Create an hourly load profile chart.
        
        Args:
            hourly_loads: Dictionary with hourly load data
            date: Date for the profile (e.g., "Jul 15")
            
        Returns:
            Plotly figure with hourly load profile
        """
        # Create hour labels
        hours = list(range(24))
        hour_labels = [f"{h}:00" for h in hours]
        
        # Create figure
        fig = go.Figure()
        
        # Add total load trace
        if "total" in hourly_loads:
            fig.add_trace(go.Scatter(
                x=hour_labels,
                y=hourly_loads["total"],
                mode="lines+markers",
                name="Total Load",
                line=dict(color="rgba(55, 83, 109, 1)", width=3),
                marker=dict(size=8)
            ))
        
        # Add component load traces
        for component, loads in hourly_loads.items():
            if component == "total":
                continue
            
            # Format component name for display
            display_name = component.replace("_", " ").title()
            
            fig.add_trace(go.Scatter(
                x=hour_labels,
                y=loads,
                mode="lines+markers",
                name=display_name,
                marker=dict(size=6),
                line=dict(width=2)
            ))
        
        # Update layout
        fig.update_layout(
            title=f"Hourly Load Profile ({date})",
            xaxis_title="Hour of Day",
            yaxis_title="Load (W)",
            height=500,
            legend=dict(
                orientation="h",
                yanchor="bottom",
                y=1.02,
                xanchor="right",
                x=1
            ),
            hovermode="x unified"
        )
        
        return fig
    
    @staticmethod
    def create_daily_load_profile(daily_loads: Dict[str, List[float]], 
                                 month: str = "July") -> go.Figure:
        """
        Create a daily load profile chart for a month.
        
        Args:
            daily_loads: Dictionary with daily load data
            month: Month name
            
        Returns:
            Plotly figure with daily load profile
        """
        # Get number of days in month
        month_num = list(calendar.month_name).index(month)
        year = datetime.now().year
        num_days = calendar.monthrange(year, month_num)[1]
        
        # Create day labels
        days = list(range(1, num_days + 1))
        day_labels = [f"{d}" for d in days]
        
        # Create figure
        fig = go.Figure()
        
        # Add total load trace
        if "total" in daily_loads:
            fig.add_trace(go.Scatter(
                x=day_labels,
                y=daily_loads["total"][:num_days],
                mode="lines+markers",
                name="Total Load",
                line=dict(color="rgba(55, 83, 109, 1)", width=3),
                marker=dict(size=8)
            ))
        
        # Add component load traces
        for component, loads in daily_loads.items():
            if component == "total":
                continue
            
            # Format component name for display
            display_name = component.replace("_", " ").title()
            
            fig.add_trace(go.Scatter(
                x=day_labels,
                y=loads[:num_days],
                mode="lines+markers",
                name=display_name,
                marker=dict(size=6),
                line=dict(width=2)
            ))
        
        # Update layout
        fig.update_layout(
            title=f"Daily Load Profile ({month})",
            xaxis_title="Day of Month",
            yaxis_title="Load (W)",
            height=500,
            legend=dict(
                orientation="h",
                yanchor="bottom",
                y=1.02,
                xanchor="right",
                x=1
            ),
            hovermode="x unified"
        )
        
        return fig
    
    @staticmethod
    def create_monthly_load_comparison(monthly_loads: Dict[str, List[float]],
                                      load_type: str = "cooling") -> go.Figure:
        """
        Create a monthly load comparison chart.
        
        Args:
            monthly_loads: Dictionary with monthly load data
            load_type: Type of load ("cooling" or "heating")
            
        Returns:
            Plotly figure with monthly load comparison
        """
        # Create month labels
        months = list(calendar.month_name)[1:]
        
        # Create figure
        fig = go.Figure()
        
        # Add total load bars
        if "total" in monthly_loads:
            fig.add_trace(go.Bar(
                x=months,
                y=monthly_loads["total"],
                name="Total Load",
                marker_color="rgba(55, 83, 109, 0.7)",
                opacity=0.7
            ))
        
        # Add component load bars
        for component, loads in monthly_loads.items():
            if component == "total":
                continue
            
            # Format component name for display
            display_name = component.replace("_", " ").title()
            
            fig.add_trace(go.Bar(
                x=months,
                y=loads,
                name=display_name,
                visible="legendonly"
            ))
        
        # Update layout
        title = f"Monthly {load_type.title()} Load Comparison"
        y_title = f"{load_type.title()} Load (kWh)"
        
        fig.update_layout(
            title=title,
            xaxis_title="Month",
            yaxis_title=y_title,
            height=500,
            legend=dict(
                orientation="h",
                yanchor="bottom",
                y=1.02,
                xanchor="right",
                x=1
            ),
            hovermode="x unified"
        )
        
        return fig
    
    @staticmethod
    def create_annual_load_distribution(annual_loads: Dict[str, float],
                                       load_type: str = "cooling") -> go.Figure:
        """
        Create an annual load distribution pie chart.
        
        Args:
            annual_loads: Dictionary with annual load data by component
            load_type: Type of load ("cooling" or "heating")
            
        Returns:
            Plotly figure with annual load distribution
        """
        # Extract components and values
        components = []
        values = []
        
        for component, load in annual_loads.items():
            if component == "total":
                continue
            
            # Format component name for display
            display_name = component.replace("_", " ").title()
            components.append(display_name)
            values.append(load)
        
        # Create pie chart
        fig = go.Figure(data=[go.Pie(
            labels=components,
            values=values,
            hole=0.3,
            textinfo="label+percent",
            insidetextorientation="radial"
        )])
        
        # Update layout
        title = f"Annual {load_type.title()} Load Distribution"
        
        fig.update_layout(
            title=title,
            height=500,
            legend=dict(
                orientation="h",
                yanchor="bottom",
                y=1.02,
                xanchor="right",
                x=1
            )
        )
        
        return fig
    
    @staticmethod
    def create_peak_load_analysis(peak_loads: Dict[str, Dict[str, Any]],
                                 load_type: str = "cooling") -> go.Figure:
        """
        Create a peak load analysis chart.
        
        Args:
            peak_loads: Dictionary with peak load data
            load_type: Type of load ("cooling" or "heating")
            
        Returns:
            Plotly figure with peak load analysis
        """
        # Extract peak load data
        components = []
        values = []
        times = []
        
        for component, data in peak_loads.items():
            if component == "total":
                continue
            
            # Format component name for display
            display_name = component.replace("_", " ").title()
            components.append(display_name)
            values.append(data["value"])
            times.append(data["time"])
        
        # Create bar chart
        fig = go.Figure(data=[go.Bar(
            x=components,
            y=values,
            text=times,
            textposition="auto",
            hovertemplate="<b>%{x}</b><br>Peak Load: %{y:.0f} W<br>Time: %{text}<extra></extra>"
        )])
        
        # Update layout
        title = f"Peak {load_type.title()} Load Analysis"
        y_title = f"Peak {load_type.title()} Load (W)"
        
        fig.update_layout(
            title=title,
            xaxis_title="Component",
            yaxis_title=y_title,
            height=500
        )
        
        return fig
    
    @staticmethod
    def create_load_duration_curve(hourly_loads: List[float],
                                  load_type: str = "cooling") -> go.Figure:
        """
        Create a load duration curve.
        
        Args:
            hourly_loads: List of hourly loads for the year
            load_type: Type of load ("cooling" or "heating")
            
        Returns:
            Plotly figure with load duration curve
        """
        # Sort loads in descending order
        sorted_loads = sorted(hourly_loads, reverse=True)
        
        # Create hour indices
        hours = list(range(1, len(sorted_loads) + 1))
        
        # Create figure
        fig = go.Figure(data=[go.Scatter(
            x=hours,
            y=sorted_loads,
            mode="lines",
            line=dict(color="rgba(55, 83, 109, 1)", width=2),
            fill="tozeroy",
            fillcolor="rgba(55, 83, 109, 0.2)"
        )])
        
        # Update layout
        title = f"{load_type.title()} Load Duration Curve"
        x_title = "Hours"
        y_title = f"{load_type.title()} Load (W)"
        
        fig.update_layout(
            title=title,
            xaxis_title=x_title,
            yaxis_title=y_title,
            height=500,
            xaxis=dict(
                type="log",
                range=[0, math.log10(len(hours))]
            )
        )
        
        return fig
    
    @staticmethod
    def create_heat_map(hourly_data: List[List[float]],
                       x_labels: List[str],
                       y_labels: List[str],
                       title: str,
                       colorscale: str = "Viridis") -> go.Figure:
        """
        Create a heat map visualization.
        
        Args:
            hourly_data: 2D list of hourly data
            x_labels: Labels for x-axis
            y_labels: Labels for y-axis
            title: Chart title
            colorscale: Colorscale for the heatmap
            
        Returns:
            Plotly figure with heat map
        """
        # Create figure
        fig = go.Figure(data=go.Heatmap(
            z=hourly_data,
            x=x_labels,
            y=y_labels,
            colorscale=colorscale,
            colorbar=dict(title="Load (W)")
        ))
        
        # Update layout
        fig.update_layout(
            title=title,
            height=600,
            xaxis=dict(
                title="Hour of Day",
                tickmode="array",
                tickvals=list(range(0, 24, 2)),
                ticktext=[f"{h}:00" for h in range(0, 24, 2)]
            ),
            yaxis=dict(
                title="Day",
                autorange="reversed"
            )
        )
        
        return fig
    
    @staticmethod
    def display_time_based_visualization(cooling_loads: Dict[str, Any] = None,
                                        heating_loads: Dict[str, Any] = None) -> None:
        """
        Display time-based visualization in Streamlit.
        
        Args:
            cooling_loads: Dictionary with cooling load data
            heating_loads: Dictionary with heating load data
        """
        st.header("Time-Based Visualization")
        
        # Check if load data exists
        if cooling_loads is None and heating_loads is None:
            st.warning("No load data available for visualization.")
            
            # Create sample data for demonstration
            st.info("Using sample data for demonstration.")
            
            # Generate sample cooling loads
            cooling_loads = {
                "hourly": {
                    "total": [1000 + 500 * math.sin(h * math.pi / 12) + 1000 * math.sin(h * math.pi / 6) for h in range(24)],
                    "walls": [300 + 150 * math.sin(h * math.pi / 12) for h in range(24)],
                    "roofs": [400 + 200 * math.sin(h * math.pi / 12) for h in range(24)],
                    "windows": [500 + 300 * math.sin(h * math.pi / 6) for h in range(24)],
                    "internal": [200 + 100 * math.sin(h * math.pi / 8) for h in range(24)]
                },
                "daily": {
                    "total": [2000 + 1000 * math.sin(d * math.pi / 15) for d in range(1, 32)],
                    "walls": [600 + 300 * math.sin(d * math.pi / 15) for d in range(1, 32)],
                    "roofs": [800 + 400 * math.sin(d * math.pi / 15) for d in range(1, 32)],
                    "windows": [1000 + 500 * math.sin(d * math.pi / 15) for d in range(1, 32)]
                },
                "monthly": {
                    "total": [1000, 1200, 1500, 2000, 2500, 3000, 3500, 3200, 2800, 2000, 1500, 1200],
                    "walls": [300, 350, 400, 500, 600, 700, 800, 750, 650, 500, 400, 350],
                    "roofs": [400, 450, 500, 600, 700, 800, 900, 850, 750, 600, 500, 450],
                    "windows": [500, 550, 600, 700, 800, 900, 1000, 950, 850, 700, 600, 550]
                },
                "annual": {
                    "total": 25000,
                    "walls": 6000,
                    "roofs": 8000,
                    "windows": 9000,
                    "internal": 2000
                },
                "peak": {
                    "total": {"value": 3500, "time": "Jul 15, 15:00"},
                    "walls": {"value": 800, "time": "Jul 15, 16:00"},
                    "roofs": {"value": 900, "time": "Jul 15, 14:00"},
                    "windows": {"value": 1000, "time": "Jul 15, 15:00"},
                    "internal": {"value": 200, "time": "Jul 15, 17:00"}
                }
            }
            
            # Generate sample heating loads
            heating_loads = {
                "hourly": {
                    "total": [3000 - 1000 * math.sin(h * math.pi / 12) for h in range(24)],
                    "walls": [900 - 300 * math.sin(h * math.pi / 12) for h in range(24)],
                    "roofs": [1200 - 400 * math.sin(h * math.pi / 12) for h in range(24)],
                    "windows": [1500 - 500 * math.sin(h * math.pi / 12) for h in range(24)]
                },
                "daily": {
                    "total": [3000 - 1000 * math.sin(d * math.pi / 15) for d in range(1, 32)],
                    "walls": [900 - 300 * math.sin(d * math.pi / 15) for d in range(1, 32)],
                    "roofs": [1200 - 400 * math.sin(d * math.pi / 15) for d in range(1, 32)],
                    "windows": [1500 - 500 * math.sin(d * math.pi / 15) for d in range(1, 32)]
                },
                "monthly": {
                    "total": [3500, 3200, 2800, 2000, 1500, 1000, 800, 1000, 1500, 2000, 2800, 3500],
                    "walls": [1050, 960, 840, 600, 450, 300, 240, 300, 450, 600, 840, 1050],
                    "roofs": [1400, 1280, 1120, 800, 600, 400, 320, 400, 600, 800, 1120, 1400],
                    "windows": [1750, 1600, 1400, 1000, 750, 500, 400, 500, 750, 1000, 1400, 1750]
                },
                "annual": {
                    "total": 25000,
                    "walls": 7500,
                    "roofs": 10000,
                    "windows": 12500,
                    "infiltration": 5000
                },
                "peak": {
                    "total": {"value": 3500, "time": "Jan 15, 06:00"},
                    "walls": {"value": 1050, "time": "Jan 15, 06:00"},
                    "roofs": {"value": 1400, "time": "Jan 15, 06:00"},
                    "windows": {"value": 1750, "time": "Jan 15, 06:00"},
                    "infiltration": {"value": 500, "time": "Jan 15, 06:00"}
                }
            }
        
        # Create tabs for different visualizations
        tab1, tab2, tab3, tab4, tab5 = st.tabs([
            "Hourly Profiles", 
            "Monthly Comparison", 
            "Annual Distribution", 
            "Peak Load Analysis",
            "Heat Maps"
        ])
        
        with tab1:
            st.subheader("Hourly Load Profiles")
            
            # Add load type selector
            load_type = st.radio(
                "Select Load Type",
                ["cooling", "heating"],
                horizontal=True,
                key="hourly_profile_type"
            )
            
            # Add date selector
            date = st.selectbox(
                "Select Date",
                ["Jan 15", "Apr 15", "Jul 15", "Oct 15"],
                index=2,
                key="hourly_profile_date"
            )
            
            # Get appropriate load data
            if load_type == "cooling":
                hourly_data = cooling_loads.get("hourly", {})
            else:
                hourly_data = heating_loads.get("hourly", {})
            
            # Create and display chart
            fig = TimeBasedVisualization.create_hourly_load_profile(hourly_data, date)
            st.plotly_chart(fig, use_container_width=True)
            
            # Add daily profile option
            st.subheader("Daily Load Profiles")
            
            # Add month selector
            month = st.selectbox(
                "Select Month",
                list(calendar.month_name)[1:],
                index=6,  # July
                key="daily_profile_month"
            )
            
            # Get appropriate load data
            if load_type == "cooling":
                daily_data = cooling_loads.get("daily", {})
            else:
                daily_data = heating_loads.get("daily", {})
            
            # Create and display chart
            fig = TimeBasedVisualization.create_daily_load_profile(daily_data, month)
            st.plotly_chart(fig, use_container_width=True)
        
        with tab2:
            st.subheader("Monthly Load Comparison")
            
            # Add load type selector
            load_type = st.radio(
                "Select Load Type",
                ["cooling", "heating"],
                horizontal=True,
                key="monthly_comparison_type"
            )
            
            # Get appropriate load data
            if load_type == "cooling":
                monthly_data = cooling_loads.get("monthly", {})
            else:
                monthly_data = heating_loads.get("monthly", {})
            
            # Create and display chart
            fig = TimeBasedVisualization.create_monthly_load_comparison(monthly_data, load_type)
            st.plotly_chart(fig, use_container_width=True)
            
            # Add download button for CSV
            monthly_df = pd.DataFrame(monthly_data)
            monthly_df.index = list(calendar.month_name)[1:]
            
            csv = monthly_df.to_csv().encode('utf-8')
            st.download_button(
                label=f"Download Monthly {load_type.title()} Loads as CSV",
                data=csv,
                file_name=f"monthly_{load_type}_loads.csv",
                mime="text/csv"
            )
        
        with tab3:
            st.subheader("Annual Load Distribution")
            
            # Add load type selector
            load_type = st.radio(
                "Select Load Type",
                ["cooling", "heating"],
                horizontal=True,
                key="annual_distribution_type"
            )
            
            # Get appropriate load data
            if load_type == "cooling":
                annual_data = cooling_loads.get("annual", {})
            else:
                annual_data = heating_loads.get("annual", {})
            
            # Create and display chart
            fig = TimeBasedVisualization.create_annual_load_distribution(annual_data, load_type)
            st.plotly_chart(fig, use_container_width=True)
            
            # Display annual total
            total = annual_data.get("total", 0)
            st.metric(f"Total Annual {load_type.title()} Load", f"{total:,.0f} kWh")
            
            # Add download button for CSV
            annual_df = pd.DataFrame({"Component": list(annual_data.keys()), "Load (kWh)": list(annual_data.values())})
            
            csv = annual_df.to_csv(index=False).encode('utf-8')
            st.download_button(
                label=f"Download Annual {load_type.title()} Loads as CSV",
                data=csv,
                file_name=f"annual_{load_type}_loads.csv",
                mime="text/csv"
            )
        
        with tab4:
            st.subheader("Peak Load Analysis")
            
            # Add load type selector
            load_type = st.radio(
                "Select Load Type",
                ["cooling", "heating"],
                horizontal=True,
                key="peak_load_type"
            )
            
            # Get appropriate load data
            if load_type == "cooling":
                peak_data = cooling_loads.get("peak", {})
            else:
                peak_data = heating_loads.get("peak", {})
            
            # Create and display chart
            fig = TimeBasedVisualization.create_peak_load_analysis(peak_data, load_type)
            st.plotly_chart(fig, use_container_width=True)
            
            # Display peak total
            peak_total = peak_data.get("total", {}).get("value", 0)
            peak_time = peak_data.get("total", {}).get("time", "")
            
            st.metric(f"Peak {load_type.title()} Load", f"{peak_total:,.0f} W")
            st.write(f"Peak Time: {peak_time}")
            
            # Add download button for CSV
            peak_df = pd.DataFrame({
                "Component": list(peak_data.keys()),
                "Peak Load (W)": [data.get("value", 0) for data in peak_data.values()],
                "Time": [data.get("time", "") for data in peak_data.values()]
            })
            
            csv = peak_df.to_csv(index=False).encode('utf-8')
            st.download_button(
                label=f"Download Peak {load_type.title()} Loads as CSV",
                data=csv,
                file_name=f"peak_{load_type}_loads.csv",
                mime="text/csv"
            )
        
        with tab5:
            st.subheader("Heat Maps")
            
            # Add load type selector
            load_type = st.radio(
                "Select Load Type",
                ["cooling", "heating"],
                horizontal=True,
                key="heat_map_type"
            )
            
            # Add month selector
            month = st.selectbox(
                "Select Month",
                list(calendar.month_name)[1:],
                index=6,  # July
                key="heat_map_month"
            )
            
            # Generate heat map data
            month_num = list(calendar.month_name).index(month)
            year = datetime.now().year
            num_days = calendar.monthrange(year, month_num)[1]
            
            # Get appropriate hourly data
            if load_type == "cooling":
                hourly_data = cooling_loads.get("hourly", {}).get("total", [])
            else:
                hourly_data = heating_loads.get("hourly", {}).get("total", [])
            
            # Create 2D array for heat map
            heat_map_data = []
            for day in range(1, num_days + 1):
                # Generate hourly data with day-to-day variation
                day_factor = 1 + 0.2 * math.sin(day * math.pi / 15)
                day_data = [load * day_factor for load in hourly_data]
                heat_map_data.append(day_data)
            
            # Create hour and day labels
            hour_labels = list(range(24))
            day_labels = list(range(1, num_days + 1))
            
            # Create and display heat map
            title = f"{load_type.title()} Load Heat Map ({month})"
            colorscale = "Hot" if load_type == "cooling" else "Ice"
            
            fig = TimeBasedVisualization.create_heat_map(heat_map_data, hour_labels, day_labels, title, colorscale)
            st.plotly_chart(fig, use_container_width=True)
            
            # Add explanation
            st.info(
                "The heat map shows the hourly load pattern for each day of the selected month. "
                "Darker colors indicate higher loads. This visualization helps identify peak load periods "
                "and daily/weekly patterns."
            )


# Create a singleton instance
time_based_visualization = TimeBasedVisualization()

# Example usage
if __name__ == "__main__":
    import streamlit as st
    
    # Display time-based visualization with sample data
    time_based_visualization.display_time_based_visualization()