from multiprocessing import Pool, cpu_count
import os
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from simulation import run_simulation
from data_processing import get_controller_files

# Constants and setup
initial_conditions = {
    "small_perturbation": (0.1*np.pi, 0.0, 0.0, 0.0),
    "large_perturbation": (-np.pi, 0.0, 0.0, 0),
    "overshoot_vertical_test": (-0.1*np.pi, 2*np.pi, 0.0, 0.0),
    "overshoot_angle_test": (0.2*np.pi, 2*np.pi, 0.0, 0.3*np.pi),
    "extreme_perturbation": (4*np.pi, 0.0, 0.0, 0),
}
loss_functions = ["constant", "linear", "quadratic", "cubic", "inverse", "inverse_squared", "inverse_cubed"]
loss_functions_mirrored = ["linear", "quadratic", "cubic", "inverse", "inverse_squared", "inverse_cubed"]
loss_functions_mirrored = [i+"_mirrored" for i in loss_functions_mirrored]
loss_functions = loss_functions + loss_functions_mirrored
epoch_range = (0, 100)  # Start and end of epoch range
epoch_step = 1          # Interval between epochs
dt = 0.02                # Time step for simulation
num_steps = 500          # Number of steps in each simulation

# Original individual 3D plot function
def plot_3d_epoch_evolution(epochs, theta_over_epochs, desired_theta, save_path, title, num_steps, dt):
    fig = plt.figure(figsize=(7, 5))
    ax = fig.add_subplot(111, projection='3d')
    time_steps = np.arange(num_steps) * dt

    theta_values = np.concatenate(theta_over_epochs)
    theta_min = np.min(theta_values)
    theta_max = np.max(theta_values)

    desired_range_min = max(theta_min, desired_theta - 1.5 * np.pi)
    desired_range_max = min(theta_max, desired_theta + 1.5 * np.pi)

    for epoch, theta_vals in reversed(list(zip(epochs, theta_over_epochs))):
        masked_theta_vals = np.array(theta_vals)
        masked_theta_vals[(masked_theta_vals < desired_range_min) | (masked_theta_vals > desired_range_max)] = np.nan
        ax.plot([epoch] * len(time_steps), time_steps, masked_theta_vals)

    epochs_array = np.array(epochs)
    ax.plot(epochs_array, [time_steps.max()] * len(epochs_array), [desired_theta] * len(epochs_array),
            color='r', linestyle='--', linewidth=2, label='Desired Theta at End Time')

    ax.set_xlabel("Epoch")
    ax.set_ylabel("Time (s)")
    ax.set_zlabel("Theta (rad)")
    # ax.set_zscale('symlog')
    ax.set_title(title)
    ax.set_zlim(desired_range_min, desired_range_max)
    ax.view_init(elev=20, azim=-135)

    if not os.path.exists(os.path.dirname(save_path)):
        os.makedirs(os.path.dirname(save_path))
    plt.savefig(save_path, dpi=300)
    plt.close()
    print(f"Saved plot as '{save_path}'.")

# Helper function for composite 3D plotting on an existing axis
def plot_3d_epoch_evolution_on_axis(ax, epochs, theta_over_epochs, desired_theta, title, num_steps, dt):
    time_steps = np.arange(num_steps) * dt
    theta_values = np.concatenate(theta_over_epochs)
    theta_min = np.min(theta_values)
    theta_max = np.max(theta_values)
    
    desired_range_min = max(theta_min, desired_theta - 1.5 * np.pi)
    desired_range_max = min(theta_max, desired_theta + 1.5 * np.pi)
    
    for epoch, theta_vals in reversed(list(zip(epochs, theta_over_epochs))):
        masked_theta_vals = np.array(theta_vals)
        masked_theta_vals[(masked_theta_vals < desired_range_min) | (masked_theta_vals > desired_range_max)] = np.nan
        ax.plot([epoch] * len(time_steps), time_steps, masked_theta_vals)
    
    epochs_array = np.array(epochs)
    ax.plot(epochs_array, [time_steps.max()] * len(epochs_array), [desired_theta] * len(epochs_array),
            color='r', linestyle='--', linewidth=2, label='Desired Theta')
    
    ax.set_xlabel("Epoch")
    ax.set_ylabel("Time (s)")
    ax.set_zlabel("Theta (rad)")
    # ax.set_zscale('symlog')
    ax.set_title(title)
    ax.set_zlim(desired_range_min, desired_range_max)
    ax.view_init(elev=20, azim=-135)

# Composite plot function that accepts a list of loss functions to include
def plot_composite_loss_functions_for_condition(all_results, condition_name, desired_theta, num_steps, dt, loss_list, save_path):
    # total rows = 1 (constant) + len(loss_list)
    total_rows = 1 + len(loss_list)
    fig = plt.figure(figsize=(12, 3 * total_rows))
    gs = gridspec.GridSpec(total_rows, 2)
    
    # Top row: "constant" loss function spanning both columns
    ax_const = fig.add_subplot(gs[0, :], projection='3d')
    epochs_const, theta_const = all_results["constant"][condition_name]
    plot_3d_epoch_evolution_on_axis(ax_const, epochs_const, theta_const, desired_theta,
                                    "constant", num_steps, dt)
    
    # For each loss function in the provided loss_list, plot the regular and its mirrored version
    for i, loss in enumerate(loss_list):
        # Left subplot: regular loss function
        ax_left = fig.add_subplot(gs[i+1, 0], projection='3d')
        if condition_name in all_results.get(loss, {}):
            epochs_loss, theta_loss = all_results[loss][condition_name]
            plot_3d_epoch_evolution_on_axis(ax_left, epochs_loss, theta_loss, desired_theta,
                                            loss, num_steps, dt)
        else:
            ax_left.set_title(f"No data for {loss}")
        
        # Right subplot: mirrored loss function
        mirrored_loss = loss + "_mirrored"
        ax_right = fig.add_subplot(gs[i+1, 1], projection='3d')
        if condition_name in all_results.get(mirrored_loss, {}):
            epochs_mir, theta_mir = all_results[mirrored_loss][condition_name]
            plot_3d_epoch_evolution_on_axis(ax_right, epochs_mir, theta_mir, desired_theta,
                                            mirrored_loss, num_steps, dt)
        else:
            ax_right.set_title(f"No data for {mirrored_loss}")
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=300)
    plt.close()
    print(f"Saved composite plot to {save_path}")

# Main execution
if __name__ == "__main__":
    all_results = {}  # Dictionary to store results by loss function

    for condition_name, initial_condition in initial_conditions.items():
        condition_text = f"IC_{'_'.join(map(lambda x: str(round(x, 2)), initial_condition))}"
        desired_theta = initial_condition[-1]
        condition_path = f"/home/judson/Neural-Networks-in-GNC/inverted_pendulum/analysis/time_weighting/{condition_name}"
        os.makedirs(condition_path, exist_ok=True)
        
        # For each loss function, run simulation and collect results
        for loss_function in loss_functions:
            directory = f"/home/judson/Neural-Networks-in-GNC/inverted_pendulum/training/time_weighting/{loss_function}/controllers"
            controllers = get_controller_files(directory, epoch_range, epoch_step)
            tasks = [(c, initial_condition, directory, dt, num_steps) for c in controllers]
            
            print("Starting worker processes")
            with Pool(min(cpu_count(), 16)) as pool:
                results = pool.map(run_simulation, tasks)
            
            results.sort(key=lambda x: x[0])
            epochs, state_histories, torque_histories = zip(*results)
            theta_over_epochs = [[state[0] for state in history] for history in state_histories]
            
            if loss_function not in all_results:
                all_results[loss_function] = {}
            all_results[loss_function][condition_name] = (epochs, theta_over_epochs)
            
            # Optional: save individual 3D plots
            print(f"Plotting the 3D epoch evolution for {loss_function} under {condition_text}")
            title = f"Pendulum Angle Evolution for {loss_function} and {condition_text}"
            save_path = os.path.join(condition_path, "epoch_evolution", f"{loss_function}.png")
            plot_3d_epoch_evolution(epochs, theta_over_epochs, desired_theta, save_path, title, num_steps, dt)
            print("")
        
        # Create composite figure for linear, quadratic, cubic (with constant at the top)
        loss_list1 = ["linear", "quadratic", "cubic"]
        composite_save_path1 = os.path.join(condition_path, "composite_epoch_evolution_linear_quadratic_cubic.png")
        plot_composite_loss_functions_for_condition(all_results, condition_name, desired_theta, num_steps, dt, loss_list1, composite_save_path1)
        
        # Create composite figure for inverse, inverse_squared, inverse_cubed (with constant at the top)
        loss_list2 = ["inverse", "inverse_squared", "inverse_cubed"]
        composite_save_path2 = os.path.join(condition_path, "composite_epoch_evolution_inverse_losses.png")
        plot_composite_loss_functions_for_condition(all_results, condition_name, desired_theta, num_steps, dt, loss_list2, composite_save_path2)
        
        print(f"Completed plotting for all loss functions under {condition_name} condition.\n")