"""
01 — Train the canonical POD-NN surrogate.

Trains on the first TRAIN_SIZE=100 samples, holds out the last
TEST_SIZE=400 for evaluation, retains N_MODES=4 POD modes.
Saves trained artifacts to 2_AnalyzedData/trained_model/ for use by
the plotting scripts (02, 05, 07, 08).
"""

from sklearn.metrics import r2_score, mean_squared_error, mean_absolute_error

from _utils import (
    load_data, train_pod_nn, predict_field, save_artifacts, seed_all,
    TRAIN_SIZE, TEST_SIZE, N_MODES, EPOCHS,
)


def main():
    seed_all()
    print(f"Training POD-NN  |  N_train={TRAIN_SIZE}  N_test={TEST_SIZE}  "
          f"r={N_MODES}  epochs<={EPOCHS}")

    parameters, fields, _, _ = load_data()
    X_train = parameters[:TRAIN_SIZE]
    Y_train = fields[:TRAIN_SIZE]
    X_test  = parameters[-TEST_SIZE:]
    Y_test  = fields[-TEST_SIZE:]

    trained = train_pod_nn(X_train, Y_train, n_modes=N_MODES, epochs=EPOCHS)

    var = trained["pca"].explained_variance_ratio_.sum() * 100
    print(f"  POD variance explained: {var:.2f}%")
    print(f"  Stopped at epoch:       {len(trained['history']['loss'])}")

    Y_pred = predict_field(trained, X_test)
    r2  = r2_score(Y_test.flatten(), Y_pred.flatten())
    rmse = mean_squared_error(Y_test.flatten(), Y_pred.flatten()) ** 0.5
    mae = mean_absolute_error(Y_test.flatten(), Y_pred.flatten())
    print(f"\nTest set metrics ({TEST_SIZE} samples):")
    print(f"  R2   = {r2:.6f}")
    print(f"  RMSE = {rmse:.4f} K")
    print(f"  MAE  = {mae:.4f} K")

    save_artifacts(trained)


if __name__ == "__main__":
    main()