import torch
import torch.nn as nn
from torchdiffeq import odeint
import numpy as np
import os
import shutil
import csv
import math
import matplotlib.pyplot as plt
import pandas as pd

from PendulumController import PendulumController

# List of controller file names to validate.
# Replace these paths with your actual controller file paths.
controller_file_names = [
    "/home/judson/Neural-Networks-in-GNC/inverted_pendulum/training/base_loss_learning_rate_sweep/one_fourth/lr_0.300/controllers/controller_199.pth",
    "/home/judson/Neural-Networks-in-GNC/inverted_pendulum/training/base_loss_learning_rate_sweep/four/lr_0.200/controllers/controller_200.pth"
]

# Constants for simulation
m = 10.0
g = 9.81
R = 1.0
dt = 0.01      # time step
T = 20.0       # total simulation time
num_steps = int(T / dt)
t_eval = np.linspace(0, T, num_steps)

# In-sample validation cases: [theta0, omega0, alpha0, desired_theta]
in_sample_cases = [
    (1/6 * np.pi,    0.0, 0.0, 0.0),
    (-1/6 * np.pi,   0.0, 0.0, 0.0),
    (2/3 * np.pi,    0.0, 0.0, 0.0),
    (-2/3 * np.pi,   0.0, 0.0, 0.0),
    (0.0, 1/3 * np.pi,      0.0, 0.0),
    (0.0, -1/3 * np.pi,     0.0, 0.0),
    (0.0, 2 * np.pi,        0.0, 0.0),
    (0.0, -2 * np.pi,       0.0, 0.0),
    (0.0, 0.0, 0.0,     2*np.pi),
    (0.0, 0.0, 0.0,     -2*np.pi),
    (0.0, 0.0, 0.0,     1/2 * np.pi),
    (0.0, 0.0, 0.0,     -1/2 * np.pi),
    (0.0, 0.0, 0.0,     1/3 * np.pi),
    (0.0, 0.0, 0.0,     -1/3 * np.pi),
    (1/4 * np.pi,    1 * np.pi, 0.0, 0.0),
    (-1/4 * np.pi,   -1 * np.pi, 0.0, 0.0),
    (1/2 * np.pi,    -1 * np.pi, 0.0, 1/3 * np.pi),
    (-1/2 * np.pi,   1 * np.pi, 0.0, -1/3 * np.pi),
    (1/4 * np.pi,    1 * np.pi, 0.0, 2 * np.pi),
    (-1/4 * np.pi,   -1 * np.pi, 0.0, 2 * np.pi),
    (1/2 * np.pi,    -1 * np.pi, 0.0, 4 * np.pi),
    (-1/2 * np.pi,   1 * np.pi, 0.0, -4 * np.pi),
]

# Out-of-sample validation cases (generated previously)
out_sample_cases = [
    (-2.198958, -4.428501, 0.450833, 0.000000),
    (1.714196, -0.769896, 0.202738, 0.000000),
    (0.241195, -5.493715, 0.438996, 0.000000),
    (0.030605, 4.901513, -0.479243, 0.000000),
    (1.930445, -1.301926, -0.454050, 0.000000),
    (-0.676063, 4.246865, 0.036303, 0.000000),
    (0.734920, -5.925202, 0.047097, 0.000000),
    (-3.074471, -3.535424, 0.315438, 0.000000),
    (-0.094486, 6.111091, 0.150525, 0.000000),
    (-1.647671, 5.720526, 0.334181, 0.000000),
    (-2.611260, 5.087704, 0.045460, -3.610785),
    (1.654137, 0.982081, -0.192725, 1.003872),
    (-2.394899, 3.550547, -0.430938, 3.261897),
    (0.474917, 0.555166, -0.285173, 1.866752),
    (-0.640369, -4.678490, -0.340663, 3.150098),
    (1.747517, -3.248204, -0.001520, 1.221787),
    (2.505283, -2.875006, -0.065617, -3.690269),
    (1.337244, 2.221707, 0.044979, -2.459730),
    (1.531012, 2.230981, -0.291206, -1.924535),
    (-1.065792, 4.320740, 0.075405, -1.550644),
]

# Define the ODE integration step using RK4, returning the new state and the computed torque.
def pendulum_ode_step(state, dt, desired_theta):
    theta, omega, alpha = state

    def compute_torque(th, om, al):
        inp = torch.tensor([[th, om, al, desired_theta]], dtype=torch.float32)
        with torch.no_grad():
            torque = controller(inp)
            torque = torch.clamp(torque, -250, 250)
        return float(torque)

    def derivatives(state, torque):
        th, om, al = state
        a = (g / R) * np.sin(th) + torque / (m * R**2)
        return np.array([om, a, 0])  # dtheta, domega, dalpha

    # RK4 Steps
    torque1 = compute_torque(theta, omega, alpha)
    k1 = dt * derivatives(state, torque1)

    state_k2 = state + 0.5 * k1
    torque2 = compute_torque(state_k2[0], state_k2[1], state_k2[2])
    k2 = dt * derivatives(state_k2, torque2)

    state_k3 = state + 0.5 * k2
    torque3 = compute_torque(state_k3[0], state_k3[1], state_k3[2])
    k3 = dt * derivatives(state_k3, torque3)

    state_k4 = state + k3
    torque4 = compute_torque(state_k4[0], state_k4[1], state_k4[2])
    k4 = dt * derivatives(state_k4, torque4)

    new_state = state + (k1 + 2*k2 + 2*k3 + k4) / 6.0
    final_torque = (torque1 + 2*torque2 + 2*torque3 + torque4) / 6.0

    return new_state, final_torque

# Function to run validation (in-sample or out-of-sample) for a given controller.
def run_validation(controller, cases, t_eval, dt, num_steps, validation_type, save_dir):
    losses = []
    final_thetas = []
    os.makedirs(save_dir, exist_ok=True)
    for idx, (theta0, omega0, alpha0, desired_theta) in enumerate(cases):
        state = np.array([theta0, omega0, alpha0])
        theta_vals, omega_vals, alpha_vals, torque_vals = [], [], [], []
        for _ in range(num_steps):
            theta_vals.append(state[0])
            omega_vals.append(state[1])
            alpha_vals.append(state[2])
            state, torque = pendulum_ode_step(state, dt, desired_theta)
            torque_vals.append(torque)
        theta_vals = np.array(theta_vals)
        omega_vals = np.array(omega_vals)
        alpha_vals = np.array(alpha_vals)
        torque_vals = np.array(torque_vals)
        final_theta = theta_vals[-1]
        final_thetas.append(final_theta)
        loss = 1e3 * np.mean((theta_vals - desired_theta)**2)
        losses.append(loss)
        
        # Plot results
        fig, ax1 = plt.subplots(figsize=(10, 6))
        ax1.plot(t_eval, theta_vals, label="theta", color="blue")
        ax1.plot(t_eval, omega_vals, label="omega", color="green")
        ax1.plot(t_eval, alpha_vals, label="alpha", color="red")
        ax1.axhline(desired_theta, label="Desired Theta", color="black")
        ax1.set_xlabel("Time [s]")
        ax1.set_ylabel("Theta, Omega, Alpha")
        ax1.grid(True)
        ax1.legend(loc="upper left")
        ax2 = ax1.twinx()
        ax2.plot(t_eval, torque_vals, label="Torque", color="purple", linestyle="--")
        ax2.set_ylabel("Torque [Nm]")
        ax2.legend(loc="upper right")
        plt.title(f"Case {idx+1}: IC (theta={theta0:.3f}, omega={omega0:.3f}, alpha={alpha0:.3f})")
        plt.tight_layout()
        filename = f"{idx+1}_theta0_{theta0:.3f}_omega0_{omega0:.3f}_alpha0_{alpha0:.3f}_desired_{desired_theta:.3f}_final_{final_theta:.3f}.png"
        plt.savefig(os.path.join(save_dir, filename))
        plt.close()
        print(f"Saved {validation_type} case {idx+1}")
    
    df = pd.DataFrame(cases, columns=["theta0", "omega0", "alpha0", "desired_theta"])
    df["final_theta"] = final_thetas
    df["loss"] = losses
    df.insert(0, "Run #", range(1, len(cases) + 1))
    df.to_csv(os.path.join(save_dir, "summary.csv"), index=False)
    print(f"Average Loss for {validation_type}: {np.mean(losses):.4f}")
    desired_vals = np.array([case[3] for case in cases])
    avg_abs_err = np.mean(np.abs(np.array(final_thetas) - desired_vals))
    print(f"Average abs(Final Theta - Desired Theta) for {validation_type}: {avg_abs_err:.4f}")
    return df

# Main validation loop for each controller file
base_validation_dir = "validation"
os.makedirs(base_validation_dir, exist_ok=True)

for controller_file in controller_file_names:
    controller_id = os.path.splitext(os.path.basename(controller_file))[0]
    controller_dir = os.path.join(base_validation_dir, controller_id)
    in_sample_dir = os.path.join(controller_dir, "in-sample")
    out_sample_dir = os.path.join(controller_dir, "out-of-sample")
    os.makedirs(in_sample_dir, exist_ok=True)
    os.makedirs(out_sample_dir, exist_ok=True)
    
    # Load controller
    controller = PendulumController()
    controller.load_state_dict(torch.load(controller_file))
    controller.eval()
    print(f"{controller_file} loaded as {controller_id}")
    
    # In-sample validation
    print("Performing in-sample validation")
    df_in = run_validation(controller, in_sample_cases, t_eval, dt, num_steps, "in-sample", in_sample_dir)
    
    # Out-of-sample validation
    print("Performing out-of-sample validation")
    df_out = run_validation(controller, out_sample_cases, t_eval, dt, num_steps, "out-of-sample", out_sample_dir)

print("Validation complete. Check the 'validation' directory for plots and summaries.")