# %% [markdown]
# ---
# title: Can it spot an ordinary value in the wrong context?
# categories:
# - Reconstruction
# proof-id: P076
# description: Learn to reconstruct a numeric answer from its context. Use a large reconstruction error
#   as a warning that the value does not fit.
# execute:
#   enabled: false
#   eval: false
# code-fold: true
# toc: false
# proof-readout:
# - label: Learned context
#   metric:
#   - informative
#   - anomaly_auc
#   format: auc
# - label: Shuffled training
#   metric:
#   - shuffled_training
#   - anomaly_auc
#   format: auc
# - label: False alarms on normal records
#   metric:
#   - informative
#   - normal_false_positive_rate
#   format: percent
# ---
#
# ## Example
#
# ```yaml
# group: amber
# x: 1.0
# y: 0.24
# ```
#
# The answer is shown here for explanation; it is hidden from the model when scored.
#
# ```{typst}
# //| label: fig-proof-contextual-consistency
# //| fig-cap: "Amber cards are hidden prediction targets. Other fields provide the input."
# //| fig-alt: "Model tree with record, group, x, y. Amber cards are hidden prediction targets. Other fields provide the input."
# #tree(node("record", kind: "root", children: (
#   node("group", type: "Enum", detail: "Selects a relationship"),
#   node("x", type: "Number", detail: "Visible measurement"),
#   node("y", kind: "target", type: "Number", detail: "Hidden for scoring"),
# )))
# ```
#
# ## Comparison
#
# Shuffle answers among test records so their individual distribution stays unchanged. Compare with an independently shuffled-training model.
#
# ## Result
#
# {{< proof P076 status >}}
#
# All three seeds detect this constructed inconsistency. The warning score is not a calibrated anomaly probability, and rare valid values need their own false-alarm check.
#
# <details class="proof-details">
# <summary>Experiment details and code</summary>
#
# {{< proof P076 evidence >}}
#
# ### Run this experiment
#
# {{< proof P076 script >}}
#
# ### Process, examples, and observability
#
# Draw a balanced group and `x ~ Uniform(-2, 2)`. Set
# `y = slope[group] * x + offset[group] + Normal(0, 0.1)`. The model observes
# only group and x when reconstructing y. Slopes and offsets are evaluator
# constants, never model inputs. Each split has its own random stream, with
# run seeds spaced by 100 before assigning stream offsets.
#
# Actual test anomalies permute y between records, preserving its entire distribution. Valid-tail tests restrict absolute x to [1.8, 2], within training support.
#
# ### Training, scoring, and gates
#
# Fit informative and shuffled-y models from the same initialization for 500
# AdamW updates each, batch size 128 and learning rate 0.002. Use 4,096 training
# rows, 2,048 clean validation rows, and separate 2,048-row test and tail sets.
# Select each model's residual threshold at the clean validation 95th percentile.
# No test observations choose that threshold or a checkpoint.
#
# Predeclared capability gates require informative test nRMSE below 0.20,
# anomaly AUROC at least 0.90, detection recall at least 0.75, normal false
# positives at most 0.10, and tail false positives at most 0.15. Shuffled
# training must leave AUROC at most 0.65. A univariate y-deviation baseline
# scores exactly 0.5 AUROC because its two score distributions are identical.
#
# ### Complete experiment code
#

# %%
"""P076: detect broken cross-field relationships using hidden-field reconstruction."""

from collections.abc import Iterator
from functools import partial

import lightning.pytorch as lit
import numpy as np
import pyarrow as pa
from reporting import report

import relflow as rf

PROOF_ID = "P076"
BUDGET = 500
GROUPS = ("amber", "birch", "cove", "dune")
SLOPES = np.array([0.7, -0.8, 1.2, -1.3])
OFFSETS = np.array([-0.5, 0.5, -0.3, 0.3])
TRAIN_ROWS = 4096
TEST_ROWS = 2048


def records(*, rows: int, seed: int, shuffled: bool = False, tail: bool = False) -> Iterator[dict]:
    rng = np.random.default_rng(seed)
    group = np.resize(np.arange(len(GROUPS)), rows)
    rng.shuffle(group)
    x = rng.uniform(-2, 2, rows)
    if tail:
        x = rng.choice([-1.0, 1.0], rows) * rng.uniform(1.8, 2.0, rows)
    y = SLOPES[group] * x + OFFSETS[group] + rng.normal(0, 0.1, rows)
    if shuffled:
        y = y[rng.permutation(rows)]
    for index in range(rows):
        yield {"group": GROUPS[group[index]], "x": float(x[index]), "y": float(y[index])}


def build() -> rf.Model:
    model = rf.Model.xs(
        batch_size=128,
        group=rf.Enum(values=GROUPS, p_unavailable=0.0),
        x=rf.Number,
        y=rf.Number(mask=True, objective="mse"),
    )
    model.optimizer = rf.adamw(learning_rate=0.002)
    return model


def predict(model: rf.Model, rows: list[dict]) -> np.ndarray:
    visible = [{"group": row["group"], "x": row["x"]} for row in rows]
    result = model.predict(pa.Table.from_pylist(visible))["predictions"].to_pylist()
    return np.asarray([row["y"]["content"] for row in result], dtype=float)


def auc(normal: np.ndarray, anomalous: np.ndarray) -> float:
    """Probability that an anomaly scores above a normal row, with half credit for ties."""
    ordered = np.sort(normal)
    left = np.searchsorted(ordered, anomalous, side="left")
    right = np.searchsorted(ordered, anomalous, side="right")
    return float(np.mean((left + right) / (2 * len(normal))))


def run(seed: int, steps: int | None, accelerator: str) -> tuple[dict, dict]:
    stream = seed * 100
    validation = list(records(rows=2048, seed=stream + 2))
    test = list(records(rows=TEST_ROWS, seed=stream + 3))
    anomalous = list(records(rows=TEST_ROWS, seed=stream + 3, shuffled=True))
    tail = list(records(rows=TEST_ROWS, seed=stream + 4, tail=True))
    actual = np.asarray([row["y"] for row in test])
    broken = np.asarray([row["y"] for row in anomalous])
    train_mean = float(np.mean([row["y"] for row in records(rows=TRAIN_ROWS, seed=stream + 1)]))
    baseline_rmse = float(np.sqrt(np.mean((actual - train_mean) ** 2)))
    marginal_auc = auc(np.abs(actual - train_mean), np.abs(broken - train_mean))
    metrics = {"train_rows": TRAIN_ROWS, "test_rows_per_condition": TEST_ROWS, "marginal_auc": marginal_auc}
    checks = {
        "Anomalies preserve the exact target marginal": bool(np.array_equal(np.sort(actual), np.sort(broken))),
        "Anomalies preserve visible contexts": all(
            (a["group"], a["x"]) == (b["group"], b["x"]) for a, b in zip(test, anomalous, strict=True)
        ),
        "Marginal deviation baseline has exactly chance AUROC": abs(marginal_auc - 0.5) < 1e-12,
    }
    for shuffled in (False, True):
        lit.seed_everything(seed, workers=True)
        model = build()
        data = rf.SyntheticDataModule(
            model=model,
            train=partial(records, rows=TRAIN_ROWS, seed=stream + 1, shuffled=shuffled),
            validate=partial(records, rows=2048, seed=stream + 2, shuffled=shuffled),
            seed=seed,
        )
        trainer = lit.Trainer(
            accelerator=accelerator,
            devices=1,
            max_epochs=-1,
            max_steps=BUDGET if steps is None else min(steps, BUDGET),
            deterministic=True,
            logger=False,
            enable_checkpointing=False,
            enable_progress_bar=False,
            enable_model_summary=False,
            num_sanity_val_steps=0,
        )
        trainer.fit(model, datamodule=data)
        model.eval()
        valid_error = np.abs(predict(model, validation) - [row["y"] for row in validation])
        threshold = float(np.quantile(valid_error, 0.95))
        prediction = predict(model, test)
        clean_score = np.abs(prediction - actual)
        broken_score = np.abs(prediction - broken)
        tail_score = np.abs(predict(model, tail) - [row["y"] for row in tail])
        key = "shuffled_training" if shuffled else "informative"
        result = {
            "steps": trainer.global_step,
            "validation_threshold": threshold,
            "rmse": float(np.sqrt(np.mean(clean_score**2))),
            "baseline_rmse": baseline_rmse,
            "normalized_rmse": float(np.sqrt(np.mean(clean_score**2)) / baseline_rmse),
            "anomaly_auc": auc(clean_score, broken_score),
            "detection_recall": float(np.mean(broken_score > threshold)),
            "normal_false_positive_rate": float(np.mean(clean_score > threshold)),
            "tail_false_positive_rate": float(np.mean(tail_score > threshold)),
        }
        metrics[key] = result
        checks[f"{key}: every measurement is finite"] = bool(np.isfinite(list(result.values())).all())
        if shuffled:
            checks["Shuffled training anomaly AUROC stays below 0.65"] = result["anomaly_auc"] <= 0.65
        else:
            checks["Informative reconstruction nRMSE below 0.20"] = result["normalized_rmse"] < 0.20
            checks["Contextual anomaly AUROC reaches 0.90"] = result["anomaly_auc"] >= 0.90
            checks["Anomaly recall reaches 0.75"] = result["detection_recall"] >= 0.75
            checks["Normal false positive rate at most 0.10"] = result["normal_false_positive_rate"] <= 0.10
            checks["Rare valid tail false positive rate at most 0.15"] = result["tail_false_positive_rate"] <= 0.15
    return metrics, checks


if __name__ == "__main__":
    report(PROOF_ID, run, seed=7601)

# %% [markdown]
# </details>