"""Configuration, strict case alignment and evidence shared by all entrypoints."""
import hashlib
import json
import os
from pathlib import Path
import socket
import sys

import numpy as np
import pandas as pd
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
import yaml

PACKAGE_ROOT = Path(__file__).resolve().parents[1]
if str(PACKAGE_ROOT) not in sys.path:
    sys.path.insert(0, str(PACKAGE_ROOT))
from model.wavy_channel_rf import WavyChannelRF


def load_config(path):
    config_path = Path(path).resolve()
    with config_path.open() as stream:
        config = yaml.safe_load(stream)
    data = config["data"]
    if tuple(data["features"]) != WavyChannelRF.FEATURES:
        raise ValueError("The seven feature names/order must follow the formal protocol")
    if tuple(data["targets"]) != WavyChannelRF.TARGETS:
        raise ValueError("Target names/order must be Nuavg and DelP_Pa")
    params = config["model"]["rf_params"]
    for key, expected in {"n_estimators": 800, "min_samples_leaf": 2, "random_state": 42}.items():
        if params[key] != expected:
            raise ValueError(f"Formal protocol requires {key}={expected}")
    if not config["model"]["independent_targets"] or not config["experiment"]["full_fit"]:
        raise ValueError("Formal experiment requires two independent full fits")
    if data["normalization"] != "none" or data["target_transform"] != "none":
        raise ValueError("No data/target transformation is configured for this reproduction")
    if config["experiment"]["epochs"] is not None or config["experiment"]["early_stopping"]:
        raise ValueError("Epochs and early stopping do not apply to this Random Forest experiment")
    base = config_path.parent.parent
    data_root = Path(data["root"])
    data["root"] = str(data_root if data_root.is_absolute() else base / data_root)
    for key in ["results_dir", "checkpoint"]:
        p = Path(config["output"][key])
        config["output"][key] = str(p if p.is_absolute() else base / p)
    config["config_path"] = str(config_path)
    return config


def ensure_compute_node(config):
    host = socket.gethostname()
    if config["experiment"]["require_compute_node"]:
        if not os.environ.get("SLURM_JOB_ID") or "login" in host.lower():
            raise RuntimeError(f"Execution requires a SLURM compute node; current host={host}")
    print(f"[runtime] host={host}, SLURM_JOB_ID={os.environ.get('SLURM_JOB_ID')}", flush=True)
    return {"host": host, "slurm_job_id": os.environ.get("SLURM_JOB_ID"),
            "slurm_cpus_per_task": os.environ.get("SLURM_CPUS_PER_TASK")}


def sha256(path):
    digest = hashlib.sha256()
    with Path(path).open("rb") as stream:
        for chunk in iter(lambda: stream.read(1048576), b""):
            digest.update(chunk)
    return digest.hexdigest()


def load_fixed_dataset(config):
    data = config["data"]
    root = Path(data["root"])
    paths = {"dataset": root / data["dataset"], "split_file": root / data["split_file"]}
    key, split_col = data["join_key"], data["split_column"]
    frame = pd.read_csv(paths["dataset"], dtype={key: str})
    split = pd.read_csv(paths["split_file"], dtype={key: str})
    required = [key] + data["features"] + data["targets"]
    missing = set(required) - set(frame.columns)
    if missing or not {key, split_col}.issubset(split.columns):
        raise ValueError(f"CSV schema mismatch, missing dataset columns: {sorted(missing)}")
    if split_col in frame.columns:
        raise ValueError("Dataset must not already contain split labels")
    for name, table in [("dataset", frame), ("split", split)]:
        if table[key].isna().any() or table[key].duplicated().any():
            raise ValueError(f"{name}: case_uid must be nonnull and unique")
    if set(frame[key]) != set(split[key]):
        raise ValueError("Dataset and split case_uid sets differ")
    if split[split_col].isna().any() or set(split[split_col]) != {"train", "test"}:
        raise ValueError("Split labels must be exactly train/test")
    before = frame[key].tolist()
    joined = frame.merge(split[[key, split_col]], on=key, how="left", sort=False, validate="one_to_one")
    if joined[key].tolist() != before:
        raise ValueError("Merge unexpectedly reordered dataset rows")
    for column in data["features"] + data["targets"]:
        joined[column] = pd.to_numeric(joined[column], errors="raise")
    if not np.isfinite(joined[data["features"] + data["targets"]].to_numpy(dtype=float)).all():
        raise ValueError("Dataset contains NaN/Inf; no rows may be silently removed")
    counts = {"total": len(joined), **{s: int((joined[split_col] == s).sum()) for s in ["train", "test"]}}
    if counts != data["expected_counts"]:
        raise ValueError(f"Fixed split counts {counts} != expected {data['expected_counts']}")
    manifest = {"files": {name: {"path": str(path), "sha256": sha256(path)} for name, path in paths.items()},
                "counts": counts, "join_key": key, "features": data["features"], "targets": data["targets"],
                "feature_units": ["1", "1", "1", "1", "mm", "mm", "mm"],
                "target_units": ["1", "Pa"], "row_order": "original dataset order",
                "sampling": "none", "normalization": "none", "target_transform": "none",
                "distinct_input_rows": int(joined[data["features"]].drop_duplicates().shape[0])}
    print(f"[data] matched by {key}: total={counts['total']}, train={counts['train']}, "
          f"test={counts['test']}; all rows retained", flush=True)
    return joined, manifest


def metrics_for(true, predicted, targets):
    true = np.asarray(true, dtype=float)
    predicted = np.asarray(predicted, dtype=float)
    if true.shape != predicted.shape or true.ndim != 2 or true.shape[1] != len(targets):
        raise ValueError("Metric input shape mismatch")
    if not np.isfinite(true).all() or not np.isfinite(predicted).all():
        raise ValueError("Metric input contains NaN/Inf")
    return {target: {"MAE": float(mean_absolute_error(true[:, i], predicted[:, i])),
                     "RMSE": float(np.sqrt(mean_squared_error(true[:, i], predicted[:, i]))),
                     "R2": float(r2_score(true[:, i], predicted[:, i]))}
            for i, target in enumerate(targets)}


def print_metrics(label, metrics, n):
    for target, values in metrics.items():
        print(f"[{label}] {target} n={n}: MAE={values['MAE']:.10f} "
              f"RMSE={values['RMSE']:.10f} R2={values['R2']:.10f}", flush=True)


def write_json(path, value):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, indent=2, ensure_ascii=False, allow_nan=False) + "\n")


def write_metrics(directory, metrics, name="metrics"):
    directory = Path(directory)
    write_json(directory / f"{name}.json", metrics)
    pd.DataFrame([{"target": t, **m} for t, m in metrics.items()]).to_csv(directory / f"{name}.csv", index=False)


def prediction_frame(frame, predicted, config):
    data = config["data"]
    cols = [data["join_key"], data["split_column"]] + data["features"]
    output = frame[cols].reset_index(drop=True).copy()
    for i, target in enumerate(data["targets"]):
        output[f"{target}_true"] = frame[target].to_numpy()
        output[f"{target}_pred"] = predicted[:, i]
        output[f"{target}_residual"] = predicted[:, i] - frame[target].to_numpy()
    return output