# -*- coding: utf-8 -*-
"""plotting.py

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1ILADgRrYqkAEj5jyymO50ZvzDzVdfD6g
"""

# plotting.py
# Functions for generating Plotly figures from processed strategy data.

import plotly.express as px
import plotly.graph_objects as go
import pandas as pd
import numpy as np
import traceback
from utils import create_empty_figure # Import helper

def generate_figures_for_strategy(strategy_results):
    """
    Generates standard Plotly figures for a single strategy's results.

    Args:
        strategy_results: Dictionary containing processed data for one strategy,
                          as returned by process_single_file. Expected keys include:
                          'filename', 'equity_df', 'drawdown_df', 'benchmark_df',
                          'trades_df', 'exposure_series', 'turnover_df'.

    Returns:
        A dictionary containing Plotly figure objects:
        'equity_fig', 'drawdown_fig', 'benchmark_fig', 'pnl_hist_fig',
        'duration_hist_fig', 'exposure_fig', 'turnover_fig'.
        Uses empty figures if data is missing or invalid.
    """
    figures = {
        "equity_fig": create_empty_figure("Equity Curve"),
        "drawdown_fig": create_empty_figure("Drawdown Curve"),
        "benchmark_fig": create_empty_figure("Equity vs Benchmark"),
        "pnl_hist_fig": create_empty_figure("P/L Distribution"),
        "duration_hist_fig": create_empty_figure("Trade Duration Distribution"),
        "exposure_fig": create_empty_figure("Exposure"),
        "turnover_fig": create_empty_figure("Portfolio Turnover")
    }
    filename = strategy_results.get("filename", "Strategy") # Get filename for titles

    try:
        # --- Equity Curve ---
        equity_df = strategy_results.get("equity_df")
        if equity_df is not None and not equity_df.empty and 'Time' in equity_df.columns and 'Equity' in equity_df.columns:
             # Ensure Time is datetime
             equity_df['Time'] = pd.to_datetime(equity_df['Time'])
             fig = px.line(equity_df, x='Time', y='Equity', title=f'Equity Curve ({filename})')
             fig.update_layout(yaxis_title="Portfolio Value")
             figures["equity_fig"] = fig

        # --- Drawdown Curve ---
        drawdown_df = strategy_results.get("drawdown_df")
        if drawdown_df is not None and not drawdown_df.empty and 'Time' in drawdown_df.columns and 'Drawdown' in drawdown_df.columns:
             # Ensure Time is datetime
             drawdown_df['Time'] = pd.to_datetime(drawdown_df['Time'])
             # Convert drawdown to percentage for plotting
             drawdown_df['Drawdown_pct'] = drawdown_df['Drawdown'] * 100
             fig = px.area(drawdown_df, x='Time', y='Drawdown_pct', title=f'Drawdown Curve (%) ({filename})', labels={'Drawdown_pct': 'Drawdown (%)'})
             fig.update_layout(yaxis_title="Drawdown (%)")
             figures["drawdown_fig"] = fig

        # --- Equity vs Benchmark ---
        benchmark_df = strategy_results.get("benchmark_df")
        # Requires both equity and benchmark data
        if equity_df is not None and not equity_df.empty and 'Time' in equity_df.columns and 'Equity' in equity_df.columns and \
           benchmark_df is not None and not benchmark_df.empty and 'Time' in benchmark_df.columns and 'Benchmark' in benchmark_df.columns:
            try:
                # Ensure Time columns are datetime
                equity_df['Time'] = pd.to_datetime(equity_df['Time'])
                benchmark_df['Time'] = pd.to_datetime(benchmark_df['Time'])

                # Merge on Time after setting as index
                equity_indexed = equity_df.set_index('Time')['Equity']
                benchmark_indexed = benchmark_df.set_index('Time')['Benchmark']

                # Combine, handling potential different start/end dates
                combined = pd.concat([equity_indexed, benchmark_indexed], axis=1, keys=['Equity', 'Benchmark'], join='outer')

                # Normalize to start at 1 (or 100) for comparison
                # Check if first row has NaN values after outer join
                first_valid_index = combined.first_valid_index()
                if first_valid_index is not None:
                    # Normalize using the first non-NaN value for each column
                    normalized_equity = (combined['Equity'] / combined['Equity'].loc[combined['Equity'].first_valid_index()])#.fillna(method='ffill') # Optional fill
                    normalized_benchmark = (combined['Benchmark'] / combined['Benchmark'].loc[combined['Benchmark'].first_valid_index()])#.fillna(method='ffill') # Optional fill

                    # Create figure and add traces
                    fig = go.Figure()
                    fig.add_trace(go.Scatter(x=normalized_equity.index, y=normalized_equity, mode='lines', name='Strategy Equity'))
                    fig.add_trace(go.Scatter(x=normalized_benchmark.index, y=normalized_benchmark, mode='lines', name='Benchmark'))
                    fig.update_layout(title=f'Normalized Equity vs Benchmark ({filename})', xaxis_title='Date', yaxis_title='Normalized Value (Start = 1)')
                    figures["benchmark_fig"] = fig
                else:
                    print("Could not normalize Equity vs Benchmark: No valid starting point found after merge.")
                    figures["benchmark_fig"] = create_empty_figure(f"Equity vs Benchmark ({filename}) - Normalization Failed")

            except Exception as merge_err:
                 print(f"Error merging/plotting Equity vs Benchmark: {merge_err}")
                 figures["benchmark_fig"] = create_empty_figure(f"Equity vs Benchmark ({filename}) - Error")


        # --- Trade P/L Distribution ---
        trades_df = strategy_results.get("trades_df")
        if trades_df is not None and not trades_df.empty and 'profitLoss' in trades_df.columns:
            # Ensure profitLoss is numeric
            trades_df['profitLoss'] = pd.to_numeric(trades_df['profitLoss'], errors='coerce')
            valid_pnl = trades_df['profitLoss'].dropna()
            if not valid_pnl.empty:
                fig = px.histogram(valid_pnl, title=f'Trade Profit/Loss Distribution ({filename})', labels={'value': 'Profit/Loss'})
                figures["pnl_hist_fig"] = fig

        # --- Trade Duration Distribution ---
        # Uses 'duration_days' calculated in processing.py
        if trades_df is not None and not trades_df.empty and 'duration_days' in trades_df.columns:
             # Ensure duration_days is numeric
             trades_df['duration_days'] = pd.to_numeric(trades_df['duration_days'], errors='coerce')
             valid_duration = trades_df['duration_days'].dropna()
             if not valid_duration.empty:
                fig = px.histogram(valid_duration, title=f'Trade Duration Distribution (Days) ({filename})', labels={'value': 'Duration (Days)'})
                figures["duration_hist_fig"] = fig

        # --- Exposure Chart ---
        # Exposure data format varies; this is a basic example assuming a dict of series
        exposure_series_dict = strategy_results.get("exposure_series")
        if exposure_series_dict and isinstance(exposure_series_dict, dict):
            fig = go.Figure()
            exposure_plotted = False
            for series_name, series_data in exposure_series_dict.items():
                if 'values' in series_data and isinstance(series_data['values'], list):
                    # Process this specific series using the timeseries helper
                    exposure_df = process_timeseries_chart(series_data['values'], series_name)
                    if not exposure_df.empty:
                        # Plot as area chart if 'Exposure' in name, else line
                        plot_type = 'area' if 'Exposure' in series_name else 'scatter'
                        fill_type = 'tozeroy' if plot_type == 'area' else None
                        fig.add_trace(go.Scatter(x=exposure_df.index, y=exposure_df[series_name],
                                                 mode='lines', name=series_name, fill=fill_type))
                        exposure_plotted = True
            if exposure_plotted:
                fig.update_layout(title=f'Exposure ({filename})', xaxis_title='Date', yaxis_title='Value / % Exposure')
                figures["exposure_fig"] = fig
            else:
                 figures["exposure_fig"] = create_empty_figure(f"Exposure ({filename}) - No PlotData")
        else:
             figures["exposure_fig"] = create_empty_figure(f"Exposure ({filename}) - Data Missing/Invalid")


        # --- Portfolio Turnover ---
        turnover_df = strategy_results.get("turnover_df")
        if turnover_df is not None and not turnover_df.empty and 'Time' in turnover_df.columns and 'Turnover' in turnover_df.columns:
             # Ensure Time is datetime
             turnover_df['Time'] = pd.to_datetime(turnover_df['Time'])
             fig = px.line(turnover_df, x='Time', y='Turnover', title=f'Portfolio Turnover ({filename})')
             fig.update_layout(yaxis_title="Turnover")
             figures["turnover_fig"] = fig

    except Exception as e:
        print(f"Error generating figures for {filename}: {e}")
        traceback.print_exc()
        # Keep default empty figures on error

    return figures


def generate_manual_risk_figures(analysis_results, filename="Strategy"):
    """
    Generates Plotly figures from manually calculated risk analysis results.

    Args:
        analysis_results: Dictionary containing results from calculate_manual_risk_stats.
                          Expected keys: 'monthly_returns_table_for_heatmap', 'rolling_vol_df'.
        filename: Name of the strategy for figure titles.

    Returns:
        A dictionary containing Plotly figure objects:
        'monthly_heatmap_fig', 'rolling_vol_fig'.
        Uses empty figures if data is missing or invalid.
    """
    figures = {
        "monthly_heatmap_fig": create_empty_figure(f"Monthly Returns Heatmap ({filename})"),
        "rolling_vol_fig": create_empty_figure(f"Rolling Volatility ({filename})")
    }

    try:
        # --- Monthly Returns Heatmap ---
        # Expects percentages (values * 100) from calculate_manual_risk_stats
        monthly_ret_table = analysis_results.get("monthly_returns_table_for_heatmap")
        if monthly_ret_table is not None and not monthly_ret_table.empty:
            z = monthly_ret_table.values # The percentage values
            x = monthly_ret_table.columns # Month names
            y = monthly_ret_table.index   # Years

            # Create heatmap
            fig = go.Figure(data=go.Heatmap(
                z=z, x=x, y=y,
                colorscale='RdYlGn', # Red-Yellow-Green scale, good for returns
                zmid=0,             # Center color scale around zero
                # Format text labels shown on the heatmap cells
                text=monthly_ret_table.applymap(lambda v: f'{v:.1f}%' if pd.notna(v) else '').values,
                texttemplate="%{text}", # Use the formatted text
                hoverongaps=False,      # Don't show hover info for gaps
                colorbar=dict(title='Monthly Return (%)') # Add color bar title
                ))
            fig.update_layout(
                 title=f'Monthly Returns (%) ({filename})',
                 yaxis_nticks=len(y), # Ensure all years are shown as ticks
                 yaxis_title="Year",
                 yaxis_autorange='reversed' # Show earlier years at the top
                 )
            figures["monthly_heatmap_fig"] = fig

        # --- Rolling Volatility Plot ---
        rolling_vol_df = analysis_results.get("rolling_vol_df")
        # Check if DataFrame exists, is not empty, and has the 'Time' column
        if rolling_vol_df is not None and not rolling_vol_df.empty and 'Time' in rolling_vol_df.columns:
            # Ensure Time is datetime
            rolling_vol_df['Time'] = pd.to_datetime(rolling_vol_df['Time'])

            fig = go.Figure()
            colors = px.colors.qualitative.Plotly # Get a qualitative color sequence
            i = 0 # Color index
            vol_plotted = False
            # Iterate through columns starting with 'vol_'
            for col in rolling_vol_df.columns:
                if col.startswith('vol_'):
                    window_label = col.split('_')[1] # Extract window label (e.g., '3M')
                    # Plot volatility as percentage
                    fig.add_trace(go.Scatter(
                         x=rolling_vol_df['Time'],
                         y=rolling_vol_df[col] * 100, # Convert to percentage
                         mode='lines',
                         name=f'Rolling Vol ({window_label})',
                         line=dict(color=colors[i % len(colors)]) # Cycle through colors
                         ))
                    i += 1
                    vol_plotted = True

            # Update layout if at least one volatility series was plotted
            if vol_plotted:
                fig.update_layout(
                    title=f'Annualized Rolling Volatility ({filename})',
                    xaxis_title='Date',
                    yaxis_title='Volatility (%)' # Y-axis label as percentage
                    )
                figures["rolling_vol_fig"] = fig
            else:
                 figures["rolling_vol_fig"] = create_empty_figure(f"Rolling Volatility ({filename}) - No Plot Data")
        else:
             figures["rolling_vol_fig"] = create_empty_figure(f"Rolling Volatility ({filename}) - Data Missing/Invalid")


    except Exception as e:
        print(f"Error generating manual risk figures for {filename}: {e}")
        traceback.print_exc()
        # Keep default empty figures on error

    return figures