#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Train a TensorFlow model from a directory of ROOT files.

Inputs:
    Tb
    thetab
    vZ
    MBeam
    MTarget
    MLight
    MHeavy

Outputs:
    beamEnergy
    Ex
"""
import numpy as np
from pathlib import Path

import joblib
import pandas as pd
import uproot

from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

import tensorflow as tf
from tensorflow.keras import Sequential
from tensorflow.keras.layers import Dense
from tensorflow.keras.callbacks import EarlyStopping

ROOT_FOLDER = Path("/Users/jamesszalkie/ANASEN_analysis/Armory/Training_Data_Ne/")

TREE_NAME = "tree1"

INPUT_BRANCHES = [
    "Tb",
    "thetab",
    "vZ",
    "MBeam",
    "MTarget",
    "MLight",
    "MHeavy",
]

OUTPUT_BRANCHES = [
    "beamEnergy",
    "Ex",
]

MODEL_NAME = "beam_predictor.keras"

print("Reading ROOT files...")

dfs = []

for root_file in sorted(ROOT_FOLDER.glob("*.root")):

    print(f"  {root_file.name}")

    with uproot.open(root_file) as f:

        tree = f[TREE_NAME]

        arrays = tree.arrays(
            INPUT_BRANCHES + OUTPUT_BRANCHES,
            library="np"
        )

        dfs.append(pd.DataFrame(arrays))

dataset = pd.concat(dfs, ignore_index=True)

# Remove events with missing or invalid values
dataset = dataset.replace([np.inf, -np.inf], np.nan)
dataset = dataset.dropna()

print(f"Training on {len(dataset)} complete events.")

print(dataset.describe())

print()

print(dataset.isna().sum())

print()

print(np.isinf(dataset).sum())

print()
print("Total events:", len(dataset))

X = dataset[INPUT_BRANCHES].values
Y = dataset[OUTPUT_BRANCHES].values

X_train, X_test, Y_train, Y_test = train_test_split(
    X,
    Y,
    test_size=0.20,
    random_state=42,
)

input_scaler = StandardScaler()
output_scaler = StandardScaler()

X_train = input_scaler.fit_transform(X_train)
X_test = input_scaler.transform(X_test)

Y_train = output_scaler.fit_transform(Y_train)
Y_test = output_scaler.transform(Y_test)

model = Sequential([
    Dense(128, activation="relu"),
    Dense(128, activation="relu"),
    Dense(64, activation="relu"),
    Dense(32, activation="relu"),
    Dense(2)
])

model.build((None, len(INPUT_BRANCHES)))

model.compile(
    optimizer="adam",
    loss="mse",
    metrics=["mae"]
)

model.summary()

early_stop = EarlyStopping(
    monitor="val_loss",
    patience=20,
    restore_best_weights=True
)

history = model.fit(
    X_train,
    Y_train,
    epochs=25,
    batch_size=512,
    validation_split=0.20,
    callbacks=[early_stop],
    verbose=1,
)

loss, mae = model.evaluate(
    X_test,
    Y_test,
    verbose=0,
)

pred_scaled = model.predict(X_test, verbose=0)

pred = output_scaler.inverse_transform(pred_scaled)
truth = output_scaler.inverse_transform(Y_test)

print()
print(f"Test Loss : {loss:.6f}")
print(f"Test MAE  : {mae:.6f}")

beam_error = np.mean(np.abs(pred[:,0] - truth[:,0]))
Ex_error = np.mean(np.abs(pred[:,1] - truth[:,1]))

print(f"Beam Energy MAE: {beam_error:.6f} MeV")
print(f"Ex MAE: {Ex_error:.6f} MeV")

model.save(MODEL_NAME)

joblib.dump(input_scaler, "input_scaler.pkl")
joblib.dump(output_scaler, "output_scaler.pkl")

print()
print("Training complete.")
print("Saved:")
print("   ", MODEL_NAME)
print("    input_scaler.pkl")
print("    output_scaler.pkl")