import plotly.graph_objects as go
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import seaborn as sns
import networkx as nx

def plot_forecast(result):
    """Interactive backtest plot with zoom and pan functionality using Plotly"""
    forecast = result["forecast"]
    actual = result["actual"]

    # Convert to numpy arrays and flatten if needed
    forecast = np.array(forecast).flatten()
    actual = np.array(actual).flatten()
    
    # Ensure both arrays have the same length
    min_len = min(len(forecast), len(actual))
    forecast = forecast[:min_len]
    actual = actual[:min_len]

    # Create time indices
    time_indices = np.arange(len(actual))
    
    # Initialize Plotly figure
    fig = go.Figure()

    if len(actual) == 0 or len(forecast) == 0:
        fig.add_annotation(
            x=0.5, y=0.5, xref="paper", yref="paper",
            text="No data available for plotting",
            showarrow=False, font=dict(size=12)
        )
        return fig

    # Plot full historical actual
    fig.add_trace(go.Scatter(
        x=time_indices, y=actual,
        mode='lines', name="Historical Actual",
        line=dict(color="blue", width=2), opacity=0.7
    ))
    
    # Plot full historical forecast
    fig.add_trace(go.Scatter(
        x=time_indices, y=forecast,
        mode='lines', name="Historical Forecast",
        line=dict(color="orange", width=2, dash="dash"), opacity=0.7
    ))
    
    if len(actual) > 1 and len(forecast) > 1:
        last_idx = len(actual) - 1
        
        # Highlight last day actual segment
        last_actual_segment = [float(actual[last_idx-1]), float(actual[last_idx])]
        last_time_segment = [time_indices[last_idx-1], time_indices[last_idx]]
        fig.add_trace(go.Scatter(
            x=last_time_segment, y=last_actual_segment,
            mode='lines', name="Last Day Actual",
            line=dict(color="blue", width=4), showlegend=False
        ))
        
        # Add markers for last day comparison
        fig.add_trace(go.Scatter(
            x=[last_idx], y=[float(actual[last_idx])],
            mode='markers', name="Last Day Actual",
            marker=dict(color="blue", size=10, line=dict(color="darkblue", width=2)),
            showlegend=False
        ))
        fig.add_trace(go.Scatter(
            x=[last_idx], y=[float(forecast[last_idx])],
            mode='markers', name="Last Day Predicted",
            marker=dict(color="red", size=10, line=dict(color="darkred", width=2)),
            showlegend=False
        ))
        
        # Add value annotations for last day
        actual_val = float(actual[last_idx])
        forecast_val = float(forecast[last_idx])
        fig.add_annotation(
            x=last_idx, y=actual_val,
            text=f"Actual: {actual_val:.2f}",
            showarrow=True, arrowhead=1, ax=20, ay=-30,
            font=dict(size=10, color="white"),
            bgcolor="blue", opacity=0.8, bordercolor="darkblue"
        )
        fig.add_annotation(
            x=last_idx, y=forecast_val,
            text=f"Predicted: {forecast_val:.2f}",
            showarrow=True, arrowhead=1, ax=20, ay=30,
            font=dict(size=10, color="white"),
            bgcolor="red", opacity=0.8, bordercolor="darkred"
        )
    elif len(actual) == 1:
        # Handle single point case
        fig.add_trace(go.Scatter(
            x=[0], y=[float(actual[0])],
            mode='markers', name="Actual",
            marker=dict(color="blue", size=10), showlegend=False
        ))
        fig.add_trace(go.Scatter(
            x=[0], y=[float(forecast[0])],
            mode='markers', name="Predicted",
            marker=dict(color="red", size=10), showlegend=False
        ))

    # Configure layout
    fig.update_layout(
        xaxis_title="Time Index",
        yaxis_title="Value",
        showlegend=True,
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.1,
            xanchor="center",
            x=0.5
        ),
        hovermode="x unified",
        plot_bgcolor="white",
        grid=dict(rows=1, columns=1),
        xaxis=dict(showgrid=True, gridcolor="rgba(0,0,0,0.1)", gridwidth=0.8),
        yaxis=dict(showgrid=True, gridcolor="rgba(0,0,0,0.1)", gridwidth=0.8),
        margin=dict(t=50)  # Reduced top margin to accommodate legend
    )

    return fig


def plot_future_forecast(df, result, future_df):
    """Interactive future forecast plot with zoom, pan and hover functionality using Plotly"""
    # Initialize Plotly figure
    fig = go.Figure()

    # Validate and convert data
    if df.empty or 'Date' not in df.columns or 'value' not in df.columns:
        fig.add_annotation(
            x=0.5, y=0.5, xref="paper", yref="paper",
            text="No valid historical data available",
            showarrow=False, font=dict(size=12)
        )
        return fig
    
    # Plot historical data
    dates = pd.to_datetime(df['Date'])
    values = np.array(df['value']).flatten()
    fig.add_trace(go.Scatter(
        x=dates, y=values,
        mode='lines', name="Historical Data",
        line=dict(color="blue", width=2.5), opacity=0.9
    ))
    
    if "latest_prediction" in result and len(result["latest_prediction"]) > 0:
        # Convert predictions to flat array
        predictions = np.array(result["latest_prediction"]).flatten()
        
        # Create future dates
        last_date = dates.iloc[-1] if len(dates) > 0 else pd.Timestamp.now()
        horizon = len(predictions)
        
        try:
            future_dates = pd.date_range(start=last_date + pd.Timedelta(days=1), periods=horizon, freq='B')
        except:
            # Fallback to daily frequency if business day fails
            future_dates = pd.date_range(start=last_date + pd.Timedelta(days=1), periods=horizon, freq='D')
        
        if len(values) > 0 and len(predictions) > 0:
            # Create connection from last historical point to first prediction
            connection_dates = [last_date, future_dates[0]]
            connection_values = [float(values[-1]), float(predictions[0])]
            fig.add_trace(go.Scatter(
                x=connection_dates, y=connection_values,
                mode='lines', name="Connection",
                line=dict(color="orange", width=2, dash="dot"), opacity=0.7, showlegend=False
            ))
        
        # Plot forecast
        predictions_float = [float(p) for p in predictions]
        fig.add_trace(go.Scatter(
            x=future_dates, y=predictions_float,
            mode='lines+markers', name="Forecast",
            line=dict(color="orange", width=3),
            marker=dict(size=8, color="orange", line=dict(color="darkorange", width=2)),
            opacity=0.9
        ))
        
        # Plot actual future values if available
        if not future_df.empty and "future_actuals" in result and 'Date' in future_df.columns and 'value' in future_df.columns:
            actual_future_dates = pd.to_datetime(future_df['Date'])
            actual_future_values = np.array(future_df['value']).flatten()
            actual_future_values_float = [float(v) for v in actual_future_values]
            
            fig.add_trace(go.Scatter(
                x=actual_future_dates, y=actual_future_values_float,
                mode='lines+markers', name="Actual Future",
                line=dict(color="green", width=3),
                marker=dict(size=8, color="green", line=dict(color="darkgreen", width=2)),
                opacity=0.9
            ))

    # Configure layout
    fig.update_layout(
        xaxis_title="Date",
        yaxis_title="Stock Price",
        showlegend=True,
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.1,
            xanchor="center",
            x=0.5
        ),
        hovermode="x unified",
        plot_bgcolor="white",
        grid=dict(rows=1, columns=1),
        xaxis=dict(showgrid=True, gridcolor="rgba(0,0,0,0.1)", gridwidth=0.8),
        yaxis=dict(showgrid=True, gridcolor="rgba(0,0,0,0.1)", gridwidth=0.8),
        margin=dict(t=50)
    )

    return fig


def plot_metrics_precision(result):
    """Plot precision metrics using Plotly"""
    metrics = {k: v for k, v in result['metrics'].items() if k in ['R² (%)', 'Explained Variance (%)', 'MDA (%)'] and v is not None}
    if not metrics:
        fig = go.Figure()
        fig.add_annotation(
            x=0.5, y=0.5, xref="paper", yref="paper",
            text="No valid precision metrics available",
            showarrow=False, font=dict(size=12)
        )
        return fig
    
    # Create bar plot
    fig = go.Figure()
    fig.add_trace(go.Bar(
        x=list(metrics.keys()),
        y=list(metrics.values()),
        marker_color=sns.color_palette("Blues_d", len(metrics)).as_hex(),
        text=[f"{v:.2f}%" for v in metrics.values()],
        textposition='auto'
    ))
    
    # Configure layout
    max_val = max(metrics.values(), default=100)
    min_val = min(metrics.values(), default=0)
    fig.update_layout(
        yaxis_title="Value (%)",
        showlegend=False,
        plot_bgcolor="white",
        yaxis=dict(range=[min(min_val - 5, -10), max_val + 10], showgrid=True, gridcolor="rgba(0,0,0,0.1)", gridwidth=0.8),
        xaxis=dict(showgrid=False),
        margin=dict(t=50)
    )

    return fig


def plot_metrics_risk(result):
    """Plot risk metrics using Plotly"""
    metrics = {k: v for k, v in result['metrics'].items() if k in ['RMSE', 'MAE', 'MAPE (%)', 'MASE'] and v is not None}
    if not metrics:
        fig = go.Figure()
        fig.add_annotation(
            x=0.5, y=0.5, xref="paper", yref="paper",
            text="No valid risk metrics available",
            showarrow=False, font=dict(size=12)
        )
        return fig
    
    # Create bar plot
    fig = go.Figure()
    fig.add_trace(go.Bar(
        x=list(metrics.keys()),
        y=list(metrics.values()),
        marker_color=sns.color_palette("Reds_d", len(metrics)).as_hex(),
        text=[f"{v:.2f}" for v in metrics.values()],
        textposition='auto'
    ))
    
    # Configure layout
    max_val = max(metrics.values(), default=1)
    fig.update_layout(
        yaxis_title="Value",
        showlegend=False,
        plot_bgcolor="white",
        yaxis=dict(range=[0, max_val + 0.2 * max_val], showgrid=True, gridcolor="rgba(0,0,0,0.1)", gridwidth=0.8),
        xaxis=dict(showgrid=False),
        margin=dict(t=50)
    )

    return fig


def plot_loss_curve(result):
    """Plot loss curve using Plotly"""
    train_losses = result.get('train_loss', [])
    val_losses = result.get('val_loss', [])
    
    fig = go.Figure()
    fig.add_trace(go.Scatter(
        x=list(range(len(train_losses))), y=train_losses,
        mode='lines', name="Train Loss",
        line=dict(color="blue", width=2)
    ))
    if val_losses:
        fig.add_trace(go.Scatter(
            x=list(range(len(val_losses))), y=val_losses,
            mode='lines', name="Validation Loss",
            line=dict(color="orange", width=2)
        ))
    
    # Configure layout
    fig.update_layout(
        xaxis_title="Epoch",
        yaxis_title="Loss (MSE)",
        showlegend=True,
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.1,
            xanchor="center",
            x=0.5
        ),
        plot_bgcolor="white",
        grid=dict(rows=1, columns=1),
        xaxis=dict(showgrid=True, gridcolor="rgba(0,0,0,0.1)", gridwidth=0.8),
        yaxis=dict(showgrid=True, gridcolor="rgba(0,0,0,0.1)", gridwidth=0.8),
        margin=dict(t=50)
    )

    return fig


def plot_model_architecture(result):
    """Plot model architecture using matplotlib (static, as Plotly is less suited for network graphs)"""
    fig = plt.figure(figsize=(10, 6))
    ax = fig.add_subplot(111)
    ax.axis('off')
    G = nx.DiGraph()

    if "architecture" not in result:
        ax.text(0.5, 0.5, "No architecture details available", ha='center', va='center', fontsize=12)
        return fig

    arch = result["architecture"]
    model_name = arch["model_name"]
    num_layers = arch["num_layers"]
    hidden_units = arch["hidden_units"]
    dropout = arch["dropout"]
    batch_size = arch["batch_size"]
    input_size = arch["input_size"]
    output_size = arch["output_size"]

    # Handle model-specific hidden units for visualization
    if model_name == "MLPModel":
        hidden_nodes = min(hidden_units[0], 5)
        units_label = f"{hidden_units[0]},{hidden_units[1]}"
    elif model_name == "CNNModel":
        hidden_nodes = 5
        units_label = f"{hidden_units} filters"
    elif model_name == "TransformerModel":
        hidden_nodes = min(hidden_units, 5)
        units_label = f"{hidden_units}"
    else:
        hidden_nodes = min(hidden_units, 5)
        units_label = f"{hidden_units}"

    # Simplified block diagram for complex models
    if model_name in ["CNNModel", "HybridModel", "CNN_GRU"]:
        G = nx.DiGraph()
        pos = {}
        nodes = []
        y_pos = 0.5
        layer_width = 1.0 / 4

        if model_name == "CNNModel":
            components = [
                ("Input", f"{input_size} units"),
                ("Conv1D", f"{hidden_units} filters"),
                ("MaxPool", ""),
                ("Output", f"{output_size} units")
            ]
        elif model_name == "HybridModel":
            components = [
                ("Input", f"{input_size} units"),
                ("Conv1D", "32 filters"),
                (f"BiLSTM ({num_layers} layers)", f"{hidden_units*2} units"),
                ("Output", f"{output_size} units")
            ]
        elif model_name == "CNN_GRU":
            components = [
                ("Input", f"{input_size} units"),
                ("Conv1D", "32 filters"),
                (f"GRU ({num_layers} layers)", f"{hidden_units} units"),
                ("Output", f"{output_size} units")
            ]

        for i, (comp, label) in enumerate(components):
            G.add_node(comp, layer=comp)
            pos[comp] = (i * layer_width, y_pos)
            nodes.append([comp])
            if i > 0:
                G.add_edge(components[i-1][0], comp)

        nx.draw(G, pos, ax=ax, with_labels=False, node_color='lightblue', edge_color='gray', 
                node_size=2000, node_shape='s', arrowsize=10)

        for node, (x, y) in pos.items():
            label = [comp[1] for comp in components if comp[0] == node][0]
            ax.text(x, y + 0.05, f"{node}\n{label}", ha='center', va='bottom', fontsize=8,
                    bbox=dict(facecolor='white', alpha=0.8, edgecolor='black'))

    else:
        max_nodes_display = 5
        input_nodes = min(input_size, max_nodes_display)
        output_nodes = min(output_size, max_nodes_display)

        nodes = []
        pos = {}
        layer_width = 1.0 / (num_layers + 2)
        y_pos = 0.5

        for i in range(input_nodes):
            node = f"input_{i}"
            G.add_node(node, layer="input")
            pos[node] = (0, y_pos + (i - input_nodes / 2) * 0.1)
        nodes.append([f"input_{i}" for i in range(input_nodes)])

        for layer in range(num_layers):
            layer_nodes = []
            for i in range(hidden_nodes):
                node = f"hidden_{layer}_{i}"
                G.add_node(node, layer=f"hidden_{layer+1}")
                pos[node] = ((layer + 1) * layer_width, y_pos + (i - hidden_nodes / 2) * 0.1)
                layer_nodes.append(node)
            nodes.append(layer_nodes)

        output_layer_nodes = []
        for i in range(output_nodes):
            node = f"output_{i}"
            G.add_node(node, layer="output")
            pos[node] = ((num_layers + 1) * layer_width, y_pos + (i - output_nodes / 2) * 0.1)
            output_layer_nodes.append(node)
        nodes.append(output_layer_nodes)

        for layer in range(len(nodes) - 1):
            for src in nodes[layer]:
                for dst in nodes[layer + 1]:
                    G.add_edge(src, dst)

        nx.draw(G, pos, ax=ax, with_labels=False, node_color='lightblue', edge_color='gray', 
                node_size=500, arrowsize=10)

        for node in G.nodes(data=True):
            layer = node[1]['layer']
            x, y = pos[node[0]]
            if layer.startswith("hidden"):
                label = f"Layer {layer.split('_')[1]}: {units_label} units"
            elif layer == "input":
                label = f"Input: {input_size} units"
            elif layer == "output":
                label = f"Output: {output_size} units"
            ax.text(x, y + 0.05, label, ha='center', va='bottom', fontsize=8)

    # Add model details as annotation
    details = f"Dropout: {dropout:.2f}\nBatch Size: {batch_size}"
    ax.text(0.5, 0.05, details, ha='center', va='bottom', fontsize=10, transform=ax.transAxes,
            bbox=dict(facecolor='white', alpha=0.8, edgecolor='black'))

    plt.tight_layout()
    return fig