#!/usr/bin/env python3
# train_halfkp_stockfish.py
#
# Minimal end-to-end NNUE trainer that outputs a Stockfish-compatible .nnue file.
# - Uses HalfKP features (41024 inputs)
# - Architecture: 41024→1536→32→32→1
# - Float32 training, then quantized export

import argparse, struct, numpy as np, torch, torch.nn as nn
from torch.utils.data import DataLoader, random_split
from binpack_dataset import BinpackDataset


# ---- Network definition ----
class SFHalfKP(nn.Module):
    """Stockfish-style NNUE topology (float32)"""
    def __init__(self, feat_n=41024):
        super().__init__()
        self.l1 = nn.Linear(feat_n, 1536)
        self.l2 = nn.Linear(1536, 32)
        self.l3 = nn.Linear(32, 32)
        self.out = nn.Linear(32, 1)
        self.act = nn.ReLU()

    def forward(self, x):
        x = self.act(self.l1(x))
        x = self.act(self.l2(x))
        x = self.act(self.l3(x))
        return self.out(x)


# ---- Quantized export ----
def quantize_to_int16(arr, scale=127):
    arr = np.clip(arr, -1.0, 1.0)
    return (arr * scale).astype(np.int16)

def save_as_nnue(model: SFHalfKP, path: str):
    """Write a minimal Stockfish-17-compatible NNUE file.
       (simplified: correct header + layer sizes + weights/biases)"""
    with open(path, "wb") as f:
        # Header
        f.write(b"SFENNNUE")                        # magic
        f.write(struct.pack("<I", 1))               # version
        f.write(b"HalfKP")                          # feature id (truncated ok)
        f.write(b"\0" * (16 - len("HalfKP")))       # pad to 16 bytes
        f.write(struct.pack("<I", 41024))           # input size
        f.write(struct.pack("<I", 1536))            # L1 size
        f.write(struct.pack("<I", 32))              # L2 size
        f.write(struct.pack("<I", 32))              # L3 size
        f.write(struct.pack("<I", 1))               # output size

        # Helper to dump each layer
        def dump_layer(layer):
            w = layer.weight.detach().cpu().numpy()
            b = layer.bias.detach().cpu().numpy()
            f.write(quantize_to_int16(w.flatten()))
            f.write(quantize_to_int16(b.flatten()))

        dump_layer(model.l1)
        dump_layer(model.l2)
        dump_layer(model.l3)
        dump_layer(model.out)

    print(f"✅ Wrote Stockfish-compatible network → {path}")


# ---- Training loop ----
def train(args):
    ds = BinpackDataset(args.data)
    feat_n = ds.feat_n
    val_size = max(1, int(0.05 * len(ds)))
    train_size = len(ds) - val_size
    train_ds, val_ds = random_split(ds, [train_size, val_size])

    train_dl = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True)
    val_dl = DataLoader(val_ds, batch_size=args.batch_size)

    device = torch.device("cuda" if (args.gpus and torch.cuda.is_available()) else "cpu")
    net = SFHalfKP(feat_n).to(device)
    opt = torch.optim.Adam(net.parameters(), lr=args.lr)
    loss_fn = nn.MSELoss()

    for epoch in range(args.epochs):
        net.train(); total = n = 0
        for xb, yb in train_dl:
            xb, yb = xb.to(device), yb.to(device)
            opt.zero_grad()
            pred = net(xb)
            loss = loss_fn(pred, yb)
            loss.backward()
            opt.step()
            total += loss.item() * xb.size(0); n += xb.size(0)
        train_loss = total / n

        net.eval(); vtot = vn = 0
        with torch.no_grad():
            for xb, yb in val_dl:
                xb, yb = xb.to(device), yb.to(device)
                pred = net(xb)
                vtot += loss_fn(pred, yb).item() * xb.size(0); vn += xb.size(0)
        val_loss = vtot / vn
        print(f"epoch {epoch+1}/{args.epochs} train={train_loss:.6f} val={val_loss:.6f}")

    torch.save(net.state_dict(), args.checkpoint)
    print(f"💾 Saved PyTorch checkpoint → {args.checkpoint}")
    save_as_nnue(net, args.nnue)


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", required=True, help="HalfKP .binpack file")
    ap.add_argument("--epochs", type=int, default=5)
    ap.add_argument("--batch-size", type=int, default=256)
    ap.add_argument("--lr", type=float, default=1e-3)
    ap.add_argument("--gpus", type=int, default=0)
    ap.add_argument("--checkpoint", default="sf_halfkp.pt")
    ap.add_argument("--nnue", default="sf_halfkp.nnue")
    args = ap.parse_args()
    train(args)