chart_plotly.py

"""
Interactive Plotly Chart for Elliott Wave Analysis - TWII
Generates a professional financial analysis dashboard with:
  - Candlestick chart with Elliott Wave annotations and zigzag overlay
  - Volume bars with color coding
  - Equity curve from backtest results
"""

import sys
import os

# Ensure the script directory is on the path so we can import the backtest module
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import numpy as np
import pandas as pd
import plotly.graph_objects as go
from plotly.subplots import make_subplots

from elliott_wave_backtest import (
    fetch_data,
    find_swing_points,
    merge_swing_points,
    detect_impulse_waves,
    detect_corrective_waves,
    run_backtest,
    SYMBOL,
    PERIOD,
    INTERVAL,
    INITIAL_CAPITAL,
)

# ── Output paths ────────────────────────────────────────────────
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
HTML_PATH = os.path.join(BASE_DIR, "chart_plotly.html")
PNG_PATH = os.path.join(BASE_DIR, "chart_plotly.png")

# ── Color palette ───────────────────────────────────────────────
COLORS = {
    "bg": "#0e1117",
    "paper": "#161b22",
    "grid": "rgba(255,255,255,0.06)",
    "text": "#c9d1d9",
    "text_dim": "#8b949e",
    "candle_up": "#3fb950",
    "candle_down": "#f85149",
    "volume_up": "rgba(63,185,80,0.35)",
    "volume_down": "rgba(248,81,73,0.35)",
    "zigzag": "#58a6ff",
    "impulse": "#3fb950",
    "corrective": "#f0883e",
    "buy_marker": "#3fb950",
    "sell_marker": "#f85149",
    "equity": "#a371f7",
    "equity_fill": "rgba(163,113,247,0.12)",
    "accent_cyan": "#79c0ff",
}


def build_figure(df: pd.DataFrame, result: dict) -> go.Figure:
    """Build the full interactive Plotly figure with 3 subplot rows."""

    trades = result["trades"]
    zigzag = result["zigzag"]
    impulse_waves = result["impulse_waves"]
    corrective_waves = result["corrective_waves"]

    # ── Compute equity curve ────────────────────────────────────
    equity = INITIAL_CAPITAL
    equity_values = [equity]
    equity_dates = [trades[0].entry_date] if trades else [df.index[0]]
    for t in trades:
        equity *= (1 + t.pnl_pct / 100)
        equity_values.append(equity)
        if t.exit_date:
            equity_dates.append(t.exit_date)

    min_len = min(len(equity_dates), len(equity_values))
    equity_dates = equity_dates[:min_len]
    equity_values = equity_values[:min_len]

    # ── Determine candle colors for volume bars ─────────────────
    close_vals = df["Close"].values.flatten()
    open_vals = df["Open"].values.flatten()
    is_up = close_vals >= open_vals

    # ── Create subplots ─────────────────────────────────────────
    fig = make_subplots(
        rows=3, cols=1,
        shared_xaxes=True,
        vertical_spacing=0.03,
        row_heights=[0.55, 0.20, 0.25],
        subplot_titles=("", "", ""),
    )

    # ═══════════════════════════════════════════════════════════
    # ROW 1: Candlestick + Waves
    # ═══════════════════════════════════════════════════════════

    fig.add_trace(
        go.Candlestick(
            x=df.index,
            open=df["Open"].values.flatten(),
            high=df["High"].values.flatten(),
            low=df["Low"].values.flatten(),
            close=close_vals,
            increasing_line_color=COLORS["candle_up"],
            increasing_fillcolor=COLORS["candle_up"],
            decreasing_line_color=COLORS["candle_down"],
            decreasing_fillcolor=COLORS["candle_down"],
            name="TWII",
            showlegend=True,
            whiskerwidth=0.4,
        ),
        row=1, col=1,
    )

    # Zigzag line
    if zigzag:
        zz_dates = [df.index[p[0]] for p in zigzag]
        zz_prices = [p[1] for p in zigzag]
        fig.add_trace(
            go.Scatter(
                x=zz_dates,
                y=zz_prices,
                mode="lines",
                line=dict(color=COLORS["zigzag"], width=1.2, dash="dot"),
                name="ZigZag",
                opacity=0.6,
                hoverinfo="skip",
            ),
            row=1, col=1,
        )

    # ── Impulse wave annotations (numbered 0-5) ────────────────
    wave_labels_impulse = ["0", "1", "2", "3", "4", "5"]
    for w_idx, w in enumerate(impulse_waves):
        pts = w.points
        # Draw wave connecting lines
        wave_x = [df.index[p[0]] for p in pts]
        wave_y = [p[1] for p in pts]
        fig.add_trace(
            go.Scatter(
                x=wave_x, y=wave_y,
                mode="lines",
                line=dict(color=COLORS["impulse"], width=2),
                name=f"Impulse #{w_idx+1}" if w_idx == 0 else None,
                showlegend=(w_idx == 0),
                legendgroup="impulse",
                hoverinfo="skip",
                opacity=0.8,
            ),
            row=1, col=1,
        )
        # Labels
        for j, (idx, price) in enumerate(pts):
            y_shift = 18 if j % 2 == 1 else -18
            fig.add_annotation(
                x=df.index[idx], y=price,
                text=f"<b>{wave_labels_impulse[j]}</b>",
                showarrow=False,
                font=dict(size=12, color=COLORS["impulse"], family="Arial Black"),
                yshift=y_shift,
                row=1, col=1,
            )

    # ── Corrective wave annotations (A, B, C) ──────────────────
    wave_labels_corrective = ["A", "B", "C"]
    for w_idx, w in enumerate(corrective_waves):
        pts = w.points
        wave_x = [df.index[p[0]] for p in pts]
        wave_y = [p[1] for p in pts]
        fig.add_trace(
            go.Scatter(
                x=wave_x, y=wave_y,
                mode="lines",
                line=dict(color=COLORS["corrective"], width=2, dash="dash"),
                name=f"Corrective #{w_idx+1}" if w_idx == 0 else None,
                showlegend=(w_idx == 0),
                legendgroup="corrective",
                hoverinfo="skip",
                opacity=0.8,
            ),
            row=1, col=1,
        )
        for j, (idx, price) in enumerate(pts):
            if j < 3:
                y_shift = -18 if j % 2 == 0 else 18
                fig.add_annotation(
                    x=df.index[idx], y=price,
                    text=f"<b>{wave_labels_corrective[j]}</b>",
                    showarrow=False,
                    font=dict(size=12, color=COLORS["corrective"], family="Arial Black"),
                    yshift=y_shift,
                    row=1, col=1,
                )

    # ── Buy / Sell markers ──────────────────────────────────────
    buy_dates, buy_prices, buy_texts = [], [], []
    sell_dates, sell_prices, sell_texts = [], [], []
    exit_win_dates, exit_win_prices = [], []
    exit_loss_dates, exit_loss_prices = [], []

    for t in trades:
        pnl_str = f"{t.pnl_pct:+.2f}%"
        if t.direction == "long":
            buy_dates.append(t.entry_date)
            buy_prices.append(t.entry_price)
            buy_texts.append(f"BUY @ {t.entry_price:,.0f}")
        else:
            sell_dates.append(t.entry_date)
            sell_prices.append(t.entry_price)
            sell_texts.append(f"SELL @ {t.entry_price:,.0f}")

        if t.exit_date:
            if t.pnl > 0:
                exit_win_dates.append(t.exit_date)
                exit_win_prices.append(t.exit_price)
            else:
                exit_loss_dates.append(t.exit_date)
                exit_loss_prices.append(t.exit_price)

    if buy_dates:
        fig.add_trace(
            go.Scatter(
                x=buy_dates, y=buy_prices,
                mode="markers",
                marker=dict(
                    symbol="triangle-up", size=14,
                    color=COLORS["buy_marker"],
                    line=dict(color="white", width=1),
                ),
                name="Buy",
                text=buy_texts,
                hovertemplate="%{text}<extra></extra>",
            ),
            row=1, col=1,
        )

    if sell_dates:
        fig.add_trace(
            go.Scatter(
                x=sell_dates, y=sell_prices,
                mode="markers",
                marker=dict(
                    symbol="triangle-down", size=14,
                    color=COLORS["sell_marker"],
                    line=dict(color="white", width=1),
                ),
                name="Sell",
                text=sell_texts,
                hovertemplate="%{text}<extra></extra>",
            ),
            row=1, col=1,
        )

    if exit_win_dates:
        fig.add_trace(
            go.Scatter(
                x=exit_win_dates, y=exit_win_prices,
                mode="markers",
                marker=dict(symbol="x", size=10, color=COLORS["candle_up"],
                            line=dict(width=2, color=COLORS["candle_up"])),
                name="Exit (Win)",
                hovertemplate="Exit @ %{y:,.0f}<extra>Win</extra>",
            ),
            row=1, col=1,
        )

    if exit_loss_dates:
        fig.add_trace(
            go.Scatter(
                x=exit_loss_dates, y=exit_loss_prices,
                mode="markers",
                marker=dict(symbol="x", size=10, color=COLORS["candle_down"],
                            line=dict(width=2, color=COLORS["candle_down"])),
                name="Exit (Loss)",
                hovertemplate="Exit @ %{y:,.0f}<extra>Loss</extra>",
            ),
            row=1, col=1,
        )

    # ═══════════════════════════════════════════════════════════
    # ROW 2: Volume bars
    # ═══════════════════════════════════════════════════════════

    volume_vals = df["Volume"].values.flatten()
    vol_colors = [COLORS["volume_up"] if u else COLORS["volume_down"] for u in is_up]

    fig.add_trace(
        go.Bar(
            x=df.index,
            y=volume_vals,
            marker_color=vol_colors,
            name="Volume",
            showlegend=False,
            hovertemplate="Vol: %{y:,.0f}<extra></extra>",
        ),
        row=2, col=1,
    )

    # ═══════════════════════════════════════════════════════════
    # ROW 3: Equity curve
    # ═══════════════════════════════════════════════════════════

    if len(equity_dates) > 1:
        fig.add_trace(
            go.Scatter(
                x=equity_dates,
                y=equity_values,
                mode="lines",
                line=dict(color=COLORS["equity"], width=2.5),
                fill="tozeroy",
                fillcolor=COLORS["equity_fill"],
                name="Equity",
                hovertemplate="Equity: $%{y:,.0f}<extra></extra>",
            ),
            row=3, col=1,
        )

        # Baseline
        fig.add_hline(
            y=INITIAL_CAPITAL,
            line_dash="dash",
            line_color=COLORS["text_dim"],
            line_width=1,
            row=3, col=1,
        )

    # ═══════════════════════════════════════════════════════════
    # Layout and styling
    # ═══════════════════════════════════════════════════════════

    # Performance summary for subtitle
    total_trades = len(trades)
    wins = sum(1 for t in trades if t.pnl > 0)
    win_rate = (wins / total_trades * 100) if total_trades else 0
    total_return = sum(t.pnl_pct for t in trades) if trades else 0
    final_equity = equity_values[-1] if equity_values else INITIAL_CAPITAL

    dim_color = COLORS["text_dim"]
    subtitle_text = (
        f"<span style='color:{dim_color};font-size:13px'>"
        f"Trades: {total_trades}  |  "
        f"Win Rate: {win_rate:.1f}%  |  "
        f"Return: {total_return:+.2f}%  |  "
        f"Final Equity: ${final_equity:,.0f}"
        f"</span>"
    )

    fig.update_layout(
        template="plotly_dark",
        paper_bgcolor=COLORS["paper"],
        plot_bgcolor=COLORS["bg"],
        height=950,
        width=1600,
        title=dict(
            text=(
                f"<b>Elliott Wave Analysis</b> - TWII ({SYMBOL})<br>"
                f"{subtitle_text}"
            ),
            x=0.5,
            xanchor="center",
            font=dict(size=18, color=COLORS["text"], family="Segoe UI, Arial"),
        ),
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.01,
            xanchor="center",
            x=0.5,
            font=dict(size=10, color=COLORS["text"]),
            bgcolor="rgba(0,0,0,0)",
        ),
        margin=dict(l=60, r=30, t=110, b=40),
        hovermode="x unified",
        xaxis_rangeslider_visible=False,
    )

    # Remove the built-in rangeslider from candlestick (we add one on the bottom axis)
    fig.update_xaxes(
        rangeslider_visible=False,
        row=1, col=1,
    )

    # Add rangeslider on the bottom-most x-axis
    fig.update_xaxes(
        rangeslider_visible=True,
        rangeslider_thickness=0.04,
        row=3, col=1,
    )

    # Style all axes
    for i in range(1, 4):
        fig.update_xaxes(
            gridcolor=COLORS["grid"],
            zeroline=False,
            tickfont=dict(size=10, color=COLORS["text_dim"]),
            row=i, col=1,
        )
        fig.update_yaxes(
            gridcolor=COLORS["grid"],
            zeroline=False,
            tickfont=dict(size=10, color=COLORS["text_dim"]),
            tickformat=",.0f",
            row=i, col=1,
        )

    # Y-axis labels
    fig.update_yaxes(title_text="Price", title_font=dict(size=11, color=COLORS["text_dim"]), row=1, col=1)
    fig.update_yaxes(title_text="Volume", title_font=dict(size=11, color=COLORS["text_dim"]), row=2, col=1)
    fig.update_yaxes(title_text="Equity (TWD)", title_font=dict(size=11, color=COLORS["text_dim"]), row=3, col=1)

    # Hide weekends / market holidays for cleaner x-axis
    fig.update_xaxes(
        rangebreaks=[dict(bounds=["sat", "mon"])],
        row=1, col=1,
    )
    fig.update_xaxes(
        rangebreaks=[dict(bounds=["sat", "mon"])],
        row=2, col=1,
    )
    fig.update_xaxes(
        rangebreaks=[dict(bounds=["sat", "mon"])],
        row=3, col=1,
    )

    return fig


def main():
    print("=" * 60)
    print("  Elliott Wave - Interactive Plotly Chart Generator")
    print("=" * 60)

    # Download data
    df = fetch_data(SYMBOL, PERIOD, INTERVAL)

    # Run backtest (reuses detection logic)
    result = run_backtest(df)

    print(f"\nImpulse waves detected: {len(result['impulse_waves'])}")
    print(f"Corrective waves detected: {len(result['corrective_waves'])}")
    print(f"Trades executed: {len(result['trades'])}")

    # Build figure
    fig = build_figure(df, result)

    # Export HTML
    fig.write_html(
        HTML_PATH,
        include_plotlyjs="cdn",
        full_html=True,
        config={
            "scrollZoom": True,
            "displayModeBar": True,
            "modeBarButtonsToAdd": ["drawline", "drawrect", "eraseshape"],
            "toImageButtonOptions": {
                "format": "png",
                "width": 1920,
                "height": 1080,
                "scale": 2,
            },
        },
    )
    print(f"\nHTML saved: {HTML_PATH}")

    # Export PNG
    try:
        fig.write_image(PNG_PATH, width=1920, height=1080, scale=2)
        print(f"PNG  saved: {PNG_PATH}")
    except Exception as e:
        print(f"PNG export failed (kaleido/chrome issue): {e}")
        print("The HTML file is fully functional and can be opened in any browser.")

    print("\nDone.")


if __name__ == "__main__":
    main()