"""S3-02: Lead-lag multi-asset squeeze.
Quando BTC fa squeeze breakout, ETH/SOL spesso seguono.
Usa il breakout di BTC per anticipare entrata su ETH (e viceversa).
Testa anche correlazione inter-asset per conferma segnale.
"""
from __future__ import annotations
import sys
sys.path.insert(0, ".")

import numpy as np
import pandas as pd
from src.data.downloader import load_data

FEE_RT = 0.002
INITIAL = 1000
LEVERAGE = 3


def keltner_ratio(close, high, low, window=14):
    n = len(close)
    r = np.full(n, np.nan)
    for i in range(window, n):
        wc, wh, wl = close[i-window:i], high[i-window:i], low[i-window:i]
        ma = np.mean(wc)
        bb_std = np.std(wc)
        tr = np.maximum(wh-wl, np.maximum(np.abs(wh-np.roll(wc,1)), np.abs(wl-np.roll(wc,1))))
        atr = np.mean(tr[1:])
        kc = (ma+1.5*atr)-(ma-1.5*atr)
        bb = (ma+2*bb_std)-(ma-2*bb_std)
        if kc > 0: r[i] = bb/kc
    return r


def load_aligned(assets, tf):
    """Carica e allinea dati multi-asset per timestamp."""
    dfs = {}
    for asset in assets:
        try:
            if asset == "SOL":
                df = pd.read_parquet(f"data/raw/sol_{tf}.parquet")
                df["datetime"] = pd.to_datetime(df["timestamp"], unit="ms", utc=True)
            else:
                df = load_data(asset, tf)
            dfs[asset] = df
        except Exception:
            pass

    if len(dfs) < 2:
        return None

    # Allinea per timestamp
    common_ts = set(dfs[list(dfs.keys())[0]]["timestamp"].values)
    for df in dfs.values():
        common_ts &= set(df["timestamp"].values)
    common_ts = sorted(common_ts)

    aligned = {}
    for asset, df in dfs.items():
        mask = df["timestamp"].isin(common_ts)
        aligned[asset] = df[mask].sort_values("timestamp").reset_index(drop=True)

    return aligned


def detect_breakouts(close, high, low, volume, kcr, sq_thr=0.8, min_dur=5):
    """Detect squeeze breakout events."""
    events = []
    in_sq = False
    sq_start = 0
    for i in range(1, len(close)):
        if np.isnan(kcr[i]):
            continue
        is_sq = kcr[i] < sq_thr
        if is_sq and not in_sq:
            in_sq = True
            sq_start = i
        elif not is_sq and in_sq:
            in_sq = False
            if i - sq_start < min_dur:
                continue
            first_ret = (close[i] - close[i-1]) / close[i-1] if close[i-1] > 0 else 0
            if abs(first_ret) < 0.001:
                continue
            events.append({
                "idx": i,
                "duration": i - sq_start,
                "direction": 1 if first_ret > 0 else -1,
                "first_ret": first_ret,
            })
    return events


print("=" * 75)
print("  S3-02: LEAD-LAG MULTI-ASSET SQUEEZE")
print("=" * 75)

for tf in ["1h", "15m"]:
    aligned = load_aligned(["BTC", "ETH", "SOL"], tf)
    if aligned is None:
        continue

    n = len(aligned["BTC"])
    ts = pd.to_datetime(aligned["BTC"]["timestamp"], unit="ms", utc=True)

    print(f"\n  Timeframe: {tf}, Candles allineate: {n}")

    # Calcola squeeze per ogni asset
    asset_data = {}
    for asset in aligned:
        df = aligned[asset]
        c, h, l, v = df["close"].values, df["high"].values, df["low"].values, df["volume"].values
        kcr = keltner_ratio(c, h, l, 14)
        events = detect_breakouts(c, h, l, v, kcr)
        asset_data[asset] = {"close": c, "high": h, "low": l, "vol": v, "kcr": kcr, "events": events}
        print(f"    {asset}: {len(events)} squeeze breakouts")

    # ================================================================
    #  STRATEGIA A: Leader-follower
    #  Quando BTC fa breakout, entra su ETH/SOL nella stessa direzione
    # ================================================================
    print(f"\n  --- LEADER-FOLLOWER ({tf}) ---")

    for leader, follower in [("BTC", "ETH"), ("BTC", "SOL"), ("ETH", "BTC"), ("ETH", "SOL")]:
        if leader not in asset_data or follower not in asset_data:
            continue

        leader_events = asset_data[leader]["events"]
        fc = asset_data[follower]["close"]

        for hold in [3, 6]:
            for delay in [0, 1, 2]:
                yearly = {}

                for ev in leader_events:
                    i = ev["idx"] + delay
                    if i + hold >= n:
                        continue

                    # Anti-fakeout su follower
                    entry = fc[i]
                    exit_price = fc[min(i + hold, n - 1)]
                    direction = ev["direction"]
                    actual = (exit_price - entry) / entry * direction
                    net = actual * LEVERAGE - FEE_RT * LEVERAGE

                    year = ts.iloc[min(i, n-1)].year
                    if year not in yearly:
                        yearly[year] = {"w": 0, "t": 0, "pnls": []}
                    yearly[year]["t"] += 1
                    if actual > 0:
                        yearly[year]["w"] += 1
                    yearly[year]["pnls"].append(net * INITIAL)

                all_t = sum(d["t"] for d in yearly.values())
                all_w = sum(d["w"] for d in yearly.values())
                if all_t < 30:
                    continue
                acc = all_w / all_t * 100
                pnl = sum(p for d in yearly.values() for p in d["pnls"])
                worst_y = min(yearly.items(), key=lambda x: x[1]["w"]/x[1]["t"] if x[1]["t"]>0 else 0)
                worst_acc = worst_y[1]["w"]/worst_y[1]["t"]*100 if worst_y[1]["t"]>0 else 0
                tag = "✅" if acc >= 76 else ""
                print(f"    {leader}→{follower} d={delay} h={hold}: trades={all_t:5d} acc={acc:.1f}% pnl=€{pnl:+.0f} worst={worst_y[0]}({worst_acc:.0f}%) {tag}")

    # ================================================================
    #  STRATEGIA B: Consensus multi-asset
    #  Trade solo quando 2+ asset hanno squeeze breakout nello stesso momento
    # ================================================================
    print(f"\n  --- CONSENSUS MULTI-ASSET ({tf}) ---")

    # Build event map: timestamp → list of (asset, direction)
    event_map = {}
    for asset, data in asset_data.items():
        for ev in data["events"]:
            idx = ev["idx"]
            if idx not in event_map:
                event_map[idx] = []
            event_map[idx].append((asset, ev["direction"]))

    for target in ["BTC", "ETH", "SOL"]:
        if target not in asset_data:
            continue
        tc = asset_data[target]["close"]

        for min_consensus in [2, 3]:
            for window_bars in [1, 3, 5]:
                yearly = {}
                daily_done = set()

                for idx in sorted(event_map.keys()):
                    if idx + 6 >= n:
                        continue

                    day = ts.iloc[idx].strftime("%Y-%m-%d")
                    if day in daily_done:
                        continue

                    # Count consensus within window
                    nearby_events = []
                    for j in range(max(0, idx - window_bars), idx + window_bars + 1):
                        if j in event_map:
                            nearby_events.extend(event_map[j])

                    # Unique assets
                    unique_assets = set(a for a, d in nearby_events)
                    if len(unique_assets) < min_consensus:
                        continue

                    # Majority direction
                    dirs = [d for a, d in nearby_events]
                    majority = 1 if sum(dirs) > 0 else -1

                    entry = tc[idx]
                    exit_price = tc[min(idx + 3, n - 1)]
                    actual = (exit_price - entry) / entry * majority
                    net = actual * LEVERAGE - FEE_RT * LEVERAGE

                    year = ts.iloc[idx].year
                    if year not in yearly:
                        yearly[year] = {"w": 0, "t": 0, "pnls": []}
                    yearly[year]["t"] += 1
                    if actual > 0:
                        yearly[year]["w"] += 1
                    yearly[year]["pnls"].append(net * INITIAL)
                    daily_done.add(day)

                all_t = sum(d["t"] for d in yearly.values())
                all_w = sum(d["w"] for d in yearly.values())
                if all_t < 20:
                    continue
                acc = all_w / all_t * 100
                pnl = sum(p for d in yearly.values() for p in d["pnls"])
                tag = "✅" if acc >= 76 else ""
                print(f"    {target} consensus>={min_consensus} w={window_bars}: trades={all_t:4d} acc={acc:.1f}% pnl=€{pnl:+.0f} {tag}")

    # ================================================================
    #  STRATEGIA C: Correlation-weighted squeeze
    #  Peso il segnale squeeze in base alla correlazione rolling con BTC
    # ================================================================
    print(f"\n  --- CORRELATION-WEIGHTED ({tf}) ---")

    for target in ["ETH", "SOL"]:
        if target not in asset_data:
            continue
        tc = asset_data[target]["close"]
        btc_c = asset_data["BTC"]["close"]

        # Rolling correlation
        corr_window = 48  # 48 bars
        rolling_corr = np.full(n, np.nan)
        ret_t = np.diff(np.log(np.where(tc == 0, 1e-10, tc)))
        ret_b = np.diff(np.log(np.where(btc_c == 0, 1e-10, btc_c)))
        for i in range(corr_window, len(ret_t)):
            c_val = np.corrcoef(ret_t[i-corr_window:i], ret_b[i-corr_window:i])[0, 1]
            rolling_corr[i + 1] = c_val if np.isfinite(c_val) else 0

        events = asset_data[target]["events"]

        for corr_thr in [0.5, 0.6, 0.7, 0.8]:
            yearly = {}
            for ev in events:
                i = ev["idx"]
                if i + 3 >= n or np.isnan(rolling_corr[i]):
                    continue

                # Solo quando correlazione con BTC è alta
                if abs(rolling_corr[i]) < corr_thr:
                    continue

                entry = tc[i - 1]
                exit_price = tc[min(i + 2, n - 1)]
                actual = (exit_price - entry) / entry * ev["direction"]
                net = actual * LEVERAGE - FEE_RT * LEVERAGE

                year = ts.iloc[i].year
                if year not in yearly:
                    yearly[year] = {"w": 0, "t": 0, "pnls": []}
                yearly[year]["t"] += 1
                if actual > 0:
                    yearly[year]["w"] += 1
                yearly[year]["pnls"].append(net * INITIAL)

            all_t = sum(d["t"] for d in yearly.values())
            all_w = sum(d["w"] for d in yearly.values())
            if all_t < 20:
                continue
            acc = all_w / all_t * 100
            pnl = sum(p for d in yearly.values() for p in d["pnls"])
            tag = "✅" if acc >= 76 else ""
            print(f"    {target} corr>={corr_thr}: trades={all_t:4d} acc={acc:.1f}% pnl=€{pnl:+.0f} {tag}")
