"""Fit statistical CME networks and the precipitation-constraint GP."""

import json
import sys
from pathlib import Path

import numpy as np
import torch
import yaml


ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from model.causalmodelevaluation import (FORMAT_VERSION, LaggedPartialCorrelationCME,
                                         PrecipitationConstraintGP, asymmetric_f1)


def main():
    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
    np.random.seed(int(config["seed"]))
    data = np.load(ROOT / config["paths"]["dataset"])
    if str(data["format_version"]) != config["data"]["format_version"]:
        raise ValueError("incompatible synthetic data format")
    estimator = LaggedPartialCorrelationCME(**config["model"])
    references = [estimator.fit(series) for series in data["reference_series"]]
    model_networks, scores = [], []
    for model_series in data["model_series"]:
        networks = [estimator.fit(series) for series in model_series]
        model_networks.append(networks)
        scores.append(float(np.mean([asymmetric_f1(ref, net, int(config["evaluation"]["lag_tolerance"]))["f1"]
                                     for ref, net in zip(references, networks)])))
    gp = PrecipitationConstraintGP(int(config["seed"])).fit(np.asarray(scores), data["delta_precipitation"])
    checkpoint = {
        "version": FORMAT_VERSION, "config": config, "network_config": estimator.config(),
        "reference_networks": [network.state_dict() for network in references],
        "model_networks": [[network.state_dict() for network in networks] for networks in model_networks],
        "model_f1": torch.tensor(scores), "gp": gp.state_dict(),
        "metadata": {"method": "lagged target-history conditional regression ParCorr approximation",
                     "gradient_training": False, "tensor_layout": "source,target,lag", "time_step_days": 3}
    }
    checkpoint_path = ROOT / config["paths"]["checkpoint"]
    metrics_path = ROOT / config["paths"]["training_metrics"]
    checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
    metrics_path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(checkpoint, checkpoint_path)
    metrics_path.write_text(json.dumps({"model_f1": scores, "gp_kernel": str(gp.model.kernel_),
                                        "paper_alpha": config["paper_model"]["alpha"],
                                        "engineering_alpha": config["model"]["alpha"]}, indent=2) + "\n")
    print(f"checkpoint={checkpoint_path.relative_to(ROOT)} networks={len(references) * (1 + len(model_networks))} gp_samples={len(scores)}")


if __name__ == "__main__":
    main()