chart_matplotlib.py

"""
Bloomberg Terminal-Style Elliott Wave Chart
Professional multi-panel visualization for TWII Elliott Wave backtest results.
"""

import sys
import numpy as np
import pandas as pd
import yfinance as yf
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import matplotlib.ticker as mticker
from matplotlib.gridspec import GridSpec
from matplotlib.patches import FancyArrowPatch
from datetime import datetime

# Import wave detection logic from backtest module
sys.path.insert(0, "E:/Developer/lufftw/repo/stock")
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, Trade
)

# ── Bloomberg-style Color Palette ──────────────────────────────
BG_COLOR       = "#0D1117"
PANEL_BG       = "#161B22"
GRID_COLOR     = "#21262D"
TEXT_COLOR      = "#C9D1D9"
TEXT_DIM        = "#8B949E"
ACCENT_ORANGE   = "#FF9800"
ACCENT_CYAN     = "#00BCD4"
ACCENT_GREEN    = "#00E676"
ACCENT_RED      = "#FF1744"
IMPULSE_COLOR   = "#00E5FF"
CORRECT_COLOR   = "#FF6D00"
ZIGZAG_COLOR    = "#5C6BC0"
UP_CANDLE       = "#26A69A"
DOWN_CANDLE     = "#EF5350"
UP_VOL          = "#26A69A"
DOWN_VOL        = "#EF5350"
EQUITY_COLOR    = "#42A5F5"
DRAWDOWN_COLOR  = "#EF5350"
LONG_MARKER     = "#00E676"
SHORT_MARKER    = "#E040FB"
WIN_EXIT        = "#76FF03"
LOSS_EXIT       = "#FF1744"


def build_equity_curve(trades, df):
    """Build daily equity curve from trades list."""
    equity = INITIAL_CAPITAL
    dates = [df.index[0]]
    values = [equity]

    for t in trades:
        equity *= (1 + t.pnl_pct / 100)
        if t.exit_date is not None:
            dates.append(t.exit_date)
            values.append(equity)

    # Ensure final date is included
    if dates[-1] != df.index[-1]:
        dates.append(df.index[-1])
        values.append(equity)

    return dates, values


def compute_drawdown(dates, values):
    """Compute drawdown series from equity curve."""
    arr = np.array(values)
    running_max = np.maximum.accumulate(arr)
    drawdown_pct = (running_max - arr) / running_max * 100
    return drawdown_pct


def draw_ohlc_bars(ax, df):
    """Draw OHLC bar chart (Bloomberg style thin bars, not fat candlesticks)."""
    dates = mdates.date2num(df.index.to_pydatetime())
    opens = df["Open"].values
    highs = df["High"].values
    lows = df["Low"].values
    closes = df["Close"].values

    for i in range(len(df)):
        color = UP_CANDLE if closes[i] >= opens[i] else DOWN_CANDLE
        # High-Low vertical line
        ax.plot([dates[i], dates[i]], [lows[i], highs[i]],
                color=color, linewidth=0.8, solid_capstyle="round")
        # Open tick (left)
        ax.plot([dates[i] - 0.3, dates[i]], [opens[i], opens[i]],
                color=color, linewidth=1.0, solid_capstyle="butt")
        # Close tick (right)
        ax.plot([dates[i], dates[i] + 0.3], [closes[i], closes[i]],
                color=color, linewidth=1.0, solid_capstyle="butt")


def main():
    # ── 1. Download Data ───────────────────────────────────────
    print("Downloading TWII data...")
    df = fetch_data(SYMBOL, PERIOD, INTERVAL)

    # ── 2. Run backtest & wave detection ───────────────────────
    result = run_backtest(df)
    trades = result["trades"]
    zigzag = result["zigzag"]
    impulse_waves = result["impulse_waves"]
    corrective_waves = result["corrective_waves"]

    print(f"Impulse waves detected: {len(impulse_waves)}")
    print(f"Corrective waves detected: {len(corrective_waves)}")
    print(f"Trades executed: {len(trades)}")

    # ── 3. Build equity curve ──────────────────────────────────
    eq_dates, eq_values = build_equity_curve(trades, df)
    drawdown = compute_drawdown(eq_dates, eq_values)

    # ── 4. Create Figure ───────────────────────────────────────
    plt.style.use("dark_background")
    plt.rcParams.update({
        "font.family": "Consolas",
        "font.size": 9,
        "axes.labelsize": 10,
        "axes.titlesize": 12,
        "xtick.labelsize": 8,
        "ytick.labelsize": 8,
        "legend.fontsize": 8,
        "figure.facecolor": BG_COLOR,
        "axes.facecolor": PANEL_BG,
        "axes.edgecolor": GRID_COLOR,
        "grid.color": GRID_COLOR,
        "grid.alpha": 0.5,
        "text.color": TEXT_COLOR,
        "axes.labelcolor": TEXT_COLOR,
        "xtick.color": TEXT_DIM,
        "ytick.color": TEXT_DIM,
    })

    fig = plt.figure(figsize=(20, 13), facecolor=BG_COLOR)
    gs = GridSpec(4, 1, figure=fig, height_ratios=[5, 1.2, 1.8, 1.0],
                  hspace=0.08, left=0.06, right=0.96, top=0.93, bottom=0.05)

    ax_price = fig.add_subplot(gs[0])
    ax_vol = fig.add_subplot(gs[1], sharex=ax_price)
    ax_equity = fig.add_subplot(gs[2])
    ax_dd = fig.add_subplot(gs[3], sharex=ax_equity)

    # ════════════════════════════════════════════════════════════
    # PANEL 1: OHLC Price Chart with Elliott Wave Labels
    # ════════════════════════════════════════════════════════════
    draw_ohlc_bars(ax_price, df)

    # Zigzag overlay
    if zigzag:
        zz_dates = [df.index[p[0]] for p in zigzag]
        zz_prices = [p[1] for p in zigzag]
        ax_price.plot(zz_dates, zz_prices, color=ZIGZAG_COLOR,
                      linewidth=1.2, alpha=0.7, label="ZigZag", zorder=3)

    # Impulse wave labels (0-1-2-3-4-5)
    for w in impulse_waves:
        pts = w.points
        wave_dates = [df.index[p[0]] for p in pts]
        wave_prices = [p[1] for p in pts]
        # Draw impulse wave line
        ax_price.plot(wave_dates, wave_prices, color=IMPULSE_COLOR,
                      linewidth=2.0, alpha=0.8, zorder=4)
        labels = ["0", "1", "2", "3", "4", "5"]
        for j, (idx, price) in enumerate(pts):
            va = "bottom" if j % 2 == 1 else "top"
            offset_y = 80 if j % 2 == 1 else -80
            ax_price.annotate(
                labels[j],
                xy=(df.index[idx], price),
                xytext=(0, offset_y),
                textcoords="offset points",
                fontsize=11, fontweight="bold", color=IMPULSE_COLOR,
                ha="center", va=va,
                bbox=dict(boxstyle="round,pad=0.2", facecolor=BG_COLOR,
                          edgecolor=IMPULSE_COLOR, alpha=0.85, linewidth=0.8),
                arrowprops=dict(arrowstyle="-", color=IMPULSE_COLOR,
                                alpha=0.4, linewidth=0.6),
                zorder=6,
            )

    # Corrective wave labels (A-B-C)
    for w in corrective_waves:
        pts = w.points
        # Draw line through A, B_low, B_high, C
        wave_dates = [df.index[p[0]] for p in pts]
        wave_prices = [p[1] for p in pts]
        ax_price.plot(wave_dates, wave_prices, color=CORRECT_COLOR,
                      linewidth=2.0, alpha=0.8, linestyle="--", zorder=4)
        labels = ["A", "B", "C"]
        display_pts = [(pts[0][0], pts[0][1]),   # A start (high)
                       (pts[2][0], pts[2][1]),   # B end (bounce high)
                       (pts[3][0], pts[3][1])]   # C end (low)
        for j, (idx, price) in enumerate(display_pts):
            va = "top" if j in [0] else ("bottom" if j == 1 else "top")
            offset_y = -80 if va == "top" else 80
            ax_price.annotate(
                labels[j],
                xy=(df.index[idx], price),
                xytext=(0, offset_y),
                textcoords="offset points",
                fontsize=11, fontweight="bold", color=CORRECT_COLOR,
                ha="center", va=va,
                bbox=dict(boxstyle="round,pad=0.2", facecolor=BG_COLOR,
                          edgecolor=CORRECT_COLOR, alpha=0.85, linewidth=0.8),
                arrowprops=dict(arrowstyle="-", color=CORRECT_COLOR,
                                alpha=0.4, linewidth=0.6),
                zorder=6,
            )

    # Trade entry/exit markers
    for t in trades:
        # Entry marker
        if t.direction == "long":
            ax_price.scatter(t.entry_date, t.entry_price, marker="^",
                             color=LONG_MARKER, s=120, zorder=7,
                             edgecolors="white", linewidths=0.6)
        else:
            ax_price.scatter(t.entry_date, t.entry_price, marker="v",
                             color=SHORT_MARKER, s=120, zorder=7,
                             edgecolors="white", linewidths=0.6)
        # Exit marker
        if t.exit_date is not None:
            exit_color = WIN_EXIT if t.pnl > 0 else LOSS_EXIT
            ax_price.scatter(t.exit_date, t.exit_price, marker="X",
                             color=exit_color, s=100, zorder=7,
                             edgecolors="white", linewidths=0.5)

    ax_price.set_ylabel("Price (TWD)", fontweight="bold")
    ax_price.yaxis.set_major_formatter(mticker.FuncFormatter(
        lambda x, _: f"{x:,.0f}"))
    ax_price.grid(True, alpha=0.3, linewidth=0.5)
    ax_price.tick_params(axis="x", labelbottom=False)

    # Custom legend
    from matplotlib.lines import Line2D
    legend_elements = [
        Line2D([0], [0], color=UP_CANDLE, linewidth=2, label="Up Bar"),
        Line2D([0], [0], color=DOWN_CANDLE, linewidth=2, label="Down Bar"),
        Line2D([0], [0], color=ZIGZAG_COLOR, linewidth=1.2, label="ZigZag"),
        Line2D([0], [0], color=IMPULSE_COLOR, linewidth=2, label="Impulse (1-5)"),
        Line2D([0], [0], color=CORRECT_COLOR, linewidth=2, linestyle="--",
               label="Corrective (ABC)"),
        Line2D([0], [0], marker="^", color=LONG_MARKER, linestyle="None",
               markersize=8, label="Long Entry"),
        Line2D([0], [0], marker="v", color=SHORT_MARKER, linestyle="None",
               markersize=8, label="Short Entry"),
        Line2D([0], [0], marker="X", color=WIN_EXIT, linestyle="None",
               markersize=8, label="Win Exit"),
        Line2D([0], [0], marker="X", color=LOSS_EXIT, linestyle="None",
               markersize=8, label="Loss Exit"),
    ]
    ax_price.legend(handles=legend_elements, loc="upper left",
                    framealpha=0.85, facecolor=PANEL_BG, edgecolor=GRID_COLOR,
                    ncol=3, columnspacing=1.0)

    # ════════════════════════════════════════════════════════════
    # PANEL 2: Volume Bars
    # ════════════════════════════════════════════════════════════
    closes = df["Close"].values
    opens = df["Open"].values
    volumes = df["Volume"].values
    colors_vol = [UP_VOL if closes[i] >= opens[i] else DOWN_VOL
                  for i in range(len(df))]

    dates_num = mdates.date2num(df.index.to_pydatetime())
    ax_vol.bar(dates_num, volumes, width=0.6, color=colors_vol, alpha=0.7,
               edgecolor="none")
    ax_vol.set_ylabel("Volume", fontweight="bold")
    ax_vol.yaxis.set_major_formatter(mticker.FuncFormatter(
        lambda x, _: f"{x/1e9:.1f}B" if x >= 1e9 else (f"{x/1e6:.0f}M" if x >= 1e6 else f"{x:,.0f}")))
    ax_vol.grid(True, alpha=0.3, linewidth=0.5)
    ax_vol.tick_params(axis="x", labelbottom=False)
    # Dim the volume panel slightly
    ax_vol.set_ylim(0, max(volumes) * 1.3)

    # Format shared x-axis for price/volume panels
    ax_vol.xaxis.set_major_formatter(mdates.DateFormatter("%b %Y"))
    ax_vol.xaxis.set_major_locator(mdates.MonthLocator())

    # ════════════════════════════════════════════════════════════
    # PANEL 3: Equity Curve
    # ════════════════════════════════════════════════════════════
    if len(eq_dates) > 1:
        ax_equity.plot(eq_dates, eq_values, color=EQUITY_COLOR,
                       linewidth=1.8, zorder=3, label="Portfolio Equity")
        ax_equity.fill_between(eq_dates, INITIAL_CAPITAL, eq_values,
                               where=[v >= INITIAL_CAPITAL for v in eq_values],
                               color=ACCENT_GREEN, alpha=0.08, interpolate=True)
        ax_equity.fill_between(eq_dates, INITIAL_CAPITAL, eq_values,
                               where=[v < INITIAL_CAPITAL for v in eq_values],
                               color=ACCENT_RED, alpha=0.08, interpolate=True)
        ax_equity.axhline(INITIAL_CAPITAL, color=TEXT_DIM, linestyle="--",
                          linewidth=0.6, alpha=0.6)

        # Annotate final equity
        final_eq = eq_values[-1]
        total_return = (final_eq / INITIAL_CAPITAL - 1) * 100
        ax_equity.annotate(
            f"  {final_eq:,.0f} ({total_return:+.1f}%)",
            xy=(eq_dates[-1], final_eq),
            fontsize=9, fontweight="bold",
            color=ACCENT_GREEN if total_return >= 0 else ACCENT_RED,
            va="center",
        )

    ax_equity.set_ylabel("Equity (TWD)", fontweight="bold")
    ax_equity.yaxis.set_major_formatter(mticker.FuncFormatter(
        lambda x, _: f"{x/1e6:.2f}M" if x >= 1e6 else f"{x:,.0f}"))
    ax_equity.grid(True, alpha=0.3, linewidth=0.5)
    ax_equity.tick_params(axis="x", labelbottom=False)
    ax_equity.legend(loc="upper left", framealpha=0.85, facecolor=PANEL_BG,
                     edgecolor=GRID_COLOR)

    # ════════════════════════════════════════════════════════════
    # PANEL 4: Drawdown
    # ════════════════════════════════════════════════════════════
    if len(eq_dates) > 1:
        ax_dd.fill_between(eq_dates, 0, -drawdown, color=DRAWDOWN_COLOR,
                           alpha=0.4, step="post")
        ax_dd.plot(eq_dates, -drawdown, color=DRAWDOWN_COLOR,
                   linewidth=0.8, alpha=0.8)
        max_dd = max(drawdown)
        ax_dd.axhline(0, color=TEXT_DIM, linewidth=0.5, alpha=0.4)
        ax_dd.annotate(
            f"  Max DD: {max_dd:.1f}%",
            xy=(eq_dates[np.argmax(drawdown)], -max_dd),
            fontsize=8, fontweight="bold", color=DRAWDOWN_COLOR,
            va="top",
        )

    ax_dd.set_ylabel("Drawdown %", fontweight="bold")
    ax_dd.set_xlabel("Date", fontweight="bold")
    ax_dd.grid(True, alpha=0.3, linewidth=0.5)
    ax_dd.xaxis.set_major_formatter(mdates.DateFormatter("%b %Y"))
    ax_dd.xaxis.set_major_locator(mdates.MonthLocator(interval=2))
    plt.setp(ax_dd.xaxis.get_majorticklabels(), rotation=0, ha="center")

    # ── Title Bar ──────────────────────────────────────────────
    win_count = sum(1 for t in trades if t.pnl > 0)
    total_count = len(trades)
    win_rate = (win_count / total_count * 100) if total_count > 0 else 0
    total_pnl = sum(t.pnl_pct for t in trades)

    title_text = (
        f"ELLIOTT WAVE BACKTEST  |  {SYMBOL}  |  "
        f"{df.index[0].strftime('%Y-%m-%d')} to {df.index[-1].strftime('%Y-%m-%d')}  |  "
        f"Trades: {total_count}  |  Win Rate: {win_rate:.0f}%  |  "
        f"Return: {total_pnl:+.1f}%"
    )
    fig.suptitle(title_text, fontsize=13, fontweight="bold",
                 color=ACCENT_ORANGE, y=0.97,
                 fontfamily="Consolas")

    # Subtitle
    fig.text(0.5, 0.945,
             "Taiwan Weighted Index  |  Daily  |  Impulse (1-5) & Corrective (A-B-C) Wave Detection",
             ha="center", fontsize=9, color=TEXT_DIM, fontfamily="Consolas")

    # ── Save ───────────────────────────────────────────────────
    output_path = "E:/Developer/lufftw/repo/stock/chart_matplotlib.png"
    fig.savefig(output_path, dpi=200, facecolor=BG_COLOR,
                edgecolor="none", bbox_inches="tight")
    plt.close(fig)
    print(f"\nChart saved to: {output_path}")


if __name__ == "__main__":
    main()