"""Reload the saved model in a fresh process and predict all 922 test cases."""
import argparse
from pathlib import Path
import time

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


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", required=True)
    parser.add_argument("--checkpoint", help="Optional explicit path to saved joblib .pth bundle")
    args = parser.parse_args()
    config = load_config(args.config)
    runtime = ensure_compute_node(config)
    directory = Path(config["output"]["results_dir"])
    data, data_manifest = load_fixed_dataset(config)
    test = data.loc[data[config["data"]["split_column"]] == "test"]
    features, targets = config["data"]["features"], config["data"]["targets"]
    checkpoint = Path(args.checkpoint or config["output"]["checkpoint"])
    print(f"[inference] Loading {checkpoint}", flush=True)
    model, metadata = WavyChannelRF.load(checkpoint)
    if model.n_train_ != config["data"]["expected_counts"]["train"]:
        raise ValueError("Checkpoint did not fit the required complete training set")
    for name in ["dataset", "split_file"]:
        if metadata["data"]["files"][name]["sha256"] != data_manifest["files"][name]["sha256"]:
            raise ValueError(f"{name} changed since model fitting")
    start = time.perf_counter()
    predicted = model.predict(test[features])
    seconds = time.perf_counter() - start
    metrics = metrics_for(test[targets], predicted, targets)
    print_metrics("inference-test", metrics, len(test))
    output = prediction_frame(test, predicted, config)
    output.to_csv(directory / "inference_predictions.csv", index=False)
    write_metrics(directory, metrics, "inference_metrics")
    reference = pd.read_csv(directory / "test_predictions.csv")
    key = config["data"]["join_key"]
    if reference[key].tolist() != output[key].tolist():
        raise ValueError("Reloaded inference test case order differs from training evaluation")
    pred_cols = [f"{t}_pred" for t in targets]
    error = float(np.max(np.abs(reference[pred_cols].to_numpy() - output[pred_cols].to_numpy())))
    np.testing.assert_allclose(reference[pred_cols], output[pred_cols], atol=1e-10, rtol=1e-12)
    verification = {"status": "passed", "test_rows": len(test), "predictions_finite": True,
                    "independent_process_reload": True, "case_uids_equal": True,
                    "max_abs_prediction_difference": error, "tolerance": {"atol": 1e-10, "rtol": 1e-12},
                    "prediction_seconds": seconds, "runtime": runtime,
                    "checkpoint_sha256": sha256(checkpoint), "test_metrics": metrics}
    write_json(directory / "reload_validation.json", verification)
    print(f"[inference] Reload validation passed: n={len(test)}, max_abs_difference={error:.3g}", flush=True)


if __name__ == "__main__":
    main()