"""Train PIRD: a paraphrase-invariant, multi-signal detector.

Objective (the contribution):
  - detection loss (human vs AI)
  - augmentation: paraphrased AI is included, still labelled AI
  - invariance: consistency loss pulling P(AI|x) ~ P(AI|paraphrase(x))
Full PIRD additionally fuses Stream A (statistical) + Stream C (stylometric) features onto the
encoder embedding, and fits a calibration temperature on a held-out val split (contribution C3).
Set use_features=False for the encoder-only "PIRD-lite" ablation.
"""
from __future__ import annotations
import json
import os
import random
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer

from .model import PIRDModel
from .attacks import Paraphraser
from .features import CombinedFeatures, standardize


class _PairDataset(Dataset):
    def __init__(self, items, tok, max_len=256, use_features=False):
        self.items = items; self.tok = tok; self.max_len = max_len; self.use_features = use_features

    def __len__(self):
        return len(self.items)

    def __getitem__(self, i):
        return self.items[i]

    def collate(self, batch):
        def enc(texts):
            e = self.tok(texts, return_tensors="pt", truncation=True,
                         max_length=self.max_len, padding=True)
            return e["input_ids"], e["attention_mask"]
        ids, mask = enc([b["text"] for b in batch])
        out = {"ids": ids, "mask": mask,
               "labels": torch.tensor([b["label"] for b in batch], dtype=torch.float)}
        if self.use_features:
            out["extra"] = torch.tensor(np.stack([b["extra"] for b in batch]), dtype=torch.float)
        has_para = [b for b in batch if b.get("para")]
        if has_para:
            p_ids, p_mask = enc([b["para"] for b in has_para])
            out.update({"p_ids": p_ids, "p_mask": p_mask,
                        "p_labels": torch.tensor([b["label"] for b in has_para], dtype=torch.float)})
            if self.use_features:
                out["p_extra"] = torch.tensor(np.stack([b["p_extra"] for b in has_para]),
                                              dtype=torch.float)
        return out


def build_items(human, ai, recursive_rounds=0, seed=42, paraphraser=None):
    para = paraphraser or Paraphraser()
    print(f"paraphrasing {len(ai)} AI texts (rounds={max(1, recursive_rounds)}) ...")
    ai_para = para.paraphrase_many(ai, rounds=max(1, recursive_rounds))
    items = [{"text": h, "label": 0.0, "para": None} for h in human]
    items += [{"text": a, "label": 1.0, "para": ap} for a, ap in zip(ai, ai_para)]
    random.Random(seed).shuffle(items)
    return items


def _attach_features(items, extractor, mean, std):
    X = standardize(extractor.matrix([it["text"] for it in items]), mean, std)
    for it, x in zip(items, X):
        it["extra"] = x.astype("float32")
    pidx = [i for i, it in enumerate(items) if it.get("para")]
    if pidx:
        P = standardize(extractor.matrix([items[i]["para"] for i in pidx]), mean, std)
        for j, i in enumerate(pidx):
            items[i]["p_extra"] = P[j].astype("float32")


def _fit_temperature(model, items, tok, max_len, device, use_features):
    model.eval()
    zs, ys = [], []
    with torch.no_grad():
        for i in range(0, len(items), 16):
            chunk = items[i:i + 16]
            e = tok([c["text"] for c in chunk], return_tensors="pt", truncation=True,
                    max_length=max_len, padding=True)
            extra = (torch.tensor(np.stack([c["extra"] for c in chunk]), dtype=torch.float).to(device)
                     if use_features else None)
            z = model(e["input_ids"].to(device), e["attention_mask"].to(device), extra)
            zs.append(np.atleast_1d(z.cpu().numpy())); ys.append([c["label"] for c in chunk])
    model.train()
    z = np.concatenate(zs); y = np.concatenate(ys)
    best_T, best = 1.0, 1e9
    for T in np.linspace(0.5, 5.0, 91):
        p = np.clip(1.0 / (1.0 + np.exp(-z / T)), 1e-6, 1 - 1e-6)
        bce = -(y * np.log(p) + (1 - y) * np.log(1 - p)).mean()
        if bce < best:
            best, best_T = bce, float(T)
    return best_T


def train_pird(items, out_dir="pird_ckpt", encoder="roberta-base", epochs=3, batch_size=8,
               lr=2e-5, max_len=256, lam_inv=1.0, lam_aug=1.0, seed=42, device=None,
               use_features=True, stat_model="gpt2", val_frac=0.15):
    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)

    random.Random(seed).shuffle(items)
    n_val = max(40, int(len(items) * val_frac))
    val_items, train_items = items[:n_val], items[n_val:]

    n_extra, mean, std = 0, None, None
    if use_features:
        extractor = CombinedFeatures(stat_model, device=device)
        Xo = extractor.matrix([it["text"] for it in train_items])
        mean, std = Xo.mean(0), Xo.std(0); std[std < 1e-6] = 1.0
        n_extra = Xo.shape[1]
        print(f"[pird] fusing {n_extra} A+C features (standardized)")
        _attach_features(train_items, extractor, mean, std)
        _attach_features(val_items, extractor, mean, std)

    tok = AutoTokenizer.from_pretrained(encoder)
    model = PIRDModel(encoder, n_extra=n_extra).to(device)
    ds = _PairDataset(train_items, tok, max_len, use_features)
    dl = DataLoader(ds, batch_size=batch_size, shuffle=True, collate_fn=ds.collate)
    opt = torch.optim.AdamW(model.parameters(), lr=lr)
    bce = nn.BCEWithLogitsLoss()

    model.train()
    for epoch in range(epochs):
        run = {"det": 0.0, "aug": 0.0, "inv": 0.0}
        for step, b in enumerate(dl):
            extra = b.get("extra").to(device) if "extra" in b else None
            logits = model(b["ids"].to(device), b["mask"].to(device), extra)
            labels = b["labels"].to(device)
            loss = bce(logits, labels)
            run["det"] += loss.item()
            if "p_ids" in b:
                p_extra = b.get("p_extra").to(device) if "p_extra" in b else None
                p_logits = model(b["p_ids"].to(device), b["p_mask"].to(device), p_extra)
                p_labels = b["p_labels"].to(device)
                aug = bce(p_logits, p_labels)
                ai_logits = logits[labels == 1.0]
                m = min(len(ai_logits), len(p_logits))
                inv = ((torch.sigmoid(ai_logits[:m]) - torch.sigmoid(p_logits[:m])) ** 2).mean()
                loss = loss + lam_aug * aug + lam_inv * inv
                run["aug"] += aug.item(); run["inv"] += inv.item()
            if not torch.isfinite(loss):
                print(f"  WARN non-finite loss at step {step}, skipping"); opt.zero_grad(); continue
            opt.zero_grad(); loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
        n = len(dl)
        print(f"epoch {epoch+1}/{epochs}  det={run['det']/n:.4f}  aug={run['aug']/n:.4f}  "
              f"inv={run['inv']/n:.4f}")

    temperature = _fit_temperature(model, val_items, tok, max_len, device, use_features)
    print(f"[pird] calibration temperature T={temperature:.2f}")

    os.makedirs(out_dir, exist_ok=True)
    torch.save(model.state_dict(), os.path.join(out_dir, "pird.pt"))
    cfg = {"encoder": encoder, "n_extra": n_extra, "max_len": max_len,
           "use_features": use_features, "stat_model": stat_model, "temperature": temperature,
           "feat_mean": (mean.tolist() if mean is not None else None),
           "feat_std": (std.tolist() if std is not None else None)}
    with open(os.path.join(out_dir, "config.json"), "w") as f:
        json.dump(cfg, f)
    print(f"saved PIRD checkpoint -> {out_dir}")
    return out_dir