#!/usr/bin/env python3
"""Turnover / capacity optimization for crypto XS momentum (the one validated edge).

Daily-rebalancing a 20-day signal pays ~5x wasted turnover. Test rebalance frequency,
signal smoothing, and universe breadth (top-K) against the same square-root slippage model.
Trade-off: weekly/top-20 cuts cost but top-20 cuts breadth (momentum needs dispersion).
Report gross Sharpe, avg daily turnover, and net Sharpe-vs-AUM per config.
"""
import math
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import pit_sweep as ps  # noqa: E402
from signal_sweep import xs_weights, sharpe_t  # noqa: E402
import torch  # noqa: E402

DEV = "cuda" if torch.cuda.is_available() else "cpu"


def roll(fn, X, L):
    out = np.full_like(X, np.nan)
    for t in range(L, len(X)):
        out[t] = fn(X[t - L:t], axis=0)
    return out


def main():
    syms, days, close, qv, fund = ps.load()
    T, N = close.shape
    lc = np.log(close)
    R = np.zeros((T, N)); R[1:] = lc[1:] - lc[:-1]; R = np.where(np.isfinite(R), R, 0.0)
    fund = np.where(np.isfinite(fund), fund, 0.0)
    qv = np.nan_to_num(qv)
    dv30 = roll(np.mean, qv, 30)
    vol20 = np.nan_to_num(roll(np.std, R, 20))
    Reff = R - fund
    tradeable = np.ones((T, N), bool)
    for j in range(N):
        idx = np.where(np.isfinite(close[:, j]))[0]
        if len(idx):
            tradeable[max(0, idx[-1] - 4):idx[-1] + 1, j] = False
    Reff = np.where(tradeable, Reff, 0.0)
    raw20 = np.full_like(lc, np.nan); raw20[20:] = lc[20:] - lc[:-20]

    def univ_mask(TOPK):
        u = np.zeros((T, N), bool)
        for t in range(T):
            elig = np.where((dv30[t] > 0) & np.isfinite(close[t]))[0]
            if len(elig):
                u[t, elig[np.argsort(-dv30[t, elig])[:TOPK]]] = True
        return u

    def held_weights(TOPK, K, smooth):
        sig = raw20.copy(); sig[~univ_mask(TOPK)] = np.nan
        wt = xs_weights(sig)
        if smooth > 1:                                   # EWMA the target weights over time
            a = 2.0 / (smooth + 1)
            for t in range(1, T):
                wt[t] = a * wt[t] + (1 - a) * wt[t - 1]
        wh = wt.copy()                                   # step-rebalance every K days
        if K > 1:
            for t in range(T):
                wh[t] = wt[t - (t % K)]
        return wh

    def net_sharpe(wh, AUM, eta=1.0):
        turn = np.abs(wh[1:] - wh[:-1])
        adv = dv30[1:]; sg = vol20[1:]
        hs = np.clip(30.0 / np.sqrt(np.maximum(adv, 1.0) / 1e6), 1.0, 30.0) / 1e4
        part = np.where(adv > 0, turn * AUM / adv, 0.0)
        cost = np.sum(turn * (hs + eta * sg * np.sqrt(np.clip(part, 0, None))), axis=1)
        return sharpe_t(torch.tensor(np.sum(wh[:-1] * Reff[1:], axis=1) - cost, device=DEV, dtype=torch.float64))

    aums = [1e6, 5e6, 2e7, 5e7, 2e8, 5e8]
    print(f"\n===== TURNOVER / CAPACITY OPTIMIZATION (XS momentum-20, sqrt-impact η=1) =====")
    print(f"{'config':>26} {'grossSR':>7} {'turn%':>6} | " + " ".join(
        f"{('$'+(f'{a/1e6:.0f}M' if a<1e9 else f'{a/1e9:.1f}B')):>7}" for a in aums))
    cfgs = [
        ("top50 daily (baseline)", 50, 1, 1),
        ("top50 weekly", 50, 5, 1),
        ("top50 weekly+smooth5", 50, 5, 5),
        ("top30 weekly+smooth5", 30, 5, 5),
        ("top20 weekly+smooth5", 20, 5, 5),
        ("top20 biweekly+smooth10", 20, 10, 10),
        ("top30 biweekly+smooth10", 30, 10, 10),
    ]
    for nm, TOPK, K, sm in cfgs:
        wh = held_weights(TOPK, K, sm)
        gross = sharpe_t(torch.tensor(np.sum(wh[:-1] * Reff[1:], axis=1), device=DEV, dtype=torch.float64))
        turn = float(np.mean(np.sum(np.abs(wh[1:] - wh[:-1]), axis=1))) * 100
        row = [net_sharpe(wh, a) for a in aums]
        print(f"{nm:>26} {gross:>+7.2f} {turn:>5.1f}% | " + " ".join(f"{r:>+7.2f}" for r in row))
    print("\nVERDICT: best config's $20-50M net Sharpe = realistic deployable number at meaningful scale.")


if __name__ == "__main__":
    main()