#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Aug  5 15:52:39 2026

@author: jamesszalkie
"""

import uproot
import joblib
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
from sklearn.metrics import mean_absolute_error

# User settings

ROOT_FILE = "/Users/jamesszalkie/ANASEN_analysis/Armory/Ne18Protons.root"
TREE_NAME = "tree1"

MODEL = "beam_predictor.keras"
INPUT_SCALER = "input_scaler.pkl"
OUTPUT_SCALER = "output_scaler.pkl"

# Candidate excitation energies (MeV) for nearest-state snapping
EX_CANDIDATES = np.array([0, 0.3, 1.7, 2.4, 2.8], dtype=np.float32)

INPUT_BRANCHES = [
    "Tb",
    "thetab",
    "vZ",
    "MBeam",
    "MTarget",
    "MLight",
    "MHeavy",
    "ZBeam",
    "ZHeavy",
    "ZLight",
    "ZTarget"
]

# Load model

model = tf.keras.models.load_model(MODEL)

input_scaler = joblib.load(INPUT_SCALER)
output_scaler = joblib.load(OUTPUT_SCALER)

# Read ROOT file

with uproot.open(ROOT_FILE) as f:
    tree = f[TREE_NAME]
    arrays = tree.arrays(INPUT_BRANCHES, library="np")

X = np.column_stack([arrays[b] for b in INPUT_BRANCHES])

# Remove NaN events
mask = np.all(np.isfinite(X), axis=1)
X = X[mask]

print(f"Predicting {len(X)} events")

# Predict

X_scaled = input_scaler.transform(X)

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

pred = output_scaler.inverse_transform(pred_scaled)

beam = pred[:,0]
Ex = pred[:,1]

# Snap each predicted excitation to the nearest candidate energy.
nearest_idx = np.argmin(np.abs(Ex[:, None] - EX_CANDIDATES[None, :]), axis=1)
Ex_snapped = EX_CANDIDATES[nearest_idx]

# See whether truth branches exist

truth_beam = None
truth_Ex = None

with uproot.open(ROOT_FILE) as f:

    tree = f[TREE_NAME]
    branches = tree.keys()

    if "beamEnergy" in branches:
        truth_beam = tree["beamEnergy"].array(library="np")[mask]

    if "Ex" in branches:
        truth_Ex = tree["Ex"].array(library="np")[mask]
        
print("\n============================")
print("Model Performance")
print("============================")

if truth_beam is not None:
    beam_mae = mean_absolute_error(truth_beam, beam)
    beam_rmse = np.sqrt(np.mean((truth_beam - beam)**2))

    print(f"Beam Energy MAE : {beam_mae:.4f} MeV")
    print(f"Beam Energy RMSE: {beam_rmse:.4f} MeV")

if truth_Ex is not None:
    Ex_mae = mean_absolute_error(truth_Ex, Ex)
    Ex_rmse = np.sqrt(np.mean((truth_Ex - Ex)**2))

    Ex_snap_mae = mean_absolute_error(truth_Ex, Ex_snapped)
    Ex_snap_rmse = np.sqrt(np.mean((truth_Ex - Ex_snapped)**2))

    print(f"Excitation MAE  : {Ex_mae:.4f} MeV")
    print(f"Excitation RMSE : {Ex_rmse:.4f} MeV")
    print(f"Ex Snapped MAE  : {Ex_snap_mae:.4f} MeV")
    print(f"Ex Snapped RMSE : {Ex_snap_rmse:.4f} MeV")

print("\nEx candidate states (MeV):", ", ".join(f"{x:.3f}" for x in EX_CANDIDATES))

# Plot Beam Energy
plt.figure(figsize=(8,6))

plt.hist(
    beam,
    bins=250,
    histtype="step",
    linewidth=2,
    label="Predicted",
)

if truth_beam is not None:
    plt.hist(
        truth_beam,
        bins=250,
        histtype="step",
        linewidth=2,
        label="Truth",
    )

plt.xlabel("Beam Energy (MeV)")
plt.ylabel("Counts")
plt.title("Beam Energy")
plt.legend()
plt.tight_layout()

# Plot Excitation Energy

plt.figure(figsize=(8,6))

plt.hist(
    Ex,
    bins=250,
    histtype="step",
    linewidth=2,
    label="Predicted (raw)",
)

plt.hist(
    Ex_snapped,
    bins=250,
    histtype="step",
    linewidth=2,
    label="Predicted (snapped)",
)

if truth_Ex is not None:
    plt.hist(
        truth_Ex,
        bins=250,
        histtype="step",
        linewidth=2,
        label="Truth",
    )

plt.xlabel("Excitation Energy (MeV)")
plt.ylabel("Counts")
plt.title("Excitation Energy")
plt.legend()
plt.tight_layout()

plt.figure(figsize=(6,6))

plt.scatter(truth_beam, beam, s=2)

mn = min(truth_beam.min(), beam.min())
mx = max(truth_beam.max(), beam.max())

plt.plot([mn, mx], [mn, mx], 'k--')

plt.xlabel("True Beam Energy (MeV)")
plt.ylabel("Predicted Beam Energy (MeV)")
plt.title("Beam Energy Reconstruction")

plt.figure(figsize=(6,6))

plt.scatter(truth_Ex, Ex, s=2)

mn = min(truth_Ex.min(), Ex.min())
mx = max(truth_Ex.max(), Ex.max())

plt.plot([mn, mx], [mn, mx], 'k--')

plt.xlabel("True Excitation Energy (MeV)")
plt.ylabel("Predicted Excitation Energy (MeV)")
plt.title("Excitation Energy Reconstruction")

if truth_Ex is not None:
    plt.figure(figsize=(6,6))

    plt.scatter(truth_Ex, Ex_snapped, s=2)

    mn = min(truth_Ex.min(), Ex_snapped.min())
    mx = max(truth_Ex.max(), Ex_snapped.max())

    plt.plot([mn, mx], [mn, mx], 'k--')

    plt.xlabel("True Excitation Energy (MeV)")
    plt.ylabel("Snapped Excitation Energy (MeV)")
    plt.title("Excitation Energy Reconstruction (Snapped)")