"""Run probabilistic DINCAE reconstruction on held-out daily SST fields."""

from datetime import date
from pathlib import Path
import sys

import numpy as np
import torch
import yaml

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from model.dincae import DINCAE, build_input, output_distribution


def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    device = torch.device("cuda" if torch.cuda.is_available() and config["runtime"]["device"] != "cpu" else "cpu")
    data_path = ROOT / config["data"]["root"] / config["data"]["file"]
    with np.load(data_path) as loaded:
        data = {key: loaded[key] for key in loaded.files}
    checkpoint = torch.load(ROOT / config["paths"]["checkpoint"], map_location=device, weights_only=False)
    model = DINCAE(**checkpoint["model_config"]).to(device)
    model.load_state_dict(checkpoint["model"])
    model.eval()
    start = min(int(config["data"]["train_samples"]), len(data["timestamps"]) - 1)
    means, variances, targets, missing_masks, inputs_observed = [], [], [], [], []
    with torch.no_grad():
        for index in range(start, len(data["timestamps"])):
            timestamp = date.fromisoformat(str(data["timestamps"][index]))
            model_input = build_input(data["observed_anomaly"], data["precision"], index,
                                      data["longitude"], data["latitude"], timestamp.timetuple().tm_yday)
            output = model(torch.from_numpy(model_input[None]).to(device))
            mean, variance, _ = output_distribution(output, checkpoint["gamma"], checkpoint["delta"])
            means.append(mean[0].cpu().numpy() + data["climatology"])
            variances.append(variance[0].cpu().numpy())
            targets.append(data["sst_anomaly"][index] + data["climatology"])
            missing_masks.append((data["precision"][index] == 0) & data["ocean_mask"])
            inputs_observed.append(data["observed_anomaly"][index] + data["climatology"])
    output_path = ROOT / config["paths"]["inference"]
    output_path.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(output_path, prediction=np.asarray(means), variance=np.asarray(variances),
                        target=np.asarray(targets), missing_mask=np.asarray(missing_masks),
                        observed=np.asarray(inputs_observed), ocean_mask=data["ocean_mask"],
                        timestamps=data["timestamps"][start:], units=data["units"])
    print(f"predictions={len(means)} shape={np.asarray(means).shape} output={output_path}")


if __name__ == "__main__":
    main()