# /// script
# requires-python = ">=3.12,<3.13"
# dependencies = ["pymc==5.21.1", "pytensor==2.28.3", "arviz==0.23.4"]
# ///

import csv
import json
import os

os.environ["PYTENSOR_FLAGS"] = "cxx=,mode=FAST_COMPILE"

import numpy as np
import pymc as pm
import pytensor.tensor as pt


def score_summary(loglik, batch_length=20):
    """Total log posterior-predictive density and batch-means MC error.

    Keep chains separate when forming consecutive batches. Each draw's
    influence sums over observations before its variance is estimated.
    """
    x = np.asarray(loglik, dtype=float)
    if x.ndim != 3 or not np.isfinite(x).all():
        raise ValueError("Expected finite chain-by-draw-by-observation log densities")
    chains, draws, _ = x.shape
    if draws % batch_length or draws // batch_length < 2:
        raise ValueError("Each chain must contain at least two complete batches")
    shift = x.max(axis=(0, 1))
    density = np.exp(x - shift)
    means = density.mean(axis=(0, 1))
    pointwise = shift + np.log(means)
    influence = (density / means - 1).sum(axis=2)
    batches = influence.reshape(chains, draws // batch_length, batch_length).mean(axis=2)
    mcse = float(batches.std(ddof=1) / np.sqrt(batches.size))
    return dict(score=float(pointwise.sum()), mcse=mcse,
                pointwise=pointwise, influence=influence,
                batch_length=batch_length, batches=int(batches.size))


def main():
    from dugong_model import model

    with open("observations.csv", newline="") as stream:
        rows = list(csv.DictReader(stream))
    visible = [row for row in rows if row["split"] == "visible"]
    heldout = [row for row in rows if row["split"] == "heldout"]
    if len(visible) != 80 or len(heldout) != 20:
        raise ValueError("Expected 80 visible and 20 held-out observations")
    data = {name: np.array([float(row[name]) for row in visible]) for name in ("age", "length")}
    with model(data):
        fit = pm.sample(
            draws=200, tune=200, chains=1, cores=1,
            progressbar=False, compute_convergence_checks=False,
            random_seed=0, init="adapt_diag", jitter_max_retries=1, n_init=100,
            idata_kwargs={"log_likelihood": True}, compile_kwargs={"mode": "FAST_COMPILE"},
        )
    print(fit.posterior.isel(chain=0, draw=0))
    print(fit.posterior.mean())

    ages = np.array([float(row["age"]) for row in heldout])
    outcomes = np.array([float(row["length"]) for row in heldout])
    # The model is built without the held-out outcomes. Bind them only for scoring.
    prediction_model = model({"age": ages, "length": np.full(20, data["length"].mean())})
    target = prediction_model.observed_RVs[0]
    original_value = prediction_model.rvs_to_values[target]
    prediction_model.rvs_to_values[target] = pt.as_tensor_variable(outcomes, name=original_value.name)
    likelihood = pm.compute_log_likelihood(
        fit, model=prediction_model, extend_inferencedata=False,
        progressbar=False, compile_kwargs={"mode": "FAST_COMPILE"},
    )[target.name].values
    report = score_summary(likelihood)

    repeated = fit.posterior.isel(draw=np.repeat(np.arange(200), 20)).assign_coords(draw=np.arange(4000))
    with prediction_model:
        predictions = pm.sample_posterior_predictive(
            repeated, var_names=[target.name], random_seed=314,
            return_inferencedata=False, progressbar=False,
            compile_kwargs={"mode": "FAST_COMPILE"},
        )[target.name].reshape(4000, 20)
    q05, median, q95 = np.quantile(predictions, [0.05, 0.5, 0.95], axis=0)
    covered = (q05 <= outcomes) & (outcomes <= q95)
    fields = ["row", "age", "observed_length", "predictive_q05", "predictive_median",
              "predictive_q95", "covered", "log_predictive_density", "program_log_score", "program_mcse"]
    with open("dugongs-heldout-scored.csv", "w", newline="") as stream:
        writer = csv.writer(stream)
        writer.writerow(fields)
        for i in range(20):
            writer.writerow([i, ages[i], outcomes[i], q05[i], median[i], q95[i], bool(covered[i]),
                             report["pointwise"][i], report["score"], report["mcse"]])
    output = {key: report[key] for key in ("score", "mcse", "batch_length", "batches")}
    output.update(covered=int(covered.sum()), heldout_observations=20,
                  predictive_interval=0.9, predictive_seed=314, predictive_replicates_per_draw=20)
    with open("dugongs-score.json", "w") as stream:
        json.dump(output, stream, indent=2, allow_nan=False)
        stream.write("\n")
    print(f"Held-out log score: {report['score']:.4f}")
    print(f"90% predictive interval coverage: {covered.sum()} of 20")
    print(f"Monte Carlo standard error: {report['mcse']:.4f}")


if __name__ == "__main__":
    main()