"""train.py — the ONLY file the autoresearch agent may edit (Phase-1 contract).

Toy slice: a tiny self-supervised encoder (masked reconstruction of windowed
daily [return, realized_vol]) → FROZEN → linear probe predicts NEXT-day realized
vol → val_vol_r2 = OOS R². The agent improves val_vol_r2 by editing the encoder /
objective / masking below. Writes metrics.json (the scalar the loop reads).

    python train.py
"""
import json
import numpy as np
import pandas as pd
import torch
import torch.nn as nn

# --- agent-tunable knobs ---
WINDOW = 20
EMBED_DIM = 32
MASK_FRAC = 0.30
EPOCHS = 200
LR = 1e-3
SEED = 0
# ---------------------------

torch.manual_seed(SEED)
np.random.seed(SEED)
dev = "cuda" if torch.cuda.is_available() else "cpu"


def build():
    df = pd.read_parquet("data/processed/eurusd_daily.parquet").reset_index(drop=True)
    feats = df[["ret", "realized_vol"]].to_numpy(np.float32)
    target = df["realized_vol"].to_numpy(np.float32)  # predict NEXT-day RV
    X, y = [], []
    for t in range(WINDOW, len(df) - 1):
        X.append(feats[t - WINDOW:t])
        y.append(target[t + 1])
    X = np.stack(X); y = np.array(y, np.float32)
    n_tr = int(0.7 * len(X))                      # time-ordered OOS split
    mu, sd = X[:n_tr].mean((0, 1)), X[:n_tr].std((0, 1)) + 1e-8  # train-only stats
    X = (X - mu) / sd
    return (X[:n_tr], y[:n_tr]), (X[n_tr:], y[n_tr:])


class Encoder(nn.Module):
    def __init__(self, win, emb):
        super().__init__()
        self.net = nn.Sequential(
            nn.Flatten(), 
            nn.Linear(win * 2, 128), 
            nn.LayerNorm(128),
            nn.GELU(), 
            nn.Linear(128, emb)
        )

    def forward(self, x):
        return self.net(x)


def main():
    (Xtr, ytr), (Xte, yte) = build()
    Xtr_t = torch.tensor(Xtr, device=dev)
    enc = Encoder(WINDOW, EMBED_DIM).to(dev)
    dec = nn.Sequential(nn.Linear(EMBED_DIM, 128), nn.GELU(), nn.Linear(128, WINDOW * 2)).to(dev)
    opt = torch.optim.Adam(list(enc.parameters()) + list(dec.parameters()), lr=LR)

    for _ in range(EPOCHS):  # SSL: masked reconstruction of the window
        mask = (torch.rand_like(Xtr_t) > MASK_FRAC).float()
        rec = dec(enc((Xtr_t * mask)))
        loss = (((rec - Xtr_t.flatten(1)) ** 2) * (1 - mask.flatten(1))).mean()
        opt.zero_grad(); loss.backward(); opt.step()

    enc.eval()
    with torch.no_grad():  # FROZEN embeddings
        Etr = enc(Xtr_t).cpu().numpy()
        Ete = enc(torch.tensor(Xte, device=dev)).cpu().numpy()

    # linear probe (ridge, closed form) on frozen embeddings → val_vol_r2 (OOS R²)
    A = np.hstack([Etr, np.ones((len(Etr), 1))])
    w = np.linalg.solve(A.T @ A + 1e-3 * np.eye(A.shape[1]), A.T @ ytr)
    pred = np.hstack([Ete, np.ones((len(Ete), 1))]) @ w
    ss_res = ((yte - pred) ** 2).sum()
    ss_tot = ((yte - yte.mean()) ** 2).sum()
    val_vol_r2 = float(1 - ss_res / ss_tot)

    json.dump({"val_vol_r2": val_vol_r2, "n_test": len(yte),
               "knobs": {"WINDOW": WINDOW, "EMBED_DIM": EMBED_DIM, "MASK_FRAC": MASK_FRAC, "EPOCHS": EPOCHS}},
              open("metrics.json", "w"), indent=2)
    print("val_vol_r2 = %.4f  (n_test=%d, dev=%s)" % (val_vol_r2, len(yte), dev))


if __name__ == "__main__":
    main()