"""Tabular baseline models on the publication split.

Trains decision tree, random forest, gradient boosting, SVM and symbolic
regression (GEP-style) baselines on the same seed-17 manifest split and the
same tabular feature vector used by ``TabularANN`` (raw mix columns, rich
maskable columns with observed-mask channels, one-hot categoricals). Each
target (compressive, tensile, flexural) is fitted on the supervised train
rows and scored on the supervised test rows, matching how the GNN's
per-target metrics are computed.

The symbolic-regression baseline uses only the raw mix columns: genetic
program search does not scale to the full masked feature block, and published
GEP concrete models use compact mix-level inputs.
"""

from __future__ import annotations

import argparse
import sys
import time
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.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)


def build_models():
    return {
        "decision_tree": lambda: DecisionTreeRegressor(random_state=SEED),
        "random_forest": lambda: RandomForestRegressor(
            n_estimators=300, random_state=SEED, n_jobs=-1
        ),
        "gradient_boosting": lambda: HistGradientBoostingRegressor(random_state=SEED),
        "svm": lambda: make_pipeline(
            StandardScaler(), SVR(C=100.0, gamma="scale", epsilon=0.5)
        ),
        "gep_symbolic": lambda: SymbolicRegressor(
            population_size=1000,
            generations=20,
            tournament_size=20,
            function_set=("add", "sub", "mul", "div", "sqrt", "log"),
            parsimony_coefficient=0.001,
            random_state=SEED,
            n_jobs=-1,
        ),
    }


def metrics(y_true: np.ndarray, y_pred: np.ndarray):
    err = y_pred - y_true
    rmse = float(np.sqrt(np.mean(err**2)))
    mae = float(np.mean(np.abs(err)))
    ss_res = float(np.sum(err**2))
    ss_tot = float(np.sum((y_true - y_true.mean()) ** 2))
    return rmse, mae, 1.0 - ss_res / ss_tot


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")
    args = parser.parse_args()

    df = read_table(args.data)
    train_idx, val_idx, test_idx = load_split_indices(args.split, len(df))
    # Baselines have no validation-based early stopping; fold validation 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)}")

    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)
    records = []
    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, factory in build_models().items():
            cols = slice(0, N_RAW) if name == "gep_symbolic" else slice(None)
            model = factory()
            start = time.time()
            model.fit(X[tr, cols], y[tr])
            pred = model.predict(X[te, cols])
            rmse, mae, r2 = metrics(y[te], pred)
            print(f"  {name:18s} RMSE={rmse:7.3f} MAE={mae:7.3f} R2={r2:.4f} "
                  f"({time.time() - start:.0f}s)")
            records.append(
                {"model": name, "target": target, "n_train": len(tr),
                 "n_test": len(te), "rmse": rmse, "mae": mae, "r2": r2}
            )
            if name == "gep_symbolic":
                with open(out_dir / f"gep_program_{target}.txt", "w") as fh:
                    fh.write(str(model._program))

    table = pd.DataFrame(records)
    table.to_csv(out_dir / "baseline_metrics.csv", index=False)
    print(f"\nsaved {out_dir / 'baseline_metrics.csv'}")


if __name__ == "__main__":
    main()
