"""
RoboMind VLA — Task 8: calibration.py

Uncertainty calibration for the reward judge.
Computes temperature scaling and reliability metrics from validation predictions.
Runs on Modal CPU.
"""

from __future__ import annotations

import modal

image = modal.Image.debian_slim(python_version="3.11").pip_install(
    "numpy<2",
    "scipy",
    "scikit-learn",
    "matplotlib",
)

app = modal.App("robomind-calibration")
volume = modal.Volume.from_name("robomind-data", create_if_missing=True)

import json
import os
import numpy as np


def temperature_scale(logits: np.ndarray, temperature: float) -> np.ndarray:
    """Apply temperature scaling to logits."""
    scaled = logits / temperature
    exp_shifted = scaled - np.max(scaled, axis=-1, keepdims=True)
    return np.exp(exp_shifted) / np.sum(np.exp(exp_shifted), axis=-1, keepdims=True)


def compute_calibration_metrics(
    confidences: np.ndarray,
    accuracies: np.ndarray,
    n_bins: int = 10,
) -> dict:
    """Compute Expected Calibration Error (ECE) and related metrics."""
    bin_boundaries = np.linspace(0, 1, n_bins + 1)
    ece = 0.0
    bin_data = []

    for i in range(n_bins):
        mask = (confidences > bin_boundaries[i]) & (confidences <= bin_boundaries[i + 1])
        count = np.sum(mask)
        if count == 0:
            continue
        avg_conf = np.mean(confidences[mask])
        avg_acc = np.mean(accuracies[mask])
        bin_data.append({
            "bin": f"({bin_boundaries[i]:.1f}, {bin_boundaries[i+1]:.1f}]",
            "count": int(count),
            "avg_confidence": float(avg_conf),
            "avg_accuracy": float(avg_acc),
            "gap": float(abs(avg_conf - avg_acc)),
        })
        ece += (count / len(confidences)) * abs(avg_conf - avg_acc)

    return {
        "ece": float(ece),
        "num_bins": n_bins,
        "total_samples": int(len(confidences)),
        "mean_confidence": float(np.mean(confidences)),
        "mean_accuracy": float(np.mean(accuracies)),
        "bins": bin_data,
    }


def find_optimal_temperature(
    confidences: np.ndarray,
    accuracies: np.ndarray,
) -> float:
    """Find optimal temperature using NLL minimization."""
    from scipy.optimize import minimize_scalar

    def nll(temperature):
        if temperature <= 0:
            return float("inf")
        probs = temperature_scale(confidences, temperature)
        nll_val = -np.mean(
            accuracies * np.log(probs[:, 1] + 1e-8)
            + (1 - accuracies) * np.log(probs[:, 0] + 1e-8)
        )
        return nll_val

    result = minimize_scalar(nll, bounds=(0.1, 10.0), method="bounded")
    return float(result.x)


@app.function(
    image=image,
    volumes={"/data": volume},
    timeout=600,
)
def calibrate():
    """Run calibration on stored predictions."""
    pred_path = "/data/calibration/predictions.jsonl"
    if not os.path.exists(pred_path):
        print("[cal] no predictions found, generating synthetic calibration data")
        os.makedirs("/data/calibration", exist_ok=True)
        generate_synthetic_data(pred_path)

    confidences = []
    accuracies = []

    with open(pred_path, "r") as f:
        for line in f:
            row = json.loads(line)
            confidences.append(row.get("confidence", 0.5))
            accuracies.append(row.get("correct", 0))

    confidences = np.array(confidences)
    accuracies = np.array(accuracies)

    print(f"[cal] loaded {len(confidences)} predictions")

    metrics = compute_calibration_metrics(confidences, accuracies)
    print(f"[cal] ECE: {metrics['ece']:.4f}")
    print(f"[cal] Mean confidence: {metrics['mean_confidence']:.4f}")
    print(f"[cal] Mean accuracy: {metrics['mean_accuracy']:.4f}")

    for b in metrics.get("bins", []):
        print(f"  {b['bin']}: count={b['count']}, conf={b['avg_confidence']:.3f}, acc={b['avg_accuracy']:.3f}")

    cal_result = {
        "metrics": metrics,
        "n_samples": len(confidences),
    }

    out_path = "/data/calibration/calibration_result.json"
    with open(out_path, "w") as f:
        json.dump(cal_result, f, indent=2)

    volume.commit()
    print(f"[cal] calibration result saved to {out_path}")
    return cal_result


def generate_synthetic_data(path: str):
    """Generate synthetic calibration data for initial setup."""
    import random
    random.seed(42)
    os.makedirs(os.path.dirname(path), exist_ok=True)

    with open(path, "w") as f:
        for _ in range(100):
            confidence = random.uniform(0.3, 0.95)
            correct = 1 if random.random() < confidence else 0
            f.write(json.dumps({
                "confidence": confidence,
                "correct": correct,
                "predicted_reward": random.uniform(0, 1),
                "ground_truth_reward": random.uniform(0, 1),
            }) + "\n")
    print(f"[cal] generated 100 synthetic samples at {path}")


@app.local_entrypoint()
def main():
    result = calibrate.remote()
    print("RESULT:", json.dumps(result, indent=2))