"""Evaluate fixed populations and intervention catalogs with condition-level rows."""

from __future__ import annotations
import json, time
from pathlib import Path
import numpy as np
import torch
from pivot.evaluation.metrics import (
    correlation,
    mmd2,
    sliced_wasserstein,
    energy_distance,
    retrieval_metrics,
    bootstrap,
)
from pivot.evaluation.rewards import Reward
from pivot.evaluation.inference import (
    encode_label,
    forward_predict,
    reward_guidance,
    project_and_rerank,
)


def save_json(path, value):
    """Write strict JSON. Undefined numerical statistics become null."""

    def clean(x):
        if isinstance(x, dict):
            return {str(k): clean(v) for k, v in x.items()}
        if isinstance(x, (list, tuple)):
            return [clean(v) for v in x]
        if isinstance(x, (float, np.floating)):
            return float(x) if np.isfinite(x) else None
        if isinstance(x, np.integer):
            return int(x)
        return x

    p = Path(path)
    p.parent.mkdir(parents=True, exist_ok=True)
    p.write_text(json.dumps(clean(value), indent=2, allow_nan=False))


def evaluate(
    data,
    predict,
    output,
    method="PIVOT",
    partition="test",
    catalog="single",
    reward_kind="cosine",
    n_cells=128,
    seed=0,
    model=None,
    device="cpu",
    guidance_steps=25,
    step_size=0.5,
    k_nearest=10,
    initialization="best",
    max_targets=None,
):
    """Forward and inverse evaluation on disjoint query and reference cells.

    Feature errors use normalized expression, with each target's observed top-20
    effect genes selected once. Population metrics use PCA coordinates. Inverse
    query populations are disjoint from the reserved candidate-outcome cells.
    The target condition is the aggregation and bootstrap unit.
    """
    rng = np.random.default_rng(seed)
    ci = data.indices(partition, True)
    if not len(ci):
        raise ValueError("Evaluation needs controls in the requested partition")
    ci = rng.choice(ci, min(n_cells, len(ci)), replace=False)
    c0 = data.emb[ci]
    control_expr = np.asarray(data.Xhvg[ci].mean(0)).ravel()
    control_latent = c0.mean(0)
    candidates = (
        data.singles
        if catalog == "single"
        else data.combos if catalog == "combination" else data.perturbations
    )
    if not candidates:
        raise ValueError("Candidate catalog is empty")
    labels = data.labels(partition)
    if max_targets:
        labels = sorted(
            rng.choice(labels, min(max_targets, len(labels)), replace=False)
        )
    # Cache endpoint predictions once for each candidate. Count all candidate-cell
    # predictions, including those used to choose a warm-start embedding.
    t0 = time.perf_counter()
    predictions = {p: np.asarray(predict(c0, p)) for p in candidates}
    catalog_seconds = time.perf_counter() - t0
    query_ids = data.indices(partition, False)
    refids = data.indices("reference", False)
    reference = {
        p: data.emb[np.intersect1d(refids, data.pert_to_idx[p])] for p in candidates
    }
    if any(len(v) == 0 for v in reference.values()):
        raise ValueError("Every candidate needs independent reference outcomes")
    forward = []
    inverse = []
    for p in labels:
        ids = np.intersect1d(query_ids, data.pert_to_idx[p])
        ids = rng.choice(ids, min(n_cells, len(ids)), replace=False)
        target = data.emb[ids]
        pred = predictions.get(p)
        if pred is None:
            pred = np.asarray(predict(c0, p))
        expression = np.asarray(data.Xhvg[ids].mean(0)).ravel()
        pred_expr = data.decode_to_genes(pred).mean(0)
        effect = expression - control_expr
        pred_effect = pred_expr - control_expr
        de = np.argsort(-np.abs(effect))[: min(20, len(effect))]
        train_genes = {g for q in data.labels("train") for g in data.parse(q)}
        forward.append(
            {
                "target": p,
                "query_cells": len(ids),
                "genes_observed_in_training": all(
                    g in train_genes for g in data.parse(p)
                ),
                "mse_expression": float(np.mean((pred_expr - expression) ** 2)),
                "effect_pearson": correlation(pred_effect, effect),
                "effect_pearson_de20": correlation(pred_effect[de], effect[de]),
                "mmd2": mmd2(pred, target, data.meta["mmd_gamma"]),
                "sliced_wasserstein": sliced_wasserstein(pred, target, seed=seed),
                "energy": energy_distance(pred, target),
                "variance_ratio": float(
                    np.var(pred, axis=0).sum()
                    / max(np.var(target, axis=0).sum(), 1e-12)
                ),
            }
        )
        if p not in candidates:
            continue
        reward = Reward(
            reward_kind,
            target_sample=target,
            control_ref=control_latent,
            gamma=data.meta["mmd_gamma"],
            device=device,
        )
        score = lambda pop: float(
            reward(torch.as_tensor(pop, dtype=torch.float32, device=device))
            .mean()
            .detach()
            .cpu()
        )
        measured = {q: score(v) for q, v in reference.items()}
        scores = {q: score(v) for q, v in predictions.items()}
        ranked = sorted(candidates, key=lambda q: (-scores[q], q))
        tied = sum(abs(scores[q] - scores[ranked[0]]) < 1e-10 for q in candidates)

        def row(name, ranking, cost, seconds, continuous=None):
            selected = ranking[0]
            out = {
                "target": p,
                "search": name,
                **retrieval_metrics(ranking, p),
                "selected": selected,
                "predicted_reward": scores.get(selected),
                "measured_reward": measured[selected],
                "measured_regret": max(measured.values()) - measured[selected],
                "selected_gene_overlap": len(
                    set(data.parse(selected)) & set(data.parse(p))
                )
                / max(len(set(data.parse(p))), 1),
                "catalog_size": len(candidates),
                "candidate_cell_evaluations": cost,
                "seconds": seconds,
                "top_score_ties": tied,
            }
            if continuous is not None:
                out["continuous_reward"] = continuous
            return out

        inverse.append(
            row(
                "exhaustive",
                ranked,
                len(candidates) * len(c0),
                catalog_seconds / len(labels),
            )
        )
        if model is not None and guidance_steps > 0:
            # A best-candidate initialization uses the complete ranked catalog.
            # Random initialization avoids this search but still projects over all embeddings.
            init = (
                ranked[0] if initialization == "best" else str(rng.choice(candidates))
            )
            e = encode_label(model, data, init, device)
            ct = torch.as_tensor(c0, device=device)
            t1 = time.perf_counter()
            es = reward_guidance(model, ct, reward, e, guidance_steps, step_size)
            guided = project_and_rerank(
                model, data, candidates, es, ct, reward, k_nearest, k_nearest, device
            )
            continuous = score(forward_predict(model, ct, es).cpu().numpy())
            names = [q for q, _ in guided]
            cost = (
                (len(candidates) if initialization == "best" else 0)
                + guidance_steps
                + min(k_nearest, len(candidates))
            )
            gr = row(
                "guidance_" + initialization,
                names,
                cost * len(c0),
                time.perf_counter()
                - t1
                + (catalog_seconds / len(labels) if initialization == "best" else 0),
                continuous,
            )
            gr["initial_candidate"] = init
            gr["predicted_reward"] = guided[0][1]
            gr["projection_reward_change"] = guided[0][1] - continuous
            inverse.append(gr)
    summary = {
        "forward": {
            k: bootstrap([r[k] for r in forward], seed)
            for k in (
                "mse_expression",
                "effect_pearson",
                "effect_pearson_de20",
                "mmd2",
                "sliced_wasserstein",
                "energy",
                "variance_ratio",
            )
        },
        "inverse": {},
    }
    for search in sorted({r["search"] for r in inverse}):
        rows = [r for r in inverse if r["search"] == search]
        summary["inverse"][search] = {
            k: bootstrap([r[k] for r in rows], seed)
            for k in (
                "top1",
                "top5",
                "ndcg10",
                "measured_reward",
                "measured_regret",
                "selected_gene_overlap",
            )
        }
    result = {
        "protocol": "split-first-v1",
        "method": method,
        "data_meta": data.meta,
        "partition": partition,
        "catalog": catalog,
        "reward": reward_kind,
        "seed": seed,
        "n_controls": len(c0),
        "query_cell_ids": data.obs.iloc[query_ids].cell_id.tolist(),
        "guidance": {
            "steps": guidance_steps,
            "step_size": step_size,
            "k_nearest": k_nearest,
            "initialization": initialization,
        },
        "summary": summary,
        "forward": forward,
        "inverse": inverse,
    }
    save_json(output, result)
    return result