import yfinance as yf
import pandas as pd
import numpy as np
import mplfinance as mpf
import talib
import gradio as gr
import os
from datetime import date

# ---------------- TALib Patterns ----------------
TALIB_PATTERNS = sorted(p for p in dir(talib) if p.startswith("CDL"))
PATTERN_MAP = {p.replace("CDL", ""): p for p in TALIB_PATTERNS}
PATTERNS = ["None"] + list(PATTERN_MAP.keys())

# ---------------- Data Cleaning ----------------
def clean_ohlc(df):
    df = df.copy()

    if isinstance(df.columns, pd.MultiIndex):
        df.columns = df.columns.get_level_values(0)

    df.columns = [c.lower() for c in df.columns]
    df.index.name = None

    df = df.rename(columns={
        "open": "Open",
        "high": "High",
        "low": "Low",
        "close": "Close",
        "volume": "Volume"
    })

    df.index = pd.to_datetime(df.index, errors="coerce")

    for c in ["Open", "High", "Low", "Close", "Volume"]:
        if c in df.columns:
            df[c] = pd.to_numeric(df[c], errors="coerce")

    df = df.dropna(subset=["Open", "High", "Low", "Close"])
    df = df.sort_index()

    return df

# ---------------- Load Data ----------------
def load_data(symbol, start, end):
    if not symbol:
        return None, "❌ Symbol required"

    df = yf.download(symbol, start=start, end=end, progress=False)
    if df.empty:
        return None, "❌ No data found"

    df = clean_ohlc(df)
    return df, f"✅ Loaded {len(df)} rows"

# ---------------- Pattern Builder ----------------
def build_pattern(df, pattern):
    if pattern == "None":
        return []

    func = getattr(talib, PATTERN_MAP[pattern])
    res = func(df["Open"], df["High"], df["Low"], df["Close"])
    s = pd.Series(res, index=df.index)

    aps = []

    if (s > 0).any():
        bull = pd.Series(np.nan, index=df.index)
        bull[s > 0] = df["Low"][s > 0] * 0.98
        aps.append(mpf.make_addplot(bull, type="scatter",
                                    marker="^", color="green", markersize=90))

    if (s < 0).any():
        bear = pd.Series(np.nan, index=df.index)
        bear[s < 0] = df["High"][s < 0] * 1.02
        aps.append(mpf.make_addplot(bear, type="scatter",
                                    marker="v", color="red", markersize=90))

    return aps

# ---------------- Build Chart ----------------
def build_chart(df, symbol, pattern):
    if df is None:
        return None, "❌ Load data first"

    addplots = build_pattern(df, pattern)
    os.makedirs("/tmp", exist_ok=True)
    path = f"/tmp/{symbol}_{pattern}.png"

    fig, _ = mpf.plot(
        df,
        type="candle",
        volume="Volume" in df.columns,
        addplot=addplots if addplots else None,
        style="yahoo",
        title=f"{symbol} | {pattern}",
        figscale=1.8,
        returnfig=True
    )

    fig.savefig(path, dpi=150, bbox_inches="tight")
    return path, "✅ Chart rebuilt"

# ================= GRADIO APP =================
with gr.Blocks(fill_height=True) as app:
    gr.Markdown("## 📊 Candlestick Pattern Analyzer")

    state_df = gr.State(None)

    # -------- Load Section --------
    with gr.Row():
        symbol = gr.Textbox(label="Symbol", value="MSFT")
        start = gr.Textbox(label="Start Date", value="2024-01-01")
        end = gr.Textbox(label="End Date", value=date.today().strftime("%Y-%m-%d"))
        load_btn = gr.Button("Load Data", variant="primary")
        load_status = gr.Textbox(label="Load Status")

    # -------- Pattern + Chart --------
    with gr.Row():
        with gr.Column(scale=1):
            pattern = gr.Dropdown(PATTERNS, value="None", label="Pattern")
            build_btn = gr.Button("Build Chart")

        with gr.Column(scale=3):
            chart = gr.Image(height=600)
            build_status = gr.Textbox(label="Status")

    # -------- Events --------
    load_btn.click(
        load_data,
        inputs=[symbol, start, end],
        outputs=[state_df, load_status]
    )

    build_btn.click(
        build_chart,
        inputs=[state_df, symbol, pattern],
        outputs=[chart, build_status]
    )

    # 🔥 AUTO REBUILD WHEN PATTERN CHANGES
    pattern.change(
        build_chart,
        inputs=[state_df, symbol, pattern],
        outputs=[chart, build_status]
    )

if __name__ == "__main__":
    app.launch()