"""Pre-registered trend-indicator experiments.

Run the QQQ grid:
  uv run --with finlab --with numba --with ta-lib python research.py --market us_fund --out qqq
Run stocks with --market us --out stocks, or Taiwan with --market tw --out taiwan.
No reports are uploaded.
"""
import argparse
import itertools
import json
from datetime import date
from pathlib import Path

import finlab
import numpy as np
import pandas as pd
import talib
from finlab import FinlabDataFrame, data
from finlab.backtest import sim
from numba import njit

END = "2026-11-09"
GRIDS = {"ichimoku": ([7, 9, 11], [22, 26, 30]),
         "psar": ([0.01, 0.02, 0.03], [0.1, 0.2, 0.3]),
         "supertrend": ([7, 10, 14], [2, 3, 4]),
         "darvas": ([40, 60, 80], [2, 3, 5]),
         "heikin": ([2, 3, 4], [1, 2, 3])}


@njit
def recursive(h, l, c, auxiliary, kind, parameter):
    """Columnwise causal Supertrend or confirmed Darvas boxes."""
    state = np.zeros(c.shape, dtype=np.bool_)
    for j in range(c.shape[1]):
        upper, lower, pivot, bottom = np.nan, np.nan, np.nan, np.nan
        box_top, box_bottom, age, bullish = np.nan, np.nan, 0, False
        for i in range(1, c.shape[0]):
            if not np.isfinite(c[i, j]) or not np.isfinite(auxiliary[i, j]):
                continue
            if kind == 0:
                up = (h[i, j] + l[i, j]) / 2 + parameter * auxiliary[i, j]
                down = (h[i, j] + l[i, j]) / 2 - parameter * auxiliary[i, j]
                if not np.isfinite(upper):
                    upper, lower = up, down
                else:
                    up = up if up < upper or c[i-1, j] > upper else upper
                    down = down if down > lower or c[i-1, j] < lower else lower
                    bullish = c[i, j] >= down if bullish else c[i, j] > up
                    upper, lower = up, down
            else:
                if h[i, j] > auxiliary[i, j]:
                    pivot, bottom, age = h[i, j], l[i, j], 0
                elif np.isfinite(pivot):
                    bottom = min(bottom, l[i, j])
                    age += 1
                    if age == parameter:
                        box_top, box_bottom = pivot, bottom
                if np.isfinite(box_top):
                    if c[i, j] < box_bottom:
                        bullish = False
                    elif c[i, j] > box_top:
                        bullish = True
            state[i, j] = bullish
    return state


def signal(ohlc, kind, a, b):
    o, h, l, c = ohlc
    if kind == "ichimoku":
        middle = lambda n: (h.rolling(n).max() + l.rolling(n).min()) / 2
        conversion, base = middle(a), middle(b)
        cloud = pd.DataFrame(np.maximum(((conversion + base) / 2).shift(b), middle(2*b).shift(b)), index=c.index, columns=c.columns)
        return (c > cloud) & (conversion > base)
    if kind == "psar":
        sar = pd.DataFrame(index=c.index, columns=c.columns, dtype=float)
        for column in c:
            valid = h[column].notna() & l[column].notna()
            sar.loc[valid, column] = talib.SAR(h.loc[valid, column].to_numpy(), l.loc[valid, column].to_numpy(), acceleration=a, maximum=b)
        return c > sar
    if kind == "heikin":
        ha_close = (o + h + l + c) / 4
        # y[t] = (y[t-1] + HAclose[t-1])/2; first bar uses (open+close)/2.
        ha_open = ha_close.shift().copy()
        for column in c:
            first = c[column].first_valid_index()
            if first is not None:
                ha_open.loc[first, column] = (o.loc[first, column] + c.loc[first, column]) / 2
        ha_open = ha_open.ewm(alpha=0.5, adjust=False).mean()
        valid = (ha_close > ha_open) & (l >= ha_open - 1e-10)
        entry = valid.rolling(a).sum() == a
        exit_ = (~valid).rolling(b).sum() == b
        return pd.DataFrame(FinlabDataFrame(entry).hold_until(FinlabDataFrame(exit_)))
    if kind == "supertrend":
        tr = pd.DataFrame(np.maximum(h-l, np.maximum((h-c.shift()).abs(), (l-c.shift()).abs())), index=c.index, columns=c.columns)
        # Wilder smoothing, seeded by the arithmetic mean of the first n true ranges.
        atr = tr.copy() * np.nan
        for column in c:
            valid = tr[column].dropna()
            if len(valid) >= a:
                seed = valid.iloc[a-1:].copy()
                seed.iloc[0] = valid.iloc[:a].mean()
                atr.loc[seed.index, column] = seed.ewm(alpha=1/a, adjust=False).mean()
        auxiliary, mode = atr, 0
    else:
        auxiliary, mode = h.rolling(a).max().shift(), 1
    return pd.DataFrame(recursive(h.to_numpy(), l.to_numpy(), c.to_numpy(), auxiliary.to_numpy(), mode, b), index=c.index, columns=c.columns)


def load(market):
    data.set_market(market)
    prefix = {"tw": "etl:", "us": "us_price:", "us_fund": "us_fund_price:"}[market]
    warmup = "2007" if market == "tw" else "2016"
    frames = [pd.DataFrame(data.get(prefix + "adj_" + field).loc[warmup:END]).astype(float) for field in ("open", "high", "low", "close")]
    benchmark = frames[-1]["0050"] if market == "tw" else data.get("us_fund_price:adj_close")["QQQ"].loc[warmup:END]
    calendar = benchmark.dropna().index
    frames = [frame.reindex(calendar) for frame in frames]
    c = frames[-1]
    if market == "us_fund":
        columns, eligible = ["QQQ"], c.notna()
    elif market == "tw":
        columns = [s for s in c if s.isdigit() and len(s) == 4 and not s.startswith(("0", "91"))]
        amount = pd.DataFrame(data.get("price:成交金額")).reindex(index=c.index, columns=columns)
        eligible = amount.rolling(20).mean() > 20_000_000
    else:
        common = data.get("etl:us_common_stock_filter")
        columns = list(c.columns.intersection(common.columns))
        volume = pd.DataFrame(data.get("us_price:volume")).reindex(index=c.index, columns=columns)
        amount = (c[columns] * volume).rolling(60, min_periods=20).mean()
        extreme = c[columns].pct_change(fill_method=None).abs().ge(0.5).rolling(252, min_periods=1).max().shift().fillna(0).astype(bool)
        eligible = amount.rank(axis=1, method="first", ascending=False).le(100) & amount.gt(5_000_000) & c[columns].gt(5) & ~extreme
        columns = list(eligible.columns[eligible.any()])
    return [frame.reindex(index=c.index, columns=columns) for frame in frames], eligible.reindex(index=c.index, columns=columns).fillna(False)


def sharpe(returns):
    return float(returns.mean() / returns.std() * np.sqrt(252)) if returns.std() > 0 else 0.0


def metrics(returns):
    equity = (1 + returns).cumprod()
    years = (returns.index[-1] - returns.index[0]).days / 365.25
    return {"cagr": float(equity.iloc[-1] ** (1 / years) - 1), "sharpe": sharpe(returns), "mdd": float((equity / equity.cummax() - 1).min())}


def backtest(position, fee=None):
    costs = {} if fee is None else {"fee_ratio": fee}
    return sim(FinlabDataFrame(position), resample=None, upload=False, **costs)


def daily_returns(report, calendar):
    return report.creturn.reindex(calendar).ffill().fillna(1).pct_change().fillna(0)


def run(market, out):
    start = "2008-01-02" if market == "tw" else "2017-01-03"
    fund = pd.DataFrame(data.get("etl:adj_close" if market == "tw" else "us_fund_price:adj_close").loc[start:END])
    benchmark = fund["0050" if market == "tw" else "QQQ"].dropna()
    calendar = benchmark.index
    benchmark_returns = benchmark.pct_change().fillna(0)
    cash = benchmark_returns * 0 if market == "tw" else fund["BIL"].reindex(calendar).ffill().pct_change().fillna(0)
    prices, eligible = load(market)
    if market == "tw":
        etf = [pd.DataFrame(data.get("etl:adj_" + field).loc["2007":END, ["0050"]]).reindex(prices[-1].index) for field in ("open", "high", "low", "close")]
    elif market == "us":
        etf = None
    else:
        etf = prices
    results, curves = {}, {"benchmark": (1 + benchmark_returns).cumprod(), "BIL": (1 + cash).cumprod()}
    stock_eligible = eligible.reindex(calendar).fillna(False)
    halves = (slice(None, len(calendar) // 2), slice(len(calendar) // 2, None))
    variants = {"us": ["stocks"], "tw": ["timing", "avoid", "stocks"], "us_fund": ["timing", "avoid"]}[market]
    for kind in (list(GRIDS) if market != "tw" else ["ichimoku", "psar"]):
        aa, bb = GRIDS[kind]
        for variant in variants:
            name, rows = kind + "_" + variant, []
            ohlc = prices if variant == "stocks" else etf
            for a, b in itertools.product(aa, bb):
                state = signal(ohlc, kind, a, b).reindex(calendar).fillna(False)
                state = ~state if variant == "avoid" else state
                position = state & stock_eligible if variant == "stocks" else state
                report = backtest(position)
                returns = daily_returns(report, calendar)
                row = {"a": a, "b": b, **metrics(returns)}
                rows.append(row)
                if a == aa[1] and b == bb[1]:
                    center_returns, center_report, center_position = returns, report, position
                    weights = position.div(position.sum(axis=1).replace(0, np.nan), axis=0).fillna(0)
                    monthly_turnover = weights.diff().abs().sum(axis=1).resample("ME").sum() / 2
                    central = {
                        **row,
                        "exposure": float(position.any(axis=1).shift().fillna(False).mean()),
                        "holdings_mean": float(position.sum(axis=1).mean()),
                        "target_turnover_monthly_mean": float(monthly_turnover.mean()),
                        "target_turnover_monthly_max": float(monthly_turnover.max()),
                        "trades": int(len(report.get_trades())),
                    }
            central["halves"] = [sharpe(center_returns.iloc[s]) for s in halves]
            central["benchmark_halves"] = [sharpe(benchmark_returns.iloc[s]) for s in halves]
            central["benchmark"] = metrics(benchmark_returns)
            central["fixed_exposure_sharpe"] = sharpe(benchmark_returns * central["exposure"])
            central["excess_sharpe"] = sharpe(center_returns - cash)
            central["benchmark_excess_sharpe"] = sharpe(benchmark_returns - cash)
            best_year = int(center_returns.groupby(center_returns.index.year).apply(lambda r: (1+r).prod()-1).idxmax())
            keep = calendar.year != best_year
            central["best_year"] = best_year
            central["without_best_year"] = [sharpe(center_returns[keep]), sharpe(benchmark_returns[keep])]
            central["winning_grid"] = sum(r["sharpe"] > central["benchmark"]["sharpe"] for r in rows)
            central["grid"] = rows
            checks = [
                central["winning_grid"] >= 7,
                all(a > b for a, b in zip(central["halves"], central["benchmark_halves"])),
                central["sharpe"] > central["fixed_exposure_sharpe"],
                central["excess_sharpe"] > central["benchmark_excess_sharpe"],
                central["without_best_year"][0] > central["without_best_year"][1],
            ]
            central["passed"] = all(checks)
            curves[name] = (1 + center_returns).cumprod()
            if market != "tw":
                stress_returns = daily_returns(backtest(center_position, fee=0.001), calendar)
                central["cost_0_1pct"] = metrics(stress_returns)
                central["cost_0_1pct"]["excess_sharpe"] = sharpe(stress_returns - cash)
                central["passed"] &= (central["cost_0_1pct"]["sharpe"] > central["benchmark"]["sharpe"]
                                      and central["cost_0_1pct"]["excess_sharpe"] > central["benchmark_excess_sharpe"])
            center_report.to_html(str(out / (name + "_report.html")))
            results[name] = central
            print(market, name, "grid", central["winning_grid"], "sharpe", round(central["sharpe"], 3), "passed", central["passed"], flush=True)
    pd.DataFrame(curves).to_csv(out / "data.csv", index_label="date")
    payload = {"market": market, "start": start, "end": END, "run_date": str(date.today()), "finlab_version": finlab.__version__, "results": results}
    (out / "metrics.json").write_text(json.dumps(payload, indent=2, allow_nan=False))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--market", choices=["tw", "us", "us_fund"], required=True)
    parser.add_argument("--out", type=Path, required=True)
    args = parser.parse_args()
    args.out.mkdir(parents=True, exist_ok=True)
    finlab.login()
    run(args.market, args.out)