"""Fit the formal 3686-case experiment on a SLURM compute node."""
import argparse
from datetime import datetime, timezone
import importlib.metadata
import os
from pathlib import Path
import platform
import time

import numpy as np
from common import (load_config, ensure_compute_node, load_fixed_dataset, metrics_for,
                    print_metrics, prediction_frame, write_json, write_metrics, sha256, WavyChannelRF)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", required=True)
    args = parser.parse_args()
    config = load_config(args.config)
    runtime = ensure_compute_node(config)
    directory = Path(config["output"]["results_dir"])
    directory.mkdir(parents=True, exist_ok=True)
    started = time.perf_counter()
    data, data_manifest = load_fixed_dataset(config)
    split_col = config["data"]["split_column"]
    train, test = (data.loc[data[split_col] == split] for split in ["train", "test"])
    features, targets = config["data"]["features"], config["data"]["targets"]
    versions = {p: importlib.metadata.version(p) for p in ["numpy", "pandas", "scikit-learn", "joblib", "PyYAML", "matplotlib"]}
    print("[train] Formal full fit: two independent Random Forests, 800 trees per target; no epochs.", flush=True)
    model = WavyChannelRF(config["model"]["rf_params"])
    model.fit(train[features], train[targets])
    predictions = model.predict(test[features])
    train_predictions = model.predict(train[features])
    metrics = metrics_for(test[targets], predictions, targets)
    train_metrics = metrics_for(train[targets], train_predictions, targets)
    print_metrics("test", metrics, len(test))
    print_metrics("train", train_metrics, len(train))
    manifest = {"experiment": config["experiment"], "utc_time": datetime.now(timezone.utc).isoformat(),
                "runtime": runtime, "python": platform.python_version(), "versions": versions,
                "data": data_manifest, "fit_seconds": model.fit_seconds_,
                "fitted_train_rows": model.n_train_,
                "trees_per_target": {t: len(model.models_[t].estimators_) for t in targets},
                "actual_rf_parameters": {t: model.models_[t].get_params() for t in targets},
                "checkpoint_semantics": "sole fully fitted formal model; no validation/test model selection",
                "serialization": "joblib bundle containing two fitted sklearn forests; .pth is requested filename",
                "sources": config["reference"], "training_script_sha256": sha256(__file__),
                "model_source_sha256": sha256(Path(__file__).resolve().parents[1] / "model/wavy_channel_rf.py")}
    checkpoint = Path(config["output"]["checkpoint"])
    model.save(checkpoint, manifest)
    print(f"[save] Reloadable joblib checkpoint: {checkpoint} ({checkpoint.stat().st_size} bytes)", flush=True)
    write_metrics(directory, metrics)
    write_json(directory / "train_metrics.json", train_metrics)
    prediction_frame(test, predictions, config).to_csv(directory / "test_predictions.csv", index=False)
    model.feature_importances().to_csv(directory / "feature_importances.csv", index=False)
    write_json(directory / "data_manifest.json", data_manifest)
    write_json(directory / "config_snapshot.json", config)
    manifest.update({"test_metrics": metrics, "checkpoint": str(checkpoint), "checkpoint_sha256": sha256(checkpoint),
                     "total_seconds": time.perf_counter() - started, "execution_state": "completed"})
    write_json(directory / "run_manifest.json", manifest)
    print(f"[train] Completed formal experiment in {manifest['total_seconds']:.3f} s", flush=True)


if __name__ == "__main__":
    main()