"""
Visualization utilities for HVAC Load Calculator

This module provides enhanced visualization functions for creating interactive
and informative charts for the HVAC Load Calculator application.
"""

import plotly.express as px
import plotly.graph_objects as go
import pandas as pd
import numpy as np


def create_load_breakdown_chart(load_components, title="Load Breakdown"):
    """
    Create an enhanced pie chart for load components breakdown.
    
    Args:
        load_components (dict): Dictionary of load components and their values
        title (str): Chart title
        
    Returns:
        plotly.graph_objects.Figure: Interactive pie chart
    """
    # Remove zero values
    load_components = {k: v for k, v in load_components.items() if v > 0}
    
    # Create figure
    fig = go.Figure()
    
    # Add pie chart
    fig.add_trace(go.Pie(
        labels=list(load_components.keys()),
        values=list(load_components.values()),
        textinfo='label+percent',
        insidetextorientation='radial',
        marker=dict(
            colors=px.colors.qualitative.Bold,
            line=dict(color='white', width=2)
        ),
        pull=[0.05 if x == max(load_components.values()) else 0 for x in load_components.values()],
        hovertemplate='<b>%{label}</b><br>%{value:.1f} W<br>%{percent}<extra></extra>'
    ))
    
    # Update layout
    fig.update_layout(
        title={
            'text': title,
            'y': 0.95,
            'x': 0.5,
            'xanchor': 'center',
            'yanchor': 'top',
            'font': dict(size=20)
        },
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=-0.2,
            xanchor="center",
            x=0.5,
            font=dict(size=12)
        ),
        height=500,
        margin=dict(t=80, b=80, l=40, r=40),
        paper_bgcolor='rgba(0,0,0,0)',
        plot_bgcolor='rgba(0,0,0,0)'
    )
    
    return fig


def create_component_bar_chart(df, x_col, y_col, color_col=None, title="Component Breakdown"):
    """
    Create an enhanced bar chart for component breakdown.
    
    Args:
        df (pd.DataFrame): DataFrame containing the data
        x_col (str): Column name for x-axis
        y_col (str): Column name for y-axis
        color_col (str, optional): Column name for color grouping
        title (str): Chart title
        
    Returns:
        plotly.graph_objects.Figure: Interactive bar chart
    """
    # Create figure
    if color_col:
        fig = px.bar(
            df,
            x=x_col,
            y=y_col,
            color=color_col,
            title=title,
            color_discrete_sequence=px.colors.qualitative.Bold,
            height=500,
            text=y_col
        )
    else:
        fig = px.bar(
            df,
            x=x_col,
            y=y_col,
            title=title,
            color_discrete_sequence=px.colors.qualitative.Bold,
            height=500,
            text=y_col
        )
    
    # Update layout
    fig.update_layout(
        xaxis_title=x_col,
        yaxis_title=y_col,
        legend_title=color_col if color_col else "",
        font=dict(size=12),
        xaxis={'categoryorder': 'total descending'},
        paper_bgcolor='rgba(0,0,0,0)',
        plot_bgcolor='rgba(0,0,0,0)',
        hovermode="x unified"
    )
    
    # Add data labels
    fig.update_traces(
        texttemplate='%{y:.1f}',
        textposition='outside',
        hovertemplate='<b>%{x}</b><br>%{y:.1f}<extra></extra>'
    )
    
    # Add grid lines
    fig.update_yaxes(
        showgrid=True,
        gridwidth=1,
        gridcolor='rgba(211,211,211,0.3)'
    )
    
    return fig


def create_stacked_bar_chart(df, x_col, y_cols, names, title="Stacked Bar Chart"):
    """
    Create an enhanced stacked bar chart.
    
    Args:
        df (pd.DataFrame): DataFrame containing the data
        x_col (str): Column name for x-axis
        y_cols (list): List of column names for y-axis values
        names (list): List of names for each y-column
        title (str): Chart title
        
    Returns:
        plotly.graph_objects.Figure: Interactive stacked bar chart
    """
    # Create figure
    fig = go.Figure()
    
    # Add bars for each y column
    for i, y_col in enumerate(y_cols):
        fig.add_trace(go.Bar(
            x=df[x_col],
            y=df[y_col],
            name=names[i],
            hovertemplate=f'<b>{names[i]}</b>: %{{y:.1f}}<extra></extra>'
        ))
    
    # Update layout
    fig.update_layout(
        title=title,
        xaxis_title=x_col,
        yaxis_title="Value",
        barmode='stack',
        height=500,
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.02,
            xanchor="center",
            x=0.5
        ),
        paper_bgcolor='rgba(0,0,0,0)',
        plot_bgcolor='rgba(0,0,0,0)',
        hovermode="x unified"
    )
    
    # Add grid lines
    fig.update_yaxes(
        showgrid=True,
        gridwidth=1,
        gridcolor='rgba(211,211,211,0.3)'
    )
    
    return fig


def create_grouped_bar_chart(df, x_col, y_cols, names, title="Grouped Bar Chart"):
    """
    Create an enhanced grouped bar chart.
    
    Args:
        df (pd.DataFrame): DataFrame containing the data
        x_col (str): Column name for x-axis
        y_cols (list): List of column names for y-axis values
        names (list): List of names for each y-column
        title (str): Chart title
        
    Returns:
        plotly.graph_objects.Figure: Interactive grouped bar chart
    """
    # Create figure
    fig = go.Figure()
    
    # Add bars for each y column
    for i, y_col in enumerate(y_cols):
        fig.add_trace(go.Bar(
            x=df[x_col],
            y=df[y_col],
            name=names[i],
            hovertemplate=f'<b>{names[i]}</b>: %{{y:.1f}}<extra></extra>'
        ))
    
    # Update layout
    fig.update_layout(
        title=title,
        xaxis_title=x_col,
        yaxis_title="Value",
        barmode='group',
        height=500,
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.02,
            xanchor="center",
            x=0.5
        ),
        paper_bgcolor='rgba(0,0,0,0)',
        plot_bgcolor='rgba(0,0,0,0)',
        hovermode="x unified"
    )
    
    # Add grid lines
    fig.update_yaxes(
        showgrid=True,
        gridwidth=1,
        gridcolor='rgba(211,211,211,0.3)'
    )
    
    return fig


def create_heat_map_chart(df, x_col, y_col, z_col, title="Heat Map"):
    """
    Create an enhanced heat map chart.
    
    Args:
        df (pd.DataFrame): DataFrame containing the data
        x_col (str): Column name for x-axis
        y_col (str): Column name for y-axis
        z_col (str): Column name for z-axis (color)
        title (str): Chart title
        
    Returns:
        plotly.graph_objects.Figure: Interactive heat map chart
    """
    # Create figure
    fig = px.density_heatmap(
        df,
        x=x_col,
        y=y_col,
        z=z_col,
        title=title,
        color_continuous_scale="Viridis",
        height=500
    )
    
    # Update layout
    fig.update_layout(
        xaxis_title=x_col,
        yaxis_title=y_col,
        font=dict(size=12),
        paper_bgcolor='rgba(0,0,0,0)',
        plot_bgcolor='rgba(0,0,0,0)'
    )
    
    return fig


def create_line_chart(df, x_col, y_cols, names, title="Line Chart"):
    """
    Create an enhanced line chart.
    
    Args:
        df (pd.DataFrame): DataFrame containing the data
        x_col (str): Column name for x-axis
        y_cols (list): List of column names for y-axis values
        names (list): List of names for each y-column
        title (str): Chart title
        
    Returns:
        plotly.graph_objects.Figure: Interactive line chart
    """
    # Create figure
    fig = go.Figure()
    
    # Add lines for each y column
    for i, y_col in enumerate(y_cols):
        fig.add_trace(go.Scatter(
            x=df[x_col],
            y=df[y_col],
            mode='lines+markers',
            name=names[i],
            hovertemplate=f'<b>{names[i]}</b>: %{{y:.1f}}<extra></extra>'
        ))
    
    # Update layout
    fig.update_layout(
        title=title,
        xaxis_title=x_col,
        yaxis_title="Value",
        height=500,
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.02,
            xanchor="center",
            x=0.5
        ),
        paper_bgcolor='rgba(0,0,0,0)',
        plot_bgcolor='rgba(0,0,0,0)',
        hovermode="x unified"
    )
    
    # Add grid lines
    fig.update_yaxes(
        showgrid=True,
        gridwidth=1,
        gridcolor='rgba(211,211,211,0.3)'
    )
    
    fig.update_xaxes(
        showgrid=True,
        gridwidth=1,
        gridcolor='rgba(211,211,211,0.3)'
    )
    
    return fig


def create_sankey_diagram(nodes, links, title="Energy Flow"):
    """
    Create a Sankey diagram for energy flow visualization.
    
    Args:
        nodes (list): List of node labels
        links (dict): Dictionary with source, target, and value lists
        title (str): Chart title
        
    Returns:
        plotly.graph_objects.Figure: Interactive Sankey diagram
    """
    # Create figure
    fig = go.Figure(data=[go.Sankey(
        node=dict(
            pad=15,
            thickness=20,
            line=dict(color="black", width=0.5),
            label=nodes,
            color="blue"
        ),
        link=dict(
            source=links['source'],
            target=links['target'],
            value=links['value'],
            hovertemplate='%{source.label} → %{target.label}: %{value:.1f} W<extra></extra>'
        )
    )])
    
    # Update layout
    fig.update_layout(
        title=title,
        font=dict(size=12),
        height=600,
        paper_bgcolor='rgba(0,0,0,0)',
        plot_bgcolor='rgba(0,0,0,0)'
    )
    
    return fig


def create_gauge_chart(value, min_val, max_val, title="Gauge", threshold_values=None, threshold_colors=None):
    """
    Create a gauge chart for displaying a value within a range.
    
    Args:
        value (float): Value to display
        min_val (float): Minimum value of the range
        max_val (float): Maximum value of the range
        title (str): Chart title
        threshold_values (list, optional): List of threshold values
        threshold_colors (list, optional): List of colors for each threshold
        
    Returns:
        plotly.graph_objects.Figure: Interactive gauge chart
    """
    # Set default thresholds if not provided
    if threshold_values is None:
        threshold_values = [min_val, (min_val + max_val) / 2, max_val]
    
    if threshold_colors is None:
        threshold_colors = ["green", "yellow", "red"]
    
    # Create figure
    fig = go.Figure(go.Indicator(
        mode="gauge+number",
        value=value,
        domain={'x': [0, 1], 'y': [0, 1]},
        title={'text': title},
        gauge={
            'axis': {'range': [min_val, max_val]},
            'bar': {'color': "darkblue"},
            'steps': [
                {'range': [threshold_values[i], threshold_values[i+1]], 'color': threshold_colors[i]} 
                for i in range(len(threshold_values)-1)
            ],
            'threshold': {
                'line': {'color': "red", 'width': 4},
                'thickness': 0.75,
                'value': value
            }
        }
    ))
    
    # Update layout
    fig.update_layout(
        height=300,
        paper_bgcolor='rgba(0,0,0,0)',
        plot_bgcolor='rgba(0,0,0,0)'
    )
    
    return fig


def create_enhanced_results_visualization(results):
    """
    Create enhanced visualizations for HVAC load calculation results.
    
    Args:
        results (dict): Dictionary containing calculation results
        
    Returns:
        dict: Dictionary of plotly figures
    """
    figures = {}
    
    # Prepare data for load breakdown pie chart
    load_components = {
        'Walls': results.get('wall_loss', 0),
        'Roof': results.get('roof_loss', 0),
        'Floor': results.get('floor_loss', 0),
        'Windows & Doors': results.get('window_loss', 0),
        'Infiltration': results.get('infiltration_loss', 0),
        'Ventilation': results.get('ventilation_loss', 0) - results.get('infiltration_loss', 0)
    }
    
    # Create load breakdown pie chart
    figures['load_breakdown'] = create_load_breakdown_chart(
        load_components, 
        title="Heating Load Components"
    )
    
    # Create energy flow Sankey diagram
    if 'internal_gain' in results and results['internal_gain'] > 0:
        # Create nodes and links for Sankey diagram
        nodes = [
            "Walls", "Roof", "Floor", "Windows & Doors", 
            "Infiltration", "Ventilation", "Internal Gains", 
            "Total Heat Loss", "Net Heating Load"
        ]
        
        links = {
            'source': [0, 1, 2, 3, 4, 5, 7, 6],
            'target': [7, 7, 7, 7, 7, 7, 8, 8],
            'value': [
                load_components['Walls'],
                load_components['Roof'],
                load_components['Floor'],
                load_components['Windows & Doors'],
                load_components['Infiltration'],
                load_components['Ventilation'],
                results['internal_gain'],
                results['total_heat_loss']
            ]
        }
        
        figures['energy_flow'] = create_sankey_diagram(
            nodes, 
            links, 
            title="Heating Energy Flow"
        )
    
    # Create gauge chart for heating load per area
    if 'net_heating_load' in results and 'building_info' in results:
        floor_area = results['building_info'].get('floor_area', 80.0)
        heating_load_per_area = results['net_heating_load'] / floor_area
        
        figures['load_per_area_gauge'] = create_gauge_chart(
            heating_load_per_area,
            0,
            200,
            title="Heating Load per Area (W/m²)",
            threshold_values=[0, 50, 100, 150, 200],
            threshold_colors=["green", "lightgreen", "yellow", "orange", "red"]
        )
    
    return figures