"""
Step 3: Sweep from cached data.
Usage:
  python3 backtests/sweep_from_cache.py --rr 2   (R:R = 1:2, TP=2×SL)
  python3 backtests/sweep_from_cache.py --rr 3   (R:R = 1:3)
  python3 backtests/sweep_from_cache.py --rr 4
  python3 backtests/sweep_from_cache.py --rr 5

Output: backtests/sweep_rr{N}_results.json
"""

import sys, os, json, pickle, argparse
import numpy as np
from datetime import datetime
from itertools import product

VWAP_PERIOD = 50
MAKER_FEE = 0.0002
TAKER_FEE = 0.00055

# Fixed filters per Rick's request
# Z > 1.8, NATR > 0.5, CHOP > 45

GRID = {
    "z_entry":  [1.9, 2.0, 2.2, 2.5, 2.8, 3.0],
    "z_max":    [3.0, 3.5, 4.0, 5.0],
    "natr_min": [0.5, 0.75, 1.0],
    "natr_max": [2.0, 2.5, 3.0, 3.5],
    "chop_min": [45, 50, 55],
}

# SL values to try — TP = SL × R:R
SL_VALUES = [0.3, 0.5, 0.75, 1.0, 1.5, 2.0]


def run_one(data_cache, cfg):
    all_trades = []
    for sym, d in data_cache.items():
        h=d["h"]; l=d["l"]; c=d["c"]; z_arr=d["z"]; natr_arr=d["natr"]; chop_arr=d["chop"]
        n=len(c); active=None; cooldown=0

        for i in range(VWAP_PERIOD, n):
            z=z_arr[i]; natr=natr_arr[i]; chop=chop_arr[i]

            if active:
                sl_hit=tp_hit=z_tp=False
                if active["side"]=="LONG":
                    if l[i]<=active["sl"]: sl_hit=True; cp=active["sl"]
                    if h[i]>=active["tp"]: tp_hit=True; cp2=active["tp"]
                    if z>=-0.3 and c[i]>active["ep"]: z_tp=True; cp3=c[i]
                else:
                    if h[i]>=active["sl"]: sl_hit=True; cp=active["sl"]
                    if l[i]<=active["tp"]: tp_hit=True; cp2=active["tp"]
                    if z<=0.3 and c[i]<active["ep"]: z_tp=True; cp3=c[i]

                if sl_hit:
                    all_trades.append({"sym":sym,"pnl":_pnl(active,cp,"SL"),"reason":"SL"})
                    active=None; cooldown=i+12
                elif tp_hit:
                    all_trades.append({"sym":sym,"pnl":_pnl(active,cp2,"TP"),"reason":"TP"})
                    active=None; cooldown=i+12
                elif z_tp:
                    all_trades.append({"sym":sym,"pnl":_pnl(active,cp3,"Z-TP"),"reason":"Z-TP"})
                    active=None; cooldown=i+12
                continue

            if i<cooldown: continue
            if natr<cfg["natr_min"] or natr>cfg["natr_max"]: continue
            if chop<cfg["chop_min"]: continue
            if abs(z)>cfg["z_max"]: continue

            if z<-cfg["z_entry"]:
                ep=c[i]; active={"side":"LONG","ep":ep,"qty":(5.0*3)/ep,
                    "tp":ep*(1+cfg["tp_pct"]/100),"sl":ep*(1-cfg["sl_pct"]/100)}
            elif z>cfg["z_entry"]:
                ep=c[i]; active={"side":"SHORT","ep":ep,"qty":(5.0*3)/ep,
                    "tp":ep*(1-cfg["tp_pct"]/100),"sl":ep*(1+cfg["sl_pct"]/100)}

        if active:
            all_trades.append({"sym":sym,"pnl":_pnl(active,c[-1],"END"),"reason":"END"})

    if not all_trades or len(all_trades)<5:
        return None

    wins=[t for t in all_trades if t["pnl"]>0]
    total_pnl=sum(t["pnl"] for t in all_trades)
    gp=sum(t["pnl"] for t in wins)
    gl=abs(sum(t["pnl"] for t in all_trades if t["pnl"]<=0))
    pf=gp/gl if gl>0 else 999

    return {
        "trades":len(all_trades),
        "wr":round(len(wins)/len(all_trades)*100,1),
        "pnl":round(total_pnl,2),
        "pf":round(pf,2),
        "tp_count":sum(1 for t in all_trades if t["reason"]=="TP"),
        "sl_count":sum(1 for t in all_trades if t["reason"]=="SL"),
        "ztp_count":sum(1 for t in all_trades if t["reason"]=="Z-TP"),
        "avg_pnl":round(total_pnl/len(all_trades),4),
    }


def _pnl(trade, close_price, reason):
    if trade["side"]=="LONG": pnl=trade["qty"]*(close_price-trade["ep"])
    else: pnl=trade["qty"]*(trade["ep"]-close_price)
    ef=trade["qty"]*trade["ep"]*TAKER_FEE
    xf=trade["qty"]*close_price*(MAKER_FEE if reason=="TP" else TAKER_FEE)
    return pnl-ef-xf


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--rr", type=int, required=True, help="R:R ratio (2=1:2, 3=1:3, etc)")
    parser.add_argument("--cache", default="backtests/data_cache_7d.pkl")
    args = parser.parse_args()

    rr = args.rr
    print(f"Loading cached data from {args.cache}...")
    with open(args.cache, "rb") as f:
        raw = pickle.load(f)
    data_cache = raw["symbols"]
    print(f"  {len(data_cache)} symbols, {raw['days']}d data from {raw['date']}")

    # Build combos
    keys = list(GRID.keys())
    base_combos = list(product(*[GRID[k] for k in keys]))

    combos = []
    for vals in base_combos:
        base_cfg = dict(zip(keys, vals))
        if base_cfg["z_entry"] >= base_cfg["z_max"]: continue
        for sl in SL_VALUES:
            tp = sl * rr
            cfg = {**base_cfg, "sl_pct": sl, "tp_pct": tp}
            combos.append(cfg)

    print(f"Sweeping {len(combos)} combos for R:R 1:{rr}...")

    results = []
    for idx, cfg in enumerate(combos):
        res = run_one(data_cache, cfg)
        if res and res["pnl"] > 0:  # only save profitable
            res["config"] = cfg
            res["rr"] = rr
            results.append(res)
        if (idx+1) % 500 == 0:
            print(f"  {idx+1}/{len(combos)} — {len(results)} profitable so far")

    results.sort(key=lambda x: x["pnl"], reverse=True)

    out_path = os.path.join(os.path.dirname(__file__), f"sweep_rr{rr}_results.json")
    with open(out_path, "w") as f:
        json.dump({
            "date": datetime.now().isoformat(),
            "rr": rr,
            "total_combos": len(combos),
            "profitable_combos": len(results),
            "top_15": results[:15],
        }, f, indent=2)

    print(f"\n{'='*90}")
    print(f"R:R 1:{rr} — {len(results)} profitable / {len(combos)} total")
    print(f"{'='*90}")

    if results:
        print(f"{'#':>3} {'PnL':>8} {'WR%':>6} {'PF':>6} {'Trd':>5} | Z    Zmax TP%   SL%   CHOP NATR")
        print("-"*90)
        for i, r in enumerate(results[:15]):
            c = r["config"]
            print(f"{i+1:>3} ${r['pnl']:>7.2f} {r['wr']:>5.1f}% {r['pf']:>5.2f} {r['trades']:>5} | "
                  f"{c['z_entry']:>4.1f} {c['z_max']:>4.1f} {c['tp_pct']:>5.2f} {c['sl_pct']:>5.2f} "
                  f"{c['chop_min']:>4} {c['natr_min']:.2f}-{c['natr_max']:.1f}")
    else:
        print("❌ No profitable combos found for this R:R")

    print(f"\nSaved to {out_path}")


if __name__ == "__main__":
    main()