"""Pre-registered grid / fractional Kelly research, issue #439.

Run with finlab>=2.3.8, numpy and pandas. Uses your existing FinLab login.
All 63 settings were fixed before results; no setting is selected by performance.
Outputs daily equity data, parameter metrics and native FinLab HTML reports.
The grid is a dividend-adjusted, fixed-anchor close-to-close allocation rule;
it does not model intraday limit orders. Uninvested cash earns zero.
"""
import json
from datetime import date
from itertools import product
from pathlib import Path

import finlab
import numpy as np
import pandas as pd
from finlab import data
from finlab.backtest import sim

END = "2026-11-06"
OUT = Path(__file__).resolve().parent
SLUGS = ("grid-trading-0050-backtest", "grid-trading-qqq-backtest", "kelly-criterion-0050-backtest")
GRIDS = list(product((0.03, 0.05, 0.07), (6, 8, 10), (0.25, 0.5, 0.75)))
KELLY = list(product((0.25, 0.5, 1.0), (126, 252, 504)))


def grid_position(price, gap, count, initial):
    levels = np.trunc(np.log(price / price.iloc[0]) / np.log1p(gap))
    return (initial - levels / count).clip(0, 1).to_frame()


def leverage_split(leverage):
    """Split target exposure between 0050 and the 2x ETF 00631L."""
    extra = np.maximum(leverage - 1, 0)
    return {"0050": np.minimum(leverage, 1) - extra, "00631L": extra}


def kelly_position(prices, fraction, lookback, start):
    returns = prices["0050"].pct_change()
    excess_mean = returns.rolling(lookback).mean() - (1.01 ** (1 / 252) - 1)
    leverage = (fraction * excess_mean / returns.rolling(lookback).var()).clip(0, 2)
    # The last observed session in a completed calendar month generates a signal.
    month_ends = leverage.groupby(leverage.index.to_period("M")).tail(1).index
    month_ends = month_ends[month_ends.to_period("M") < prices.index[-1].to_period("M")]
    monthly = leverage.loc[month_ends].dropna()
    held = monthly.reindex(prices.index, method="ffill").loc[start:END]
    return pd.DataFrame(leverage_split(held)).fillna(0)


def stats(returns):
    returns = returns.dropna()
    equity = (1 + returns).cumprod()
    deviation = returns.std(ddof=1)
    return {"cagr": float(equity.iloc[-1] ** (252 / len(returns)) - 1),
            "sharpe": float(returns.mean() / deviation * np.sqrt(252)) if deviation > 0 else None,
            "mdd": float((equity / equity.cummax().clip(lower=1) - 1).min()),
            "volatility": float(deviation * np.sqrt(252))}


def run(position, cost=None, resample=None):
    kwargs = {} if cost is None else {"fee_ratio": cost, "tax_ratio": 0}
    return sim(position, resample=resample, upload=False, end_date=END, **kwargs)


def daily_returns(report):
    return pd.Series(report.creturn).loc[:END].pct_change().fillna(0)


def aligned(returns, index):
    """Align daily returns to the strategy dates, with no return on the first day."""
    returns = returns.reindex(index).fillna(0)
    returns.iloc[0] = 0
    return returns


def actual_exposure(position, prices, returns):
    """Mark target holdings forward between next-close execution dates.

    Wealth uses sim's net curve so costs are reflected in the denominator.
    On a rebalance close target allocations are reset; no reset on other days.
    The average uses the weights entering each return observation.
    """
    weights = position.reindex(prices.index, method="ffill").fillna(0)
    events = weights.ne(weights.shift()).any(axis=1).shift(1, fill_value=False).to_numpy()
    targets = weights.to_numpy(dtype=float)
    asset_returns = prices.pct_change().fillna(0).to_numpy()
    wealth_returns = returns.reindex(prices.index, fill_value=0).to_numpy()
    scale = np.array([2 if c == "00631L" else 1 for c in position.columns])
    exposure = []
    live = np.zeros(len(position.columns))
    for i in range(len(prices)):
        exposure.append(float(live @ scale))
        live = live * (1 + asset_returns[i]) / (1 + wealth_returns[i])
        if events[i]:
            live = targets[i - 1]
    return pd.Series(exposure, index=prices.index).reindex(returns.index)


def higher(left, right):
    return left is not None and right is not None and left > right


def evaluate(slug, market, prices, positions, original, bil=None, spy=None):
    out = OUT.parent / slug
    out.mkdir(parents=True, exist_ok=True)
    data.set_market(market)
    reference = prices[prices.columns[0]].pct_change()
    rows = []
    for params, position in positions:
        report = run(position)
        ret = daily_returns(report)
        benchmark = aligned(reference, ret.index)
        row = {"params": list(params), **stats(ret), "benchmark_sharpe": stats(benchmark)["sharpe"]}
        row["wins"] = row["sharpe"] > row["benchmark_sharpe"]
        rows.append(row)
        print(slug, params, round(row["sharpe"], 4), flush=True)
        if params == original:
            base_report, base_ret, base_position, base_benchmark = report, ret, position, benchmark
    benchmark = base_benchmark
    exposure = actual_exposure(base_position, prices[base_position.columns], base_ret)
    average = float(exposure.mean())
    if "00631L" in base_position:
        fixed = pd.DataFrame(leverage_split(average), index=base_position.index)
    else:
        fixed = pd.DataFrame({base_position.columns[0]: average}, index=base_position.index)
    matched_report = run(fixed, resample="D")
    matched = daily_returns(matched_report).reindex(base_ret.index).fillna(0)
    mid = len(base_ret) // 2
    halves = [{"strategy": stats(base_ret.iloc[sl]), "benchmark": stats(benchmark.iloc[sl]),
               "start": str(base_ret.iloc[sl].index[0].date()),
               "end": str(base_ret.iloc[sl].index[-1].date())}
              for sl in (slice(None, mid), slice(mid, None))]
    year_returns = (1 + base_ret).groupby(base_ret.index.year).prod() - 1
    best_year = int(year_returns.idxmax())
    keep = base_ret.index.year != best_year
    monthly_turnover = base_position.diff().abs().sum(axis=1).resample("ME").sum()
    metrics = {"run_date": str(date.today()), "finlab_version": finlab.__version__,
               "start": str(base_ret.index[0].date()), "end": str(base_ret.index[-1].date()),
               "original": list(original), "strategy": stats(base_ret), "benchmark": stats(benchmark),
               "neighbors": rows, "wins": sum(r["wins"] for r in rows), "halves": halves,
               "average_exposure": average, "fixed_exposure": stats(matched),
               "best_year": best_year, "without_best_year": stats(base_ret.loc[keep]),
               "benchmark_without_best_year": stats(benchmark.loc[keep]),
               "average_monthly_target_turnover": float(monthly_turnover.mean()),
               "max_monthly_target_turnover": float(monthly_turnover.max())}
    frame = pd.DataFrame({"strategy": (1 + base_ret).cumprod(), "benchmark": (1 + benchmark).cumprod(),
                          "fixed_exposure": (1 + matched).cumprod(), "exposure": exposure})
    if bil is not None:
        cash = aligned(bil.pct_change(), base_ret.index)
        stressed = daily_returns(run(base_position, cost=0.001)).reindex(base_ret.index).fillna(0)
        metrics.update({"cost_001": stats(stressed), "excess_bil_sharpe": stats(base_ret - cash)["sharpe"],
                        "benchmark_excess_bil_sharpe": stats(benchmark - cash)["sharpe"],
                        "bil_first_date": str(bil.first_valid_index().date())})
        frame["cost_001"] = (1 + stressed).cumprod()
    if spy is not None:
        spy_returns = aligned(spy.pct_change(), base_ret.index)
        metrics["spy"] = stats(spy_returns)
        frame["SPY"] = (1 + spy_returns).cumprod()
    metrics["passed"] = (metrics["wins"] / len(rows) >= 0.7
        and all(higher(h["strategy"]["sharpe"], h["benchmark"]["sharpe"]) for h in halves)
        and metrics["strategy"]["sharpe"] > metrics["benchmark"]["sharpe"]
        and metrics["strategy"]["sharpe"] > metrics["fixed_exposure"]["sharpe"]
        and metrics["without_best_year"]["sharpe"] > metrics["benchmark_without_best_year"]["sharpe"]
        and (bil is None or (metrics["cost_001"]["sharpe"] > metrics["benchmark"]["sharpe"]
             and metrics["excess_bil_sharpe"] > metrics["benchmark_excess_bil_sharpe"])))
    (out / "metrics.json").write_text(json.dumps(metrics, ensure_ascii=False, indent=2, allow_nan=False))
    frame.to_csv(out / "data.csv", index_label="date")
    base_position.to_csv(out / "positions.csv", index_label="date")
    base_report.to_html(str(out / "report_strategy.html"))
    matched_report.to_html(str(out / "report_fixed.html"))
    print('RESULT', slug, json.dumps({k: v for k, v in metrics.items() if k != 'neighbors'}), flush=True)


def main():
    tw = pd.DataFrame(data.get("etl:adj_close"))[["0050", "00631L"]].loc[:END]
    us = pd.DataFrame(data.get("us_fund_price:adj_close"))[["QQQ", "BIL", "SPY"]].loc[:END]
    qqq = us["QQQ"].loc["2015-07-01":].dropna()
    us_extras = {"bil": us["BIL"].dropna(), "spy": us["SPY"].dropna()}
    for slug, market, price, extras in ((SLUGS[0], "tw", tw["0050"].dropna(), {}),
                                        (SLUGS[1], "us_fund", qqq, us_extras)):
        positions = [(p, grid_position(price, *p)) for p in GRIDS]
        evaluate(slug, market, price.to_frame(), positions, (0.05, 8, 0.5), **extras)
    common = tw.dropna()
    longest = max(lookback for _, lookback in KELLY)
    ready = tw["0050"].pct_change().rolling(longest).mean().first_valid_index()
    start = max(common.index[0], ready)
    positions = [(p, kelly_position(tw, *p, start)) for p in KELLY]
    evaluate(SLUGS[2], "tw", common.loc[start:], positions, (0.5, 252))


if __name__ == "__main__":
    main()