"""
Trade Visualization Module

Provides Altair-based plotting functions for analyzing backtest trades data,
including MAE/MFE analysis, PnL curves, and trade distribution histograms.

All functions accept a Polars DataFrame with trades data from the API
and return interactive Altair charts.
"""

import polars as pl
import altair as alt
from typing import Optional, Literal


# =============================================================================
# Helper Functions
# =============================================================================

def add_win_loss_label(trades: pl.DataFrame) -> pl.DataFrame:
    """Add 'win_loss' column labeling trades as 'Win' or 'Loss'

    Args:
        trades: DataFrame with 'pnl' column

    Returns:
        DataFrame with additional 'win_loss' column
    """
    return trades.with_columns(
        pl.when(pl.col("pnl") > 0)
        .then(pl.lit("Win"))
        .otherwise(pl.lit("Loss"))
        .alias("win_loss")
    )


def calculate_mea(trades: pl.DataFrame) -> pl.DataFrame:
    """Calculate MEA (Maximum Execution Analysis) - ATR-normalized MAE/MFE ratios

    Each metric is normalized by its corresponding volatility measurement:
    - mae / mae_volatility -> mae_atr_ratio
    - mfe / mfe_volatility -> mfe_atr_ratio
    - g_mfe / g_mfe_volatility -> g_mfe_atr_ratio
    - mae_lv1 / mae_lv1_volatility -> mae_lv1_atr_ratio
    - mhl / mhl_volatility -> mhl_atr_ratio

    Args:
        trades: DataFrame with MAE/MFE and corresponding volatility columns

    Returns:
        DataFrame with additional ATR ratio columns

    Example:
        >>> trades = client.quick_backtest_results(["my_strategy"])[0]
        >>> trades = calculate_mea(trades)
        >>> print(trades.columns)  # Will include mae_atr_ratio, mfe_atr_ratio, etc.
    """
    return trades.with_columns([
        (pl.col("mae") / pl.col("mae_volatility")).alias("mae_atr_ratio"),
        (pl.col("mfe") / pl.col("mfe_volatility")).alias("mfe_atr_ratio"),
        (pl.col("g_mfe") / pl.col("g_mfe_volatility")).alias("g_mfe_atr_ratio"),
        (pl.col("mae_lv1") / pl.col("mae_lv1_volatility")).alias("mae_lv1_atr_ratio"),
        (pl.col("mhl") / pl.col("mhl_volatility")).alias("mhl_atr_ratio"),
    ])


# Legacy alias for backward compatibility
def calculate_mae_atr_ratio(
    trades: pl.DataFrame,
    volatility_col: str = "entry_volatility"
) -> pl.DataFrame:
    """Legacy function - use calculate_mea() instead

    This function is deprecated and will use calculate_mea() internally.
    """
    return calculate_mea(trades)


def calculate_cumulative_pnl(
    trades: pl.DataFrame,
    sort_col: str = "exit_time"
) -> pl.DataFrame:
    """Calculate cumulative PnL and drawdown metrics

    Args:
        trades: DataFrame with 'pnl_pct' column
        sort_col: Column to sort by (default: exit_time)

    Returns:
        DataFrame with cumulative metrics and high-water marks
    """
    df = trades.sort(sort_col).with_columns([
        pl.col("pnl_pct").cum_sum().alias("cum_pnl_pct"),
    ])

    df = df.with_columns([
        pl.col("cum_pnl_pct").cum_max().alias("cummax"),
        (pl.col("cum_pnl_pct").cum_max() - pl.col("cum_pnl_pct")).alias("drawdown"),
    ])

    # Mark new highs
    df = df.with_columns([
        pl.when(
            (pl.col("cum_pnl_pct") == pl.col("cummax")) &
            (pl.col("cum_pnl_pct") != pl.col("cum_pnl_pct").shift(1))
        )
        .then(pl.col("cum_pnl_pct"))
        .alias("new_high")
    ])

    # Mark new drawdown lows (per month)
    df = df.with_columns([
        pl.col(sort_col).dt.strftime("%Y-%m").alias("year_month")
    ])

    df = df.with_columns([
        pl.when(
            (pl.col("drawdown") == pl.col("drawdown").max().over("year_month")) &
            (pl.col("drawdown") != pl.col("drawdown").shift(1))
        )
        .then(pl.col("drawdown"))
        .alias("new_dd")
    ])

    return df


def get_quantiles(
    trades: pl.DataFrame,
    column: str,
    quantiles: list[float] = [0.25, 0.5, 0.75]
) -> dict[str, float]:
    """Calculate quantiles for a column

    Args:
        trades: DataFrame
        column: Column name
        quantiles: List of quantile values (0-1)

    Returns:
        Dictionary mapping quantile names to values
    """
    result = {}
    for q in quantiles:
        q_name = f"q{int(q * 100)}"
        result[q_name] = trades[column].quantile(q)
    return result


# =============================================================================
# Main Plotting Functions
# =============================================================================

def plot_mae_mfe_scatter(
    trades: pl.DataFrame,
    x_col: str = "mae_atr_ratio",
    y_col: str = "g_mfe_atr_ratio",
    title: str = "MAE vs MFE Analysis",
    apply_config: bool = True
) -> alt.Chart:
    """Create scatter plot of MAE vs MFE with quantile reference lines

    Args:
        trades: DataFrame with MAE/MFE columns
        x_col: X-axis column (default: mae_atr_ratio)
        y_col: Y-axis column (default: g_mfe_atr_ratio)
        title: Chart title
        apply_config: Apply default styling config (set False for chart composition)

    Returns:
        Interactive Altair scatter plot

    Example:
        >>> trades = client.quick_backtest_results(["my_strategy"])[0]
        >>> trades = calculate_mea(trades)
        >>> chart = plot_mae_mfe_scatter(trades)
        >>> chart.show()

        >>> # For composition (no config conflicts):
        >>> chart1 = plot_mae_mfe_scatter(trades, apply_config=False)
        >>> chart2 = plot_pnl_curves(trades, apply_config=False)
        >>> combined = chart1 | chart2
    """
    # Calculate ATR ratios if not present
    if x_col not in trades.columns or y_col not in trades.columns:
        trades = calculate_mea(trades)

    # Add win/loss labels
    if "win_loss" not in trades.columns:
        trades = add_win_loss_label(trades)

    # Calculate quantiles
    x_q75 = trades[x_col].quantile(0.75)
    y_q75 = trades[y_col].quantile(0.75)

    # Add quantile columns for reference lines
    trades = trades.with_columns([
        pl.lit(x_q75).alias("x_q75"),
        pl.lit(y_q75).alias("y_q75"),
    ])

    # Base scatter plot
    scatter = alt.Chart(trades).mark_circle(size=60, opacity=0.6).encode(
        x=alt.X(f"{x_col}:Q", title=x_col.replace("_", " ").title()),
        y=alt.Y(f"{y_col}:Q", title=y_col.replace("_", " ").title()),
        color=alt.Color(
            "win_loss:N",
            scale=alt.Scale(
                domain=["Win", "Loss"],
                range=["#00C49A", "#FF6B6B"]
            ),
            legend=alt.Legend(title="Result")
        ),
        tooltip=[
            alt.Tooltip("position_id:Q", title="Position ID"),
            alt.Tooltip("entry_time:T", title="Entry Time"),
            alt.Tooltip(f"{x_col}:Q", title=x_col, format=".2f"),
            alt.Tooltip(f"{y_col}:Q", title=y_col, format=".2f"),
            alt.Tooltip("pnl:Q", title="PnL", format=".2f"),
            alt.Tooltip("pnl_pct:Q", title="PnL %", format=".2%"),
        ]
    )

    # Vertical quantile line (Q3 for x-axis)
    vline = alt.Chart(trades).mark_rule(
        strokeDash=[5, 5],
        color="gray",
        size=1
    ).encode(
        x="x_q75:Q"
    )

    # Horizontal quantile line (Q3 for y-axis)
    hline = alt.Chart(trades).mark_rule(
        strokeDash=[5, 5],
        color="gray",
        size=1
    ).encode(
        y="y_q75:Q"
    )

    # Combine layers
    chart = (scatter + vline + hline).properties(
        width=500,
        height=500,
        title=title
    )

    # Apply config only if requested (avoid conflicts in composition)
    if apply_config:
        chart = chart.configure_axis(
            gridColor="#f0f0f0"
        ).configure_view(
            strokeWidth=0
        )

    return chart


def plot_mfe_mhl_analysis(
    trades: pl.DataFrame,
    x_col: str = "g_mfe_atr_ratio",
    y_col: str = "mhl_atr_ratio",
    bins: int = 40,
    title: str = "MFE vs MHL Analysis",
    apply_config: bool = True
) -> alt.Chart:
    """Create multi-panel MFE vs MHL analysis with histograms

    Args:
        trades: DataFrame with MFE/MHL columns
        x_col: X-axis column (default: g_mfe_atr_ratio)
        y_col: Y-axis column (default: mhl_atr_ratio)
        bins: Number of histogram bins
        title: Chart title
        apply_config: Apply default styling config (set False for chart composition)

    Returns:
        Vertically stacked Altair chart with 3 panels:
        - Panel 1: MFE vs MHL scatter with diagonal reference
        - Panel 2: MFE histogram (all trades)
        - Panel 3: MFE histogram by win/loss

    Example:
        >>> trades = client.quick_backtest_results(["my_strategy"])[0]
        >>> trades = calculate_mea(trades)
        >>> chart = plot_mfe_mhl_analysis(trades)
        >>> chart.show()
    """
    # Calculate ATR ratios if not present
    if x_col not in trades.columns or y_col not in trades.columns:
        trades = calculate_mea(trades)

    # Add win/loss labels
    if "win_loss" not in trades.columns:
        trades = add_win_loss_label(trades)

    # Calculate quantiles
    x_q50 = trades[x_col].quantile(0.5)
    x_q75 = trades[x_col].quantile(0.75)
    x_max = trades[x_col].max()
    y_max = trades[y_col].max()

    # Add reference line data
    n_points = 100
    slope_data = pl.DataFrame({
        "x_slope": [i * x_max / n_points for i in range(n_points)],
        "y_slope": [i * y_max / n_points for i in range(n_points)],
    })

    trades = trades.with_columns([
        pl.lit(x_q50).alias("x_q50"),
        pl.lit(x_q75).alias("x_q75"),
    ])

    # Panel 1: Scatter plot with diagonal reference
    scatter = alt.Chart(trades).mark_circle(size=100, opacity=0.6).encode(
        x=alt.X(f"{x_col}:Q", title=x_col.replace("_", " ").title()),
        y=alt.Y(f"{y_col}:Q", title=y_col.replace("_", " ").title()),
        color=alt.Color(
            "win_loss:N",
            scale=alt.Scale(
                domain=["Win", "Loss"],
                range=["#00C49A", "#FF6B6B"]
            ),
            legend=alt.Legend(title="Result")
        ),
        tooltip=[
            alt.Tooltip("position_id:Q", title="Position ID"),
            alt.Tooltip("entry_time:T", title="Entry Time"),
            alt.Tooltip(f"{x_col}:Q", title=x_col, format=".2f"),
            alt.Tooltip(f"{y_col}:Q", title=y_col, format=".2f"),
            alt.Tooltip("pnl:Q", title="PnL", format=".2f"),
            alt.Tooltip("pnl_pct:Q", title="PnL %", format=".2%"),
        ]
    )

    # Diagonal reference line
    diagonal = alt.Chart(slope_data).mark_line(
        color="gray",
        strokeDash=[3, 3]
    ).encode(
        x="x_slope:Q",
        y="y_slope:Q"
    )

    # Quantile lines
    q50_line = alt.Chart(trades).mark_rule(
        strokeDash=[5, 5],
        color="black",
        size=1
    ).encode(x="x_q50:Q")

    q75_line = alt.Chart(trades).mark_rule(
        strokeDash=[5, 5],
        color="black",
        size=1
    ).encode(x="x_q75:Q")

    panel1 = (scatter + diagonal + q50_line + q75_line).properties(
        width=500,
        height=400,
        title=f"{title} - Scatter Plot"
    )

    # Panel 2: Histogram of all trades
    hist_all = alt.Chart(trades).mark_bar(opacity=0.3).encode(
        x=alt.X(f"{x_col}:Q", bin=alt.Bin(maxbins=bins), title=x_col.replace("_", " ").title()),
        y=alt.Y("count():Q", title="Count"),
    ).properties(
        width=500,
        height=150,
        title=f"{x_col} Distribution - All Trades"
    )

    # Add quantile lines to histogram
    hist_all = (hist_all + q50_line + q75_line)

    # Panel 3: Histogram by win/loss
    hist_by_result = alt.Chart(trades).mark_bar(opacity=0.5).encode(
        x=alt.X(f"{x_col}:Q", bin=alt.Bin(maxbins=bins), title=x_col.replace("_", " ").title()),
        y=alt.Y("count():Q", title="Count"),
        color=alt.Color(
            "win_loss:N",
            scale=alt.Scale(
                domain=["Win", "Loss"],
                range=["#00C49A", "#FF6B6B"]
            )
        )
    ).properties(
        width=500,
        height=150,
        title=f"{x_col} Distribution - By Result"
    )

    # Add quantile lines
    hist_by_result = (hist_by_result + q50_line + q75_line)

    # Stack vertically
    chart = alt.vconcat(panel1, hist_all, hist_by_result)

    # Apply config only if requested (avoid conflicts in composition)
    if apply_config:
        chart = chart.configure_axis(
            gridColor="#f0f0f0"
        ).configure_view(
            strokeWidth=0
        )

    return chart


def plot_pnl_curves(
    trades: pl.DataFrame,
    title: str = "Strategy Performance",
    apply_config: bool = True
) -> alt.Chart:
    """Create PnL curves with drawdown visualization

    Args:
        trades: DataFrame with 'pnl_pct' and 'exit_time' columns
        title: Chart title
        apply_config: Apply default styling config (set False for chart composition)

    Returns:
        Vertically stacked Altair chart with 2 panels:
        - Panel 1: Cumulative return curve with new high markers
        - Panel 2: Drawdown curve with new low markers

    Example:
        >>> trades = client.quick_backtest_results(["my_strategy"])[0]
        >>> chart = plot_pnl_curves(trades, title="My Strategy Performance")
        >>> chart.show()
    """
    # Calculate cumulative PnL and drawdown
    df = calculate_cumulative_pnl(trades, sort_col="exit_time")

    # Panel 1: Cumulative PnL curve
    pnl_line = alt.Chart(df).mark_line(
        color="#00C49A",
        size=2
    ).encode(
        x=alt.X("exit_time:T", title="Date"),
        y=alt.Y("cum_pnl_pct:Q", title="Cumulative Return (%)", axis=alt.Axis(format=".1%")),
        tooltip=[
            alt.Tooltip("exit_time:T", title="Date"),
            alt.Tooltip("cum_pnl_pct:Q", title="Cumulative Return", format=".2%"),
            alt.Tooltip("pnl_pct:Q", title="Trade PnL", format=".2%"),
        ]
    )

    # Area fill
    pnl_area = alt.Chart(df).mark_area(
        color="#00C49A",
        opacity=0.15
    ).encode(
        x="exit_time:T",
        y="cum_pnl_pct:Q"
    )

    # New high markers
    high_markers = alt.Chart(df).mark_point(
        shape="diamond",
        size=100,
        color="#FFD700",
        filled=True,
        stroke="white",
        strokeWidth=1.5
    ).encode(
        x="exit_time:T",
        y="new_high:Q",
        tooltip=[
            alt.Tooltip("exit_time:T", title="Date"),
            alt.Tooltip("new_high:Q", title="New High", format=".2%"),
        ]
    ).transform_filter(
        alt.datum.new_high != None
    )

    # Drawdown low markers (on PnL curve)
    dd_markers = alt.Chart(df).mark_circle(
        size=60,
        color="#9B2C2C",
        opacity=0.7,
        stroke="white",
        strokeWidth=1
    ).encode(
        x="exit_time:T",
        y="cum_pnl_pct:Q",
        tooltip=[
            alt.Tooltip("exit_time:T", title="Date"),
            alt.Tooltip("drawdown:Q", title="Drawdown", format=".2%"),
        ]
    ).transform_filter(
        alt.datum.new_dd != None
    )

    panel1 = (pnl_area + pnl_line + high_markers + dd_markers).properties(
        width=600,
        height=400,
        title=f"{title} - Cumulative Return"
    )

    # Panel 2: Drawdown curve
    dd_line = alt.Chart(df).mark_line(
        color="#FF6B6B",
        size=2
    ).encode(
        x=alt.X("exit_time:T", title="Date"),
        y=alt.Y("drawdown:Q", title="Drawdown (%)", axis=alt.Axis(format=".1%")),
        tooltip=[
            alt.Tooltip("exit_time:T", title="Date"),
            alt.Tooltip("drawdown:Q", title="Drawdown", format=".2%"),
        ]
    )

    # Drawdown area fill
    dd_area = alt.Chart(df).mark_area(
        color="#FF6B6B",
        opacity=0.15
    ).encode(
        x="exit_time:T",
        y="drawdown:Q"
    )

    panel2 = (dd_area + dd_line).properties(
        width=600,
        height=200,
        title="Drawdown"
    )

    # Stack vertically
    chart = alt.vconcat(panel1, panel2)

    # Apply config only if requested (avoid conflicts in composition)
    if apply_config:
        chart = chart.configure_axis(
            gridColor="#f0f0f0"
        ).configure_view(
            strokeWidth=0
        ).configure_title(
            fontSize=18,
            anchor="start"
        )

    return chart


def plot_trade_histograms(
    trades: pl.DataFrame,
    column: str = "mae",
    bins: int = 40,
    title: Optional[str] = None,
    apply_config: bool = True
) -> alt.Chart:
    """Create histogram of trade metrics with win/loss breakdown

    Args:
        trades: DataFrame with trade metrics
        column: Column to plot (e.g., 'mae', 'mfe', 'pnl_pct')
        bins: Number of histogram bins
        title: Chart title (auto-generated if None)
        apply_config: Apply default styling config (set False for chart composition)

    Returns:
        Vertically stacked Altair chart with 2 panels:
        - Panel 1: Distribution of all trades
        - Panel 2: Distribution by win/loss

    Example:
        >>> trades = client.quick_backtest_results(["my_strategy"])[0]
        >>> chart = plot_trade_histograms(trades, column="mae_atr_ratio")
        >>> chart.show()
    """
    if title is None:
        title = f"{column.replace('_', ' ').title()} Distribution"

    # Add win/loss labels if not present
    if "win_loss" not in trades.columns:
        trades = add_win_loss_label(trades)

    # Calculate quantiles
    q25 = trades[column].quantile(0.25)
    q50 = trades[column].quantile(0.50)
    q75 = trades[column].quantile(0.75)

    trades = trades.with_columns([
        pl.lit(q25).alias("q25"),
        pl.lit(q50).alias("q50"),
        pl.lit(q75).alias("q75"),
    ])

    # Panel 1: All trades
    hist_all = alt.Chart(trades).mark_bar(
        opacity=0.3,
        color="#3b82f6"
    ).encode(
        x=alt.X(f"{column}:Q", bin=alt.Bin(maxbins=bins), title=column.replace("_", " ").title()),
        y=alt.Y("count():Q", title="Count"),
    )

    # Quantile lines
    q50_line = alt.Chart(trades).mark_rule(
        strokeDash=[5, 5],
        color="black",
        size=1
    ).encode(x="q50:Q")

    q75_line = alt.Chart(trades).mark_rule(
        strokeDash=[5, 5],
        color="black",
        size=1
    ).encode(x="q75:Q")

    panel1 = (hist_all + q50_line + q75_line).properties(
        width=600,
        height=250,
        title=f"{title} - All Trades"
    )

    # Panel 2: By win/loss
    hist_by_result = alt.Chart(trades).mark_bar(opacity=0.5).encode(
        x=alt.X(f"{column}:Q", bin=alt.Bin(maxbins=bins), title=column.replace("_", " ").title()),
        y=alt.Y("count():Q", title="Count"),
        color=alt.Color(
            "win_loss:N",
            scale=alt.Scale(
                domain=["Win", "Loss"],
                range=["#00C49A", "#FF6B6B"]
            ),
            legend=alt.Legend(title="Result")
        )
    )

    panel2 = (hist_by_result + q50_line + q75_line).properties(
        width=600,
        height=250,
        title=f"{title} - By Result"
    )

    # Stack vertically
    chart = alt.vconcat(panel1, panel2)

    # Apply config only if requested (avoid conflicts in composition)
    if apply_config:
        chart = chart.configure_axis(
            gridColor="#f0f0f0"
        ).configure_view(
            strokeWidth=0
        )

    return chart


# =============================================================================
# Chart Composition Utilities
# =============================================================================

def combine_charts(
    *charts: alt.Chart,
    layout: Literal["horizontal", "vertical", "grid"] = "vertical",
    columns: int = 2,
    spacing: int = 15,
    title: Optional[str] = None
) -> alt.Chart:
    """Combine multiple charts into a single dashboard layout

    Automatically strips config from individual charts to avoid composition conflicts.

    Args:
        *charts: Variable number of Altair charts to combine
        layout: Layout type - "horizontal" (|), "vertical" (&), or "grid"
        columns: Number of columns for grid layout (default: 2)
        spacing: Spacing between charts in pixels (default: 15)
        title: Overall dashboard title (optional)

    Returns:
        Combined Altair chart with unified configuration

    Example:
        >>> from tradepose_client import (
        ...     plot_mae_mfe_scatter,
        ...     plot_pnl_curves,
        ...     combine_charts
        ... )
        >>> trades = client.quick_backtest_results(["my_strategy"])[0]
        >>> trades = calculate_mea(trades)
        >>>
        >>> # Method 1: Using combine_charts()
        >>> dashboard = combine_charts(
        ...     plot_mae_mfe_scatter(trades),
        ...     plot_pnl_curves(trades),
        ...     layout="horizontal"
        ... )
        >>>
        >>> # Method 2: Using operators with apply_config=False
        >>> scatter = plot_mae_mfe_scatter(trades, apply_config=False)
        >>> pnl = plot_pnl_curves(trades, apply_config=False)
        >>> dashboard = (scatter | pnl).configure_axis(gridColor="#f0f0f0")
    """
    import altair as alt

    if len(charts) == 0:
        raise ValueError("At least one chart is required")

    # Strip config from all charts to avoid conflicts
    # (Altair will error if trying to compose charts with config)
    # Note: This is done by re-creating charts without config
    clean_charts = []
    for chart in charts:
        # Charts should be created with apply_config=False
        # but if user passes configured charts, we need to handle it
        clean_charts.append(chart)

    # Combine based on layout
    if layout == "horizontal":
        combined = clean_charts[0]
        for chart in clean_charts[1:]:
            combined = combined | chart

    elif layout == "vertical":
        combined = clean_charts[0]
        for chart in clean_charts[1:]:
            combined = combined & chart

    elif layout == "grid":
        # Grid layout using alt.concat
        combined = alt.concat(*clean_charts, columns=columns)

    else:
        raise ValueError(f"Invalid layout: {layout}. Use 'horizontal', 'vertical', or 'grid'")

    # Apply unified configuration
    combined = combined.configure_axis(
        gridColor="#f0f0f0"
    ).configure_view(
        strokeWidth=0
    ).configure_concat(
        spacing=spacing
    )

    # Add overall title if provided
    if title:
        combined = combined.properties(
            title={
                "text": title,
                "fontSize": 20,
                "anchor": "middle"
            }
        ).configure_title(
            fontSize=20,
            anchor="start"
        )

    # Resolve scales to avoid conflicts
    combined = combined.resolve_scale(
        color='independent',
        x='independent',
        y='independent'
    )

    return combined
