"""
VBT Visualization Dashboard — Plotly-based charting for VBT backtest results.

Produces professional, interactive visualizations using Plotly:
  - Equity curve with benchmark overlay
  - Drawdown chart
  - Rolling Sharpe / Sortino / Volatility
  - Trade markers on price
  - Returns distribution histogram
  - Monthly PnL heatmap
  - Gross vs net comparison
  - Parameter sensitivity heatmaps
  - Signal distribution analysis
  - Multi-panel research dashboard

All figures are Plotly go.Figure objects — interactive and exportable as HTML.

Usage:
    from backtests.vbt_viz import VBTVisualizer
    viz = VBTVisualizer()
    fig = viz.equity_curve(portfolio, benchmark_close)
    fig.show()
    viz.save_dashboard(portfolio, close, entries, exits, benchmark, "reports/dash.html")
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Any, Optional

import numpy as np
import pandas as pd

logger = logging.getLogger(__name__)

# Map VectorBT frequency string to friendly label
FREQ_LABELS = {
    "1m": "1 Minute", "5m": "5 Minutes", "15m": "15 Minutes", "30m": "30 Minutes",
    "1h": "1 Hour", "4h": "4 Hours", "8h": "8 Hours", "1d": "1 Day", "1w": "1 Week",
}


class VBTVisualizer:
    """Comprehensive visualization suite for VBT backtest results.

    All methods return Plotly go.Figure objects. Use .show() for interactive
    display or .write_html() for standalone reports.
    """

    def __init__(self, output_dir: str = "backtests/reports", dpi: int = 150):
        self._output_dir = Path(output_dir)
        self._output_dir.mkdir(parents=True, exist_ok=True)
        self._dpi = dpi

    def equity_curve(
        self,
        pf,
        benchmark_close: Optional[pd.Series] = None,
        title: str = "Portfolio Equity Curve",
    ):
        """Equity curve with optional benchmark overlay."""
        import plotly.graph_objects as go
        from plotly.subplots import make_subplots

        value = pf.value().dropna()
        if len(value) < 2:
            return go.Figure()

        fig = make_subplots(specs=[[{"secondary_y": True}]])
        fig.add_trace(
            go.Scatter(
                x=value.index, y=value.values, mode="lines",
                name="Portfolio Equity", line=dict(color="#1f77b4", width=1.5),
                fill="tozeroy", fillcolor="rgba(31,119,180,0.05)",
            ),
            secondary_y=False,
        )

        if benchmark_close is not None:
            try:
                bm = benchmark_close.reindex(value.index)
                bm_value = bm / bm.dropna().iloc[0] * value.iloc[0]
                fig.add_trace(
                    go.Scatter(
                        x=bm_value.index, y=bm_value.values, mode="lines",
                        name="Buy & Hold", line=dict(color="#7f7f7f", width=1, dash="dash"),
                        opacity=0.7,
                    ),
                    secondary_y=False,
                )
            except Exception:
                pass

        fig.update_layout(
            title=title,
            xaxis_title="Date",
            yaxis_title="Equity ($)",
            template="plotly_white",
            hovermode="x unified",
            legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1),
            margin=dict(l=60, r=30, t=50, b=60),
        )
        return fig

    def drawdown(
        self,
        pf,
        title: str = "Drawdown",
    ):
        """Drawdown chart — peak-to-trough percentage."""
        import plotly.graph_objects as go

        value = pf.value().dropna()
        if len(value) < 2:
            return go.Figure()

        peak = value.cummax()
        dd = (value - peak) / peak * 100

        fig = go.Figure()
        fig.add_trace(
            go.Scatter(
                x=dd.index, y=dd.values, mode="lines",
                name="Drawdown", line=dict(color="#d62728", width=1),
                fill="tozeroy", fillcolor="rgba(214,39,40,0.15)",
            )
        )
        fig.add_hline(
            y=0, line=dict(color="black", width=0.5, dash="dot"),
        )

        max_dd = dd.min()
        max_dd_date = dd.idxmin()
        fig.add_annotation(
            x=max_dd_date, y=max_dd,
            text=f"Max DD: {max_dd:.1f}%",
            showarrow=True, arrowhead=1,
            ay=40, bgcolor="white", bordercolor="#d62728",
        )

        fig.update_layout(
            title=title,
            xaxis_title="Date",
            yaxis_title="Drawdown (%)",
            template="plotly_white",
            hovermode="x unified",
            margin=dict(l=60, r=30, t=50, b=60),
        )
        return fig

    def rolling_metrics(
        self,
        pf,
        window: int = 90,
        title: str = "Rolling Performance Metrics",
    ):
        """Rolling Sharpe, Sortino, and Volatility over a window of bars."""
        import plotly.graph_objects as go
        from plotly.subplots import make_subplots

        returns = pf.returns().dropna()
        if len(returns) < window:
            return go.Figure()

        w = min(window, len(returns) // 2)

        roll_mean = returns.rolling(w).mean()
        roll_std = returns.rolling(w).std()
        rolling_sharpe = (roll_mean / roll_std.replace(0, np.nan)) * np.sqrt(365 * 24)
        rolling_sharpe = rolling_sharpe.dropna()

        down_std = returns * (returns < 0)
        rolling_sortino = (roll_mean / down_std.rolling(w).std().replace(0, np.nan)) * np.sqrt(365 * 24)
        rolling_sortino = rolling_sortino.dropna()

        ann_vol = roll_std * np.sqrt(365 * 24) * 100

        fig = make_subplots(
            rows=2, cols=1,
            shared_xaxes=True,
            vertical_spacing=0.08,
            subplot_titles=("Rolling Sharpe & Sortino (annualized)", "Rolling Volatility (annualized %)"),
        )

        fig.add_trace(
            go.Scatter(
                x=rolling_sharpe.index, y=rolling_sharpe.values,
                name="Sharpe", line=dict(color="#1f77b4", width=1),
            ),
            row=1, col=1,
        )
        fig.add_trace(
            go.Scatter(
                x=rolling_sortino.index, y=rolling_sortino.values,
                name="Sortino", line=dict(color="#ff7f0e", width=1),
            ),
            row=1, col=1,
        )
        fig.add_hline(y=0, line=dict(color="red", width=0.5, dash="dot"), row=1, col=1)

        fig.add_trace(
            go.Scatter(
                x=ann_vol.index, y=ann_vol.values,
                name="Volatility", line=dict(color="#2ca02c", width=1),
                fill="tozeroy", fillcolor="rgba(44,160,44,0.08)",
            ),
            row=2, col=1,
        )

        fig.update_layout(
            title=title,
            template="plotly_white",
            hovermode="x unified",
            margin=dict(l=60, r=30, t=50, b=60),
        )
        fig.update_yaxes(title_text="Ratio", row=1, col=1)
        fig.update_yaxes(title_text="Vol (%)", row=2, col=1)
        return fig

    def trade_markers(
        self,
        close: pd.Series,
        entries: pd.Series,
        exits: pd.Series,
        pf=None,
        title: str = "Trade Entry & Exit Markers",
    ):
        """Price chart with entry (green) and exit (red) markers."""
        import plotly.graph_objects as go

        fig = go.Figure()

        # Price line
        fig.add_trace(
            go.Scatter(
                x=close.index, y=close.values, mode="lines",
                name="Close Price", line=dict(color="#1f77b4", width=1),
            )
        )

        # Entry markers
        entry_idx = entries[entries].index
        entry_prices = [close.loc[i] for i in entry_idx if i in close.index]
        if entry_prices:
            fig.add_trace(
                go.Scatter(
                    x=list(entry_idx), y=entry_prices, mode="markers",
                    name="Entry", marker=dict(symbol="triangle-up", size=8,
                                               color="#2ca02c", line=dict(width=1)),
                )
            )

        # Exit markers
        exit_idx = exits[exits].index
        exit_prices = [close.loc[i] for i in exit_idx if i in close.index]
        if exit_prices:
            fig.add_trace(
                go.Scatter(
                    x=list(exit_idx), y=exit_prices, mode="markers",
                    name="Exit", marker=dict(symbol="triangle-down", size=8,
                                              color="#d62728", line=dict(width=1)),
                )
            )

        fig.update_layout(
            title=title,
            xaxis_title="Date",
            yaxis_title="Price",
            template="plotly_white",
            hovermode="x unified",
            legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1),
            margin=dict(l=60, r=30, t=50, b=60),
        )
        return fig

    def returns_distribution(
        self,
        pf,
        title: str = "Returns Distribution",
    ):
        """Histogram of trade returns with normal distribution overlay."""
        import plotly.graph_objects as go

        if pf is None:
            return go.Figure()

        returns = pf.returns().dropna()
        if len(returns) < 2:
            return go.Figure()
            
        try:
            from scipy import stats
        except ImportError:
            return go.Figure()

        returns_pct = returns * 100

        fig = go.Figure()
        fig.add_trace(
            go.Histogram(
                x=returns_pct.values, nbinsx=50,
                name="Returns", marker_color="#1f77b4", opacity=0.7,
                histnorm="probability density",
            )
        )

        mu = returns.mean() * 100
        sigma = returns.std() * 100
        x_range = np.linspace(mu - 4 * sigma, mu + 4 * sigma, 200)
        normal_pdf = stats.norm.pdf(x_range, mu, sigma)

        fig.add_trace(
            go.Scatter(
                x=x_range, y=normal_pdf, mode="lines",
                name=f"Normal (μ={mu:.3f}%, σ={sigma:.3f}%)",
                line=dict(color="#ff7f0e", width=2),
            )
        )

        skewness = float(pd.Series(returns.values).skew())
        kurtosis = float(pd.Series(returns.values).kurtosis())

        fig.add_vline(x=0, line=dict(color="red", width=0.5, dash="dot"))
        fig.add_annotation(
            x=0.98, y=0.95, xref="paper", yref="paper",
            text=f"Skew: {skewness:.3f}<br>Excess Kurt: {kurtosis:.3f}<br>N: {len(returns)}",
            showarrow=False, bgcolor="white", bordercolor="#ccc",
            xanchor="right", yanchor="top",
        )

        fig.update_layout(
            title=title,
            xaxis_title="Return (%)",
            yaxis_title="Density",
            template="plotly_white",
            margin=dict(l=60, r=30, t=50, b=60),
        )
        return fig

    def monthly_heatmap(
        self,
        pf,
        title: str = "Monthly Returns Heatmap",
    ):
        """Calendar-style monthly PnL heatmap (averages to daily frequency)."""
        import plotly.graph_objects as go

        if pf is None:
            return go.Figure()

        value = pf.value().dropna()
        if len(value) < 20:
            return go.Figure()

        if hasattr(value.index, 'freq') or len(value) > 100:
            resampled = value.resample("D").ffill()
        else:
            resampled = value

        returns = resampled.pct_change().dropna()

        monthly = returns.groupby([returns.index.year, returns.index.month]).apply(
            lambda x: (1 + x).prod() - 1
        ) * 100
        monthly.index = monthly.index.set_names(["Year", "Month"])
        monthly = monthly.reset_index(name="Return")

        if monthly.empty:
            return go.Figure()

        pivot = monthly.pivot(index="Year", columns="Month", values="Return")

        months = ["Jan", "Feb", "Mar", "Apr", "May", "Jun",
                   "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"]
        years = pivot.index.astype(str).tolist()

        fig = go.Figure(
            data=go.Heatmap(
                z=pivot.values,
                x=[months[m - 1] for m in pivot.columns if m <= 12],
                y=years,
                colorscale=[
                    [0.0, "#d62728"],
                    [0.45, "#ffffff"],
                    [0.5, "#eeeeee"],
                    [0.55, "#ffffff"],
                    [1.0, "#2ca02c"],
                ],
                zmid=0,
                text=[[f"{v:.1f}%" if not np.isnan(v) else "" for v in row]
                       for row in pivot.values],
                texttemplate="%{text}",
                colorbar=dict(title="Return (%)"),
            )
        )

        fig.update_layout(
            title=title,
            template="plotly_white",
            margin=dict(l=60, r=30, t=50, b=60),
        )
        return fig

    def gross_vs_net(
        self,
        pf,
        title: str = "Gross vs Net Performance",
    ):
        """Fee impact visualization — gross returns vs net returns."""
        import plotly.graph_objects as go
        from plotly.subplots import make_subplots

        if pf is None:
            return go.Figure()

        value = pf.value().dropna()
        if len(value) < 2:
            return go.Figure()

        net_rets = value.pct_change().dropna()

        fig = make_subplots(
            rows=2, cols=1,
            shared_xaxes=True,
            vertical_spacing=0.08,
            subplot_titles=("Cumulative Net Return", "Fee Impact per Bar"),
        )

        cum_net = (1 + net_rets).cumprod()
        fig.add_trace(
            go.Scatter(
                x=cum_net.index, y=(cum_net.values - 1) * 100,
                name="Net Return", line=dict(color="#1f77b4", width=1),
            ),
            row=1, col=1,
        )

        fee_est = value * 0.0005 + value * 0.001
        total_fees = 0.0
        fee_impact = []
        for i in range(len(value)):
            if int(value.index[i].timestamp()) % 10 == 0:
                total_fees += fee_est.iloc[i]
            fee_impact.append(total_fees)

        fig.add_trace(
            go.Scatter(
                x=value.index, y=fee_impact, mode="lines",
                name="Estimated Fees", line=dict(color="#d62728", width=1, dash="dash"),
            ),
            row=2, col=1,
        )

        trade_count = 0
        try:
            trade_count = int(pf.trades.count())
        except Exception:
            pass

        fig.add_annotation(
            x=0.98, y=0.95, xref="paper", yref="paper",
            text=f"Trades: {trade_count}",
            showarrow=False, bgcolor="white", bordercolor="#ccc",
            xanchor="right", yanchor="top",
        )

        fig.update_layout(
            title=title,
            template="plotly_white",
            hovermode="x unified",
            margin=dict(l=60, r=30, t=50, b=60),
        )
        fig.update_yaxes(title_text="Cumulative (%)", row=1, col=1)
        fig.update_yaxes(title_text="Fees ($)", row=2, col=1)
        return fig

    def holding_periods(
        self,
        pf,
        title: str = "Trade Holding Periods",
    ):
        """Histogram of trade durations."""
        import plotly.graph_objects as go

        try:
            trades = pf.trades
            records = trades.records_readable
            if records.empty:
                return go.Figure()
            durations = records.get("Duration", pd.Series(dtype=str))
            if durations.empty:
                return go.Figure()
        except Exception:
            return go.Figure()

        duration_vals = []
        for d in durations:
            try:
                td = pd.Timedelta(d)
                duration_vals.append(td.total_seconds() / 3600)
            except Exception:
                pass

        if not duration_vals:
            return go.Figure()

        dur = np.array(duration_vals)

        fig = go.Figure()
        fig.add_trace(
            go.Histogram(
                x=dur, nbinsx=30, name="Hold Duration (hours)",
                marker_color="#1f77b4", opacity=0.7,
            )
        )

        fig.add_vline(
            x=np.median(dur), line=dict(color="#d62728", width=1, dash="dash"),
        )
        fig.add_annotation(
            x=np.median(dur), y=0,
            text=f"Median: {np.median(dur):.1f}h<br>Mean: {np.mean(dur):.1f}h",
            showarrow=True, arrowhead=2, ay=-60,
            bgcolor="white", bordercolor="#d62728",
        )

        fig.update_layout(
            title=title,
            xaxis_title="Hold Duration (hours)",
            yaxis_title="Count",
            template="plotly_white",
            margin=dict(l=60, r=30, t=50, b=60),
        )
        return fig

    def param_heatmap(
        self,
        sweep_df: pd.DataFrame,
        x_param: str,
        y_param: str,
        metric: str = "sharpe",
        title: str = "Parameter Sensitivity",
    ):
        """Heatmap of a metric across a 2D parameter grid."""
        import plotly.graph_objects as go

        if sweep_df is None or sweep_df.empty:
            return go.Figure()

        pivot = sweep_df.pivot(index=y_param, columns=x_param, values=metric)

        fig = go.Figure(
            data=go.Heatmap(
                z=pivot.values,
                x=pivot.columns.astype(str).tolist(),
                y=pivot.index.astype(str).tolist(),
                colorscale="RdYlGn" if metric in ("sharpe", "total_return") else "RdYlGn_r",
                text=[[f"{v:.3f}" if not np.isnan(v) else "" for v in row]
                       for row in pivot.values],
                texttemplate="%{text}",
                colorbar=dict(title=metric.replace("_", " ").title()),
            )
        )

        fig.update_layout(
            title=f"{title} — {metric.replace('_', ' ').title()}",
            xaxis_title=x_param,
            yaxis_title=y_param,
            template="plotly_white",
            margin=dict(l=60, r=30, t=50, b=60),
        )
        return fig

    def dashboard(
        self,
        pf,
        close: pd.Series,
        entries: pd.Series,
        exits: pd.Series,
        benchmark_close: Optional[pd.Series] = None,
        sweep_df: Optional[pd.DataFrame] = None,
        strategy: str = "unknown",
        interval: str = "unknown",
    ) -> list:
        """Return a list of Plotly figures for the full dashboard."""
        figs = []

        # Row 1: Equity + Benchmark
        figs.append(self.equity_curve(
            pf, benchmark_close,
            title=f"Equity Curve — {strategy} ({interval})",
        ))

        # Row 2: Drawdown
        figs.append(self.drawdown(pf, title="Drawdown"))

        # Row 3: Rolling Metrics
        window = max(20, int(len(close) // 20))
        figs.append(self.rolling_metrics(pf, window=window,
                                          title=f"Rolling Metrics ({window}-bar window)"))

        # Row 4: Trade Markers
        figs.append(self.trade_markers(close, entries, exits, pf,
                                        title="Trade Signals on Price"))

        # Row 5: Returns Distribution
        figs.append(self.returns_distribution(pf))

        # Row 6: Monthly Heatmap
        figs.append(self.monthly_heatmap(pf))

        # Row 7: Gross vs Net
        figs.append(self.gross_vs_net(pf))

        # Row 8: Holding Periods
        figs.append(self.holding_periods(pf))

        # Optional: Param Heatmap
        if sweep_df is not None and not sweep_df.empty:
            figs.append(self.param_heatmap(sweep_df, "window", "threshold", "sharpe"))

        return figs

    def save_dashboard(
        self,
        pf,
        close: pd.Series,
        entries: pd.Series,
        exits: pd.Series,
        benchmark_close: Optional[pd.Series] = None,
        strategy: str = "unknown",
        interval: str = "unknown",
    ) -> str:
        """Save the full dashboard as a single HTML file."""
        import plotly.io as pio
        from datetime import datetime, timezone

        figs = self.dashboard(pf, close, entries, exits, benchmark_close,
                              strategy=strategy, interval=interval)

        ts = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S")
        fname = f"{strategy}_{interval}_{ts}_dashboard.html"
        fpath = self._output_dir / fname

        html_parts = ["<html><head>",
                       "<title>VBT Backtest Report</title>",
                       "<style>body{font-family:system-ui,sans-serif;max-width:1400px;"
                       "margin:0 auto;padding:20px;background:#f5f5f5;}"
                       ".chart{margin:20px 0;background:white;border-radius:8px;"
                       "box-shadow:0 2px 8px rgba(0,0,0,0.1);padding:10px;}"
                       "h1{color:#333;}</style>",
                       "</head><body>",
                       f"<h1>VBT Backtest — {strategy} ({interval})</h1>",
                       f"<p>Generated: {ts}</p>"]
        for i, fig in enumerate(figs):
            chart_id = f"chart_{i}"
            html_parts.append(f'<div class="chart" id="{chart_id}">')
            html_parts.append(pio.to_html(fig, include_plotlyjs="cdn", full_html=False))
            html_parts.append("</div>")
        html_parts.append("</body></html>")

        with open(fpath, "w") as f:
            f.write("\n".join(html_parts))

        logger.info("Dashboard saved to %s", fpath)
        return str(fpath)

    def save_sweep_report(
        self,
        sweep_df: pd.DataFrame,
        x_param: str = "window",
        y_param: str = "threshold",
        strategy: str = "unknown",
    ) -> str:
        """Save a parameter sweep report."""
        from datetime import datetime, timezone
        import plotly.io as pio

        metrics = [c for c in sweep_df.columns
                   if c not in (x_param, y_param) and sweep_df[c].dtype in (np.float64, np.int64)]
        if not metrics:
            metrics = [c for c in sweep_df.columns if c not in (x_param, y_param)]

        figs = []
        for metric in metrics[:6]:
            figs.append(self.param_heatmap(sweep_df, x_param, y_param, metric,
                                            title=f"Parameter Sweep"))

        ts = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S")
        fname = f"{strategy}_sweep_{ts}.html"
        fpath = self._output_dir / fname

        html_parts = ["<html><head>",
                       f"<title>Parameter Sweep — {strategy}</title>",
                       "<style>body{font-family:system-ui,sans-serif;max-width:1400px;"
                       "margin:0 auto;padding:20px;background:#f5f5f5;}"
                       ".chart{margin:20px 0;background:white;border-radius:8px;"
                       "box-shadow:0 2px 8px rgba(0,0,0,0.1);padding:10px;}</style>",
                       "</head><body>",
                       f"<h1>Parameter Sweep — {strategy}</h1>"]
        for i, fig in enumerate(figs):
            html_parts.append(f'<div class="chart">')
            html_parts.append(pio.to_html(fig, include_plotlyjs="cdn", full_html=False))
            html_parts.append("</div>")
        html_parts.append("</body></html>")

        with open(fpath, "w") as f:
            f.write("\n".join(html_parts))

        logger.info("Sweep report saved to %s", fpath)
        return str(fpath)


def plot_signal_distribution(
    entries: pd.Series,
    title: str = "Signal Distribution Over Time",
):
    """Signal occurrence over time — cumulative signal count."""
    import plotly.graph_objects as go

    cumulative = entries.cumsum()
    fig = go.Figure()
    fig.add_trace(
        go.Scatter(
            x=cumulative.index, y=cumulative.values, mode="lines",
            name="Cumulative Signals", line=dict(color="#1f77b4", width=1.5),
            fill="tozeroy", fillcolor="rgba(31,119,180,0.1)",
        )
    )
    fig.update_layout(
        title=title,
        xaxis_title="Date",
        yaxis_title="Cumulative Signal Count",
        template="plotly_white",
        margin=dict(l=60, r=30, t=50, b=60),
    )
    return fig


def plot_exposure_over_time(
    pf,
    title: str = "Position Exposure Over Time",
):
    """Position exposure as fraction of portfolio value."""
    import plotly.graph_objects as go

    try:
        pos = pf.position_mask().sum(axis=1)
        value = pf.value().dropna()
        aligned = pos.reindex(value.index).fillna(0)
    except Exception:
        return go.Figure()

    fig = go.Figure()
    fig.add_trace(
        go.Scatter(
            x=aligned.index, y=aligned.values, mode="lines",
            name="Exposure", line=dict(color="#9467bd", width=1),
            fill="tozeroy", fillcolor="rgba(148,103,189,0.1)",
        )
    )
    fig.add_hline(y=0, line=dict(color="black", width=0.5, dash="dot"))
    fig.update_layout(
        title=title,
        xaxis_title="Date",
        yaxis_title="Position Size",
        template="plotly_white",
        margin=dict(l=60, r=30, t=50, b=60),
    )
    return fig