"""
Step 1: Download kline data and save to disk.
Usage: python3 backtests/download_data.py --days 7 --top 35
Output: backtests/data_cache_7d.pkl
"""

import sys, os, time, pickle, argparse
import numpy as np
from datetime import datetime

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from pybit.unified_trading import HTTP

VWAP_PERIOD = 50
NATR_PERIOD = 14
CHOP_PERIOD = 14
BLACKLIST = {"BTCUSDT", "ETHUSDT", "USDCUSDT", "BTCPERP", "ETHPERP"}


def fetch_klines(session, symbol, interval="5", days=7):
    all_klines = []
    bars_needed = days * 24 * 60 // int(interval)
    end_time = int(datetime.now().timestamp() * 1000)
    while len(all_klines) < bars_needed:
        try:
            resp = session.get_kline(
                category="linear", symbol=symbol,
                interval=interval, limit=1000, end=end_time,
            )
            if resp["retCode"] != 0: break
            items = resp["result"]["list"]
            if not items: break
            for item in items:
                all_klines.append({
                    "ts": int(item[0]), "h": float(item[2]),
                    "l": float(item[3]), "c": float(item[4]),
                    "v": float(item[5]),
                })
            end_time = int(items[-1][0]) - 1
            if len(items) < 1000: break
        except Exception as e:
            print(f"  ERR {symbol}: {e}"); break
        time.sleep(0.15)
    all_klines.reverse()
    seen = set(); unique = []
    for k in all_klines:
        if k["ts"] not in seen: seen.add(k["ts"]); unique.append(k)
    return unique[-bars_needed:] if len(unique) > bars_needed else unique


def get_top_symbols(session, top_n=35, min_vol=20_000_000):
    tickers = session.get_tickers(category="linear")
    cands = []
    for t in tickers["result"]["list"]:
        sym = t["symbol"]
        if not sym.endswith("USDT") or sym in BLACKLIST: continue
        vol = float(t.get("turnover24h", 0))
        if vol >= min_vol: cands.append((sym, vol))
    cands.sort(key=lambda x: x[1], reverse=True)
    return cands[:top_n]


def calc_zvwap(h, l, c, v, period=VWAP_PERIOD):
    n = len(c); z = np.full(n, 0.0)
    for i in range(period, n):
        hi=h[i-period:i]; lo=l[i-period:i]; ci=c[i-period:i]; vi=v[i-period:i]
        tp=(hi+lo+ci)/3; ctv=np.cumsum(tp*vi); cv=np.cumsum(vi)
        cvs=np.where(cv==0,1,cv); va=ctv/cvs; vwap=va[-1]
        std=np.std(ci-va)
        if std>0: z[i]=(c[i]-vwap)/std
    return z


def calc_natr(h, l, c, period=NATR_PERIOD):
    n=len(c); natr=np.full(n,0.0); tr=np.zeros(n)
    for i in range(1,n):
        tr[i]=max(h[i]-l[i],abs(h[i]-c[i-1]),abs(l[i]-c[i-1]))
    for i in range(period,n):
        atr=np.mean(tr[i-period+1:i+1])
        natr[i]=(atr/c[i])*100 if c[i]>0 else 0
    return natr


def calc_chop(h, l, c, period=CHOP_PERIOD):
    n=len(c); chop=np.full(n,50.0); tr=np.zeros(n)
    for i in range(1,n):
        tr[i]=max(h[i]-l[i],abs(h[i]-c[i-1]),abs(l[i]-c[i-1]))
    for i in range(period,n):
        atr_sum=np.sum(tr[i-period+1:i+1])
        highest=np.max(h[i-period+1:i+1]); lowest=np.min(l[i-period+1:i+1])
        rng=highest-lowest
        if rng>0: chop[i]=100*np.log10(atr_sum/rng)/np.log10(period)
    return chop


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--days", type=int, default=7)
    parser.add_argument("--top", type=int, default=35)
    args = parser.parse_args()

    session = HTTP(testnet=False)

    print(f"Fetching top {args.top} symbols...")
    symbols = get_top_symbols(session, args.top)
    print(f"  Found {len(symbols)} symbols")

    print(f"Downloading {args.days}d klines...")
    data_cache = {}
    for idx, (sym, vol) in enumerate(symbols):
        klines = fetch_klines(session, sym, days=args.days)
        if len(klines) < VWAP_PERIOD + 50:
            print(f"  SKIP {sym} ({len(klines)} bars)")
            continue
        h=np.array([k["h"] for k in klines])
        l=np.array([k["l"] for k in klines])
        c=np.array([k["c"] for k in klines])
        v=np.array([k["v"] for k in klines])
        data_cache[sym] = {
            "h": h, "l": l, "c": c, "v": v,
            "z": calc_zvwap(h,l,c,v),
            "natr": calc_natr(h,l,c),
            "chop": calc_chop(h,l,c),
            "vol24h": vol,
        }
        print(f"  [{idx+1}/{len(symbols)}] {sym} — {len(klines)} bars, vol ${vol/1e6:.0f}M")
        time.sleep(0.2)

    out_path = os.path.join(os.path.dirname(__file__), f"data_cache_{args.days}d.pkl")
    with open(out_path, "wb") as f:
        pickle.dump({"symbols": data_cache, "days": args.days,
                      "date": datetime.now().isoformat()}, f)

    print(f"\n✅ Saved {len(data_cache)} symbols to {out_path}")
    print(f"   File size: {os.path.getsize(out_path)/1024:.0f} KB")


if __name__ == "__main__":
    main()