"""Ground-truth factor probes for synthetic-world representations.

Linear probes from a representation to the discrete scene factors measure
which world variables survive the encoder and its pooling. Probes are
fitted per modality on that modality's own training split and evaluated
on held-out scenes; scene truth is generator metadata, so this is a
within-modality diagnostic, not a cross-modal signal.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import torch

from .common import read_json, write_json
from .synth_world import COLORS, SHAPES


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--data-dir", default="artifacts/synth_v0")
    parser.add_argument("--features", required=True)
    parser.add_argument(
        "--side", choices=["vision", "text"], required=True,
        help="Selects the training split whose scenes fit the probes.",
    )
    parser.add_argument("--ridge", type=float, default=1.0)
    parser.add_argument("--output", required=True)
    return parser.parse_args()


def scene_targets(scene: dict) -> dict[str, np.ndarray | float]:
    color_presence = np.zeros(len(COLORS))
    shape_presence = np.zeros(len(SHAPES))
    for group in scene["groups"]:
        color_presence[list(COLORS).index(group["color"])] = 1.0
        shape_presence[SHAPES.index(group["shape"])] = 1.0
    return {
        "color_presence": color_presence,
        "shape_presence": shape_presence,
        "group_count": float(len(scene["groups"])),
        "object_total": float(sum(g["count"] for g in scene["groups"])),
        "relation_count": float(len(scene["relations"])),
    }


def ridge_fit(
    x: np.ndarray, y: np.ndarray, ridge: float
) -> tuple[np.ndarray, np.ndarray]:
    x = np.concatenate([x, np.ones((len(x), 1))], axis=1)
    gram = x.T @ x + ridge * np.eye(x.shape[1])
    weights = np.linalg.solve(gram, x.T @ y)
    return weights, x @ weights


def evaluate(
    features_train: np.ndarray,
    features_test: np.ndarray,
    train_targets: np.ndarray,
    test_targets: np.ndarray,
    ridge: float,
    binary: bool,
) -> dict:
    weights, _ = ridge_fit(features_train, train_targets, ridge)
    prediction = (
        np.concatenate([features_test, np.ones((len(features_test), 1))], axis=1)
        @ weights
    )
    if binary:
        accuracy = float(((prediction > 0.5) == (test_targets > 0.5)).mean())
        balanced = []
        for column in range(test_targets.shape[1]):
            truth = test_targets[:, column] > 0.5
            if truth.any() and (~truth).any():
                hit = (prediction[:, column] > 0.5) == truth
                balanced.append(
                    (hit[truth].mean() + hit[~truth].mean()) / 2.0
                )
        return {
            "accuracy": accuracy,
            "balanced_accuracy": float(np.mean(balanced)),
        }
    residual = prediction[:, 0] - test_targets[:, 0]
    variance = test_targets[:, 0].var()
    return {
        "r2": float(1.0 - residual.var() / max(variance, 1e-9)),
        "mae": float(np.abs(residual).mean()),
    }


def main() -> None:
    args = parse_args()
    manifest = read_json(Path(args.data_dir, "manifest.json"))
    scenes = read_json(Path(args.data_dir, "scenes.private.json"))["scenes"]
    state = torch.load(args.features, map_location="cpu", weights_only=False)
    lookup = {int(row): i for i, row in enumerate(state["rows"])}
    features = state["features"].float().numpy()

    split_key = "vision_only_train" if args.side == "vision" else "text_only_train"
    train_rows = [row for row in manifest[split_key] if row in lookup][:8000]
    test_rows = [row for row in manifest["test"] if row in lookup]

    x_train = features[[lookup[row] for row in train_rows]]
    x_test = features[[lookup[row] for row in test_rows]]
    report = {"features": args.features, "side": args.side, "targets": {}}
    for name in ("color_presence", "shape_presence"):
        y_train = np.stack([scene_targets(scenes[row])[name] for row in train_rows])
        y_test = np.stack([scene_targets(scenes[row])[name] for row in test_rows])
        report["targets"][name] = evaluate(
            x_train, x_test, y_train, y_test, args.ridge, binary=True
        )
    for name in ("group_count", "object_total", "relation_count"):
        y_train = np.array(
            [[scene_targets(scenes[row])[name]] for row in train_rows]
        )
        y_test = np.array([[scene_targets(scenes[row])[name]] for row in test_rows])
        report["targets"][name] = evaluate(
            x_train, x_test, y_train, y_test, args.ridge, binary=False
        )
    write_json(args.output, report)
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()