"""Reference predictors fitted with the same training cells and feature basis."""

from __future__ import annotations
import copy
import numpy as np
import torch
from sklearn.linear_model import Ridge
from pivot.training.train import make_model, TrainConfig, validation_loss
from pivot.models.encoders import build_pert_tensors
from pivot.utils.common import set_seed


class Baseline:
    """Population predictor. `predict(c0, label)` returns (n_controls, d)."""

    def __init__(
        self, data, kind="ridge", alpha=1.0, epochs=60, hidden=512, seed=0, device="cpu"
    ):
        if kind not in (
            "mean_control",
            "average_effect",
            "additive",
            "ridge",
            "endpoint_mlp",
            "conditional_mlp",
        ):
            raise ValueError(kind)
        self.data = data
        self.kind = kind
        self.device = device
        tr = data.indices("train", False)
        ctrl = data.indices("train", True)
        self.control = data.emb[ctrl].mean(0)
        self.effects = {
            p: data.emb[np.intersect1d(tr, data.pert_to_idx[p])].mean(0) - self.control
            for p in data.labels("train")
        }
        self.avg = np.mean(list(self.effects.values()), axis=0)
        self.training_info = {
            "kind": kind,
            "alpha": alpha,
            "seed": seed,
            "train_cell_count": len(tr),
        }
        if kind == "ridge":
            labels = list(self.effects)
            self.reg = Ridge(alpha=alpha, fit_intercept=False).fit(
                self.features(labels), np.stack(list(self.effects.values()))
            )
        if kind in ("endpoint_mlp", "conditional_mlp"):
            set_seed(seed)
            torch.set_num_threads(4)
            cfg = TrainConfig(hidden=hidden, depth=4, seed=seed, device=device)
            self.model = make_model(data, cfg)
            opt = torch.optim.AdamW(self.model.parameters(), lr=1e-3, weight_decay=1e-5)
            rng = np.random.default_rng(seed)
            best = float("inf")
            beststate = None
            history = []
            for epoch in range(epochs):
                self.model.train()
                for ids in np.array_split(
                    rng.permutation(tr), max(1, int(np.ceil(len(tr) / 1024)))
                ):
                    c = data.emb[rng.choice(ctrl, len(ids))]
                    if kind == "endpoint_mlp":
                        c = np.broadcast_to(self.control, c.shape).copy()
                    c = torch.as_tensor(c, device=device)
                    y = torch.as_tensor(data.emb[ids], device=device)
                    labels = data.obs.iloc[ids].perturbation.tolist()
                    g, o, m, p = build_pert_tensors(data, labels, device)
                    pred = self.model.endpoint_from_pert(c, g, o, m, p)
                    loss = (pred - y).square().sum(-1).mean()
                    opt.zero_grad()
                    loss.backward()
                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), 5.0)
                    opt.step()
                self.model.eval()
                vals = []
                vc = data.emb[data.indices("val", True)][:128]
                for p in data.labels("val"):
                    ids = np.intersect1d(
                        data.indices("val", False), data.pert_to_idx[p]
                    )
                    vals.append(
                        np.mean(
                            (self.predict(vc, p).mean(0) - data.emb[ids].mean(0)) ** 2
                        )
                    )
                score = float(np.mean(vals))
                history.append(score)
                if score < best:
                    best = score
                    beststate = copy.deepcopy(self.model.state_dict())
            self.model.load_state_dict(beststate)
            self.training_info.update(
                validation_mse=history,
                best_validation_mse=best,
                epochs=epochs,
                hidden=hidden,
            )

    def features(self, labels):
        a = np.zeros((len(labels), len(self.data.genes_vocab)), np.float32)
        for i, p in enumerate(labels):
            for g in self.data.parse(p):
                a[i, self.data.gene_to_id[g]] = 1.0
        return a

    def predict(self, c0, label):
        if self.kind in ("endpoint_mlp", "conditional_mlp"):
            c = np.asarray(c0, dtype=np.float32)
            if self.kind == "endpoint_mlp":
                c = np.broadcast_to(self.control, c.shape).copy()
            g, o, m, p = build_pert_tensors(self.data, [label], self.device)
            with torch.no_grad():
                return (
                    self.model.endpoint_from_pert(
                        torch.as_tensor(c, device=self.device), g, o, m, p
                    )
                    .cpu()
                    .numpy()
                )
        if self.kind == "mean_control":
            delta = np.zeros(self.data.d)
        elif self.kind == "average_effect":
            delta = self.avg
        elif self.kind == "additive":
            # Unsupported single genes receive the training-average effect.
            delta = sum(
                (self.effects.get(g, self.avg) for g in self.data.parse(label)),
                start=np.zeros(self.data.d),
            )
        else:
            delta = self.reg.predict(self.features([label]))[0]
        return np.asarray(c0) + delta