"""Retrain the tabular baselines and record per-epoch loss histories.

Companion to ``run_baselines.py`` (same manifest split, same feature vector,
same hyperparameters): retrains each baseline while recording train/test RMSE
on a common 1..120 "epoch" axis so all models can share one training-curve
figure with the GNN/ANN (whose per-epoch losses live in outputs/train_log.txt,
parsed here into the same folder).

None of these model families trains by literal epochs, so each is mapped onto
the 120-step axis through its own native training/complexity iterator (kept in
the ``native_iteration`` column):

* decision_tree      -- tree grown to max_depth = epoch (saturates at the
                        natural depth; the final tree equals the published
                        unrestricted one).
* random_forest      -- cumulative ensemble mean over the first
                        ceil(300*epoch/120) of the 300 published trees
                        (epoch 120 == the published forest exactly).
* gradient_boosting  -- boosting iteration (max_iter=120; the published model
                        used the default 100, so epoch 100 matches it).
* svm                -- SMO optimisation budget: refit with
                        max_iter = ceil(n_iter_full * epoch / 120)
                        (epoch 120 == the fully converged published SVR).
* gep_symbolic       -- generation, evolved incrementally via warm_start to
                        120 generations (the published run stopped at 20; the
                        warm-start RNG stream differs slightly from a fresh
                        20-generation run).

Outputs (never overwrites the published metrics):
    outputs/baselines/history/{model}_{target}_history.csv
    outputs/baselines/history/{hicon_gnn,tabular_ann}_history.csv (from log)
    outputs/baselines/baseline_loss_histories.csv   (combined tidy table)
    outputs/baselines/baseline_history_final_metrics.csv
"""

from __future__ import annotations

import argparse
import math
import re
import sys
import time
import warnings
from pathlib import Path

import numpy as np
import pandas as pd

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

from gplearn.genetic import SymbolicRegressor
from sklearn.ensemble import HistGradientBoostingRegressor, RandomForestRegressor
from sklearn.exceptions import ConvergenceWarning
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.svm import SVR
from sklearn.tree import DecisionTreeRegressor

from concrete_gnn import load_split_indices, read_table
from concrete_gnn.real_data import RAW_MIX_COLUMNS, raw_mix_vector

TARGETS = {
    "compressive": "compressive_strength_mpa",
    "tensile": "tensile_strength_mpa",
    "flexural": "flexural_strength_mpa",
}

SEED = 17
N_RAW = len(RAW_MIX_COLUMNS)
RF_TREES = 300


def rmse(y_true: np.ndarray, y_pred: np.ndarray) -> float:
    return float(np.sqrt(np.mean((y_pred - y_true) ** 2)))


def _history_rows(epochs, native, tr_rmse, te_rmse):
    return [
        {
            "epoch": e,
            "native_iteration": n,
            "train_rmse": tr,
            "test_rmse": te,
        }
        for e, n, tr, te in zip(epochs, native, tr_rmse, te_rmse)
    ]


def decision_tree_history(Xtr, ytr, Xte, yte, epochs: int):
    rows = []
    for e in range(1, epochs + 1):
        model = DecisionTreeRegressor(max_depth=e, random_state=SEED)
        model.fit(Xtr, ytr)
        rows += _history_rows(
            [e], [e], [rmse(ytr, model.predict(Xtr))], [rmse(yte, model.predict(Xte))]
        )
    return rows, model


def random_forest_history(Xtr, ytr, Xte, yte, epochs: int):
    model = RandomForestRegressor(n_estimators=RF_TREES, random_state=SEED, n_jobs=-1)
    model.fit(Xtr, ytr)
    # Per-tree predictions once; cumulative means give the ensemble-growth curve.
    tr_preds = np.stack([t.predict(Xtr) for t in model.estimators_])
    te_preds = np.stack([t.predict(Xte) for t in model.estimators_])
    tr_cum = np.cumsum(tr_preds, axis=0) / np.arange(1, RF_TREES + 1)[:, None]
    te_cum = np.cumsum(te_preds, axis=0) / np.arange(1, RF_TREES + 1)[:, None]
    rows = []
    for e in range(1, epochs + 1):
        k = math.ceil(RF_TREES * e / epochs)
        rows += _history_rows(
            [e], [k], [rmse(ytr, tr_cum[k - 1])], [rmse(yte, te_cum[k - 1])]
        )
    return rows, model


def gradient_boosting_history(Xtr, ytr, Xte, yte, epochs: int):
    model = HistGradientBoostingRegressor(random_state=SEED, max_iter=epochs)
    model.fit(Xtr, ytr)
    rows = []
    stages = zip(model.staged_predict(Xtr), model.staged_predict(Xte))
    for e, (ptr, pte) in enumerate(stages, start=1):
        rows += _history_rows([e], [e], [rmse(ytr, ptr)], [rmse(yte, pte)])
    return rows, model


def _svr_pipeline(max_iter: int = -1):
    return make_pipeline(
        StandardScaler(), SVR(C=100.0, gamma="scale", epsilon=0.5, max_iter=max_iter)
    )


def svm_history(Xtr, ytr, Xte, yte, epochs: int):
    full = _svr_pipeline()
    full.fit(Xtr, ytr)
    n_iter_full = int(np.ravel(full.named_steps["svr"].n_iter_)[0])
    rows = []
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", ConvergenceWarning)
        for e in range(1, epochs + 1):
            budget = max(1, math.ceil(n_iter_full * e / epochs))
            model = full if budget >= n_iter_full else _svr_pipeline(budget)
            if model is not full:
                model.fit(Xtr, ytr)
            rows += _history_rows(
                [e],
                [budget],
                [rmse(ytr, model.predict(Xtr))],
                [rmse(yte, model.predict(Xte))],
            )
    return rows, full


def gep_history(Xtr, ytr, Xte, yte, epochs: int):
    model = SymbolicRegressor(
        population_size=1000,
        generations=1,
        tournament_size=20,
        function_set=("add", "sub", "mul", "div", "sqrt", "log"),
        parsimony_coefficient=0.001,
        random_state=SEED,
        n_jobs=-1,
        warm_start=True,
    )
    rows = []
    for g in range(1, epochs + 1):
        model.set_params(generations=g)
        model.fit(Xtr, ytr)
        rows += _history_rows(
            [g], [g], [rmse(ytr, model.predict(Xtr))], [rmse(yte, model.predict(Xte))]
        )
    return rows, model


BUILDERS = {
    "decision_tree": decision_tree_history,
    "random_forest": random_forest_history,
    "gradient_boosting": gradient_boosting_history,
    "svm": svm_history,
    "gep_symbolic": gep_history,
}

LOG_LINE = re.compile(
    r"\[(hierarchical|tabular_ann)\] epoch=(\d+) train_loss=([\d.]+) "
    r"val_rmse=([\d.]+) val_r2=([-\d.]+)"
)
LOG_MODEL_NAMES = {"hierarchical": "hicon_gnn", "tabular_ann": "tabular_ann"}


def parse_train_log(log_path: Path, out_dir: Path) -> None:
    """Extract the GNN/ANN per-epoch losses already recorded in the train log."""
    if not log_path.exists():
        print(f"train log not found at {log_path}; skipping GNN/ANN histories")
        return
    rows = {name: [] for name in LOG_MODEL_NAMES.values()}
    for m in LOG_LINE.finditer(log_path.read_text()):
        rows[LOG_MODEL_NAMES[m.group(1)]].append(
            {
                "epoch": int(m.group(2)),
                "train_loss_standardized": float(m.group(3)),
                "val_rmse": float(m.group(4)),
                "val_r2": float(m.group(5)),
            }
        )
    for name, recs in rows.items():
        if recs:
            path = out_dir / f"{name}_history.csv"
            pd.DataFrame(recs).to_csv(path, index=False)
            print(f"saved {path} ({len(recs)} epochs, parsed from {log_path.name})")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--data", default="data/HiConGNN_processed_database.csv")
    parser.add_argument("--split", default="data/split_random_seed17.csv")
    parser.add_argument("--out", default="outputs/baselines")
    parser.add_argument("--train-log", default="outputs/train_log.txt")
    parser.add_argument("--epochs", type=int, default=120)
    args = parser.parse_args()

    df = read_table(args.data)
    train_idx, val_idx, test_idx = load_split_indices(args.split, len(df))
    # Match run_baselines.py: no validation-based stopping, fold val into train.
    train_idx = list(train_idx) + list(val_idx)

    X = np.stack([raw_mix_vector(row).numpy() for _, row in df.iterrows()])
    print(f"rows={len(df)} features={X.shape[1]} (raw mix={N_RAW}) "
          f"train={len(train_idx)} test={len(test_idx)} epochs={args.epochs}")

    hist_dir = Path(args.out) / "history"
    hist_dir.mkdir(parents=True, exist_ok=True)

    combined = []
    finals = []
    for target, column in TARGETS.items():
        y = pd.to_numeric(df[column], errors="coerce").to_numpy(dtype=float)
        tr = [i for i in train_idx if np.isfinite(y[i])]
        te = [i for i in test_idx if np.isfinite(y[i])]
        print(f"\n[{target}] train n={len(tr)} test n={len(te)}")
        for name, builder in BUILDERS.items():
            cols = slice(0, N_RAW) if name == "gep_symbolic" else slice(None)
            start = time.time()
            rows, model = builder(X[tr, cols], y[tr], X[te, cols], y[te], args.epochs)
            hist = pd.DataFrame(rows)
            hist.to_csv(hist_dir / f"{name}_{target}_history.csv", index=False)
            last = hist.iloc[-1]
            print(f"  {name:18s} final train_rmse={last.train_rmse:7.3f} "
                  f"test_rmse={last.test_rmse:7.3f} ({time.time() - start:.0f}s)")
            hist.insert(0, "target", target)
            hist.insert(0, "model", name)
            combined.append(hist)
            finals.append(
                {"model": name, "target": target, "n_train": len(tr),
                 "n_test": len(te), "final_train_rmse": float(last.train_rmse),
                 "final_test_rmse": float(last.test_rmse)}
            )

    out_dir = Path(args.out)
    pd.concat(combined, ignore_index=True).to_csv(
        out_dir / "baseline_loss_histories.csv", index=False
    )
    pd.DataFrame(finals).to_csv(
        out_dir / "baseline_history_final_metrics.csv", index=False
    )
    print(f"\nsaved {out_dir / 'baseline_loss_histories.csv'}")
    print(f"saved {out_dir / 'baseline_history_final_metrics.csv'}")

    parse_train_log(Path(args.train_log), hist_dir)


if __name__ == "__main__":
    main()
