"""SEA06 — Overnight vs Intraday session capture.

IDEA: Split the 24h day into named trading sessions:
  - ASIA:         UTC 00-08 (Tokyo, Hong Kong, Singapore)
  - EUROPE:       UTC 08-16 (London open to US open)
  - US_INTRADAY:  UTC 13-21 (NYSE hours, overlap with Europe 13-16)
  - US_OVERNIGHT: UTC 21-24 & 00-01 (NY close to Asia open)

For each 1h bar, we assign it to a session. We track the EXPANDING-WINDOW
cumulative mean return per session (causal: uses only past bars).
At bar i, we go long (+1) during the session that has had the best
mean return so far (among those with enough samples >= min_samples).
If no session qualifies, we stay flat.

This captures the historically positive session with a continuously
updating, causal estimate — no look-ahead.

Vol-target applied to the direction signal.

Grid (4 configs total to stay <= 6 total backtests):
  - min_samples in [30, 90] x 1 TF (1h) = 2 calls (each covers BTC+ETH internally)
  - We also try the "best 2 sessions" variant: go long if session is in top-2

Causal guarantee:
  - session_mean[s] at bar i = mean of r[j] for all j < i in session s
  - direction[i] assigned from session_mean BEFORE updating with r[i]
  - lib shifts target by 1 bar before multiplying by returns
"""

import sys
sys.path.insert(0, "/opt/docker/PythagorasGoal/scripts/research/alt")
import altlib as al
import numpy as np
import pandas as pd


# Session definitions: list of UTC hours belonging to each session
SESSIONS = {
    "ASIA":         list(range(0, 8)),       # 00:00-07:59 UTC
    "EUROPE":       list(range(8, 16)),      # 08:00-15:59 UTC
    "US_INTRADAY":  list(range(13, 21)),     # 13:00-20:59 UTC
    "US_OVERNIGHT": list(range(21, 24)) + list(range(0, 2)),  # 21:00-01:59 UTC
}

# Map each UTC hour (0-23) to its primary session
# (some hours overlap; assign to highest-priority session)
# Priority: US_INTRADAY > EUROPE > ASIA > US_OVERNIGHT for overlapping hours
HOUR_TO_SESSION = {}
for h in range(24):
    assigned = None
    for sess, hours in SESSIONS.items():
        if h in hours:
            if assigned is None:
                assigned = sess
            # Apply priority: prefer US_INTRADAY > EUROPE > ASIA > US_OVERNIGHT
            priority = {"US_INTRADAY": 4, "EUROPE": 3, "ASIA": 2, "US_OVERNIGHT": 1}
            if priority[sess] > priority.get(assigned, 0):
                assigned = sess
    HOUR_TO_SESSION[h] = assigned if assigned else "ASIA"

SESSION_NAMES = list(SESSIONS.keys())
N_SESS = len(SESSION_NAMES)
SESS_IDX = {s: i for i, s in enumerate(SESSION_NAMES)}


def sea06_target(df: pd.DataFrame, min_samples: int = 30, top_n: int = 1) -> np.ndarray:
    """
    Go long during bars that belong to the top-N sessions by expanding-window mean return.

    Parameters
    ----------
    min_samples : int
        Minimum number of past bars in a session before we trust its mean.
    top_n : int
        Number of sessions to consider "good" (1 = only the best, 2 = best two).
    """
    dt = pd.to_datetime(df["datetime"])
    c = df["close"].values.astype(float)
    r = al.simple_returns(c)
    n = len(df)

    hours = dt.dt.hour.values  # 0..23
    bar_session = np.array([SESS_IDX[HOUR_TO_SESSION[h]] for h in hours], dtype=int)

    # Expanding cumulative stats per session
    sess_sum = np.zeros(N_SESS, dtype=float)
    sess_cnt = np.zeros(N_SESS, dtype=int)

    direction = np.zeros(n, dtype=float)

    for i in range(n):
        s = bar_session[i]

        # Compute mean returns for all sessions that have enough samples
        means = np.full(N_SESS, np.nan)
        for si in range(N_SESS):
            if sess_cnt[si] >= min_samples:
                means[si] = sess_sum[si] / sess_cnt[si]

        # Find top-N sessions by mean return (ignore NaN)
        valid_mask = np.isfinite(means)
        if valid_mask.sum() >= 1:
            valid_indices = np.where(valid_mask)[0]
            valid_means = means[valid_indices]
            # Sort descending by mean
            sorted_idx = valid_indices[np.argsort(-valid_means)]
            top_sessions = set(sorted_idx[:top_n].tolist())

            # Only go long if current bar's session is in top-N AND its mean > 0
            if s in top_sessions and means[s] > 0:
                direction[i] = 1.0

        # Update expanding window AFTER using it (causal)
        sess_sum[s] += r[i]
        sess_cnt[s] += 1

    tgt = al.vol_target(direction, df, target_vol=0.20, vol_win_days=30, leverage_cap=2.0)
    return tgt


if __name__ == "__main__":
    results = []

    # Grid: min_samples x top_n — 4 configs, 1 TF, 2 assets = 4 calls to study_weights
    # (each study_weights call runs both BTC and ETH internally)
    grid = [
        (30, 1),
        (90, 1),
        (30, 2),
        (90, 2),
    ]

    for min_samples, top_n in grid:
        name = f"SEA06-ms{min_samples}-top{top_n}"
        rep = al.study_weights(
            name,
            lambda df, ms=min_samples, tn=top_n: sea06_target(df, min_samples=ms, top_n=tn),
            tfs=("1h",),
        )
        print(al.fmt(rep))
        print("JSON:", al.as_json(rep))
        print()

        best_cell = max(rep["cells"], key=lambda c: c.get("min_asset_holdout_sharpe", -9))
        results.append((
            rep["verdict"].get("best_holdout_sharpe", best_cell.get("min_asset_holdout_sharpe", -9)),
            rep["verdict"].get("best_full_sharpe", best_cell.get("min_asset_full_sharpe", -9)),
            name,
            rep,
        ))

    # Pick the best config by hold-out Sharpe
    results.sort(key=lambda x: (x[0], x[1]), reverse=True)
    best_hold, best_full, best_name, best_rep = results[0]

    print("\n=== BEST CONFIG ===", best_name)
    print(al.fmt(best_rep))
    print("JSON:", al.as_json(best_rep))