import os
import json
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec

# --- Helper function to filter epoch data ---
def filter_epoch_data(epochs, theta_over_epochs, epoch_range, epoch_step):
    """
    Filters the list of epochs and corresponding theta values.
    Only epochs between epoch_range[0] and epoch_range[1] (inclusive) are kept,
    and then every epoch_step element is selected.
    """
    filtered_epochs = []
    filtered_theta_over_epochs = []
    for i, ep in enumerate(epochs):
        if ep >= epoch_range[0] and ep <= epoch_range[1]:
            filtered_epochs.append(ep)
            filtered_theta_over_epochs.append(theta_over_epochs[i])
    if epoch_step > 1:
        filtered_epochs = filtered_epochs[::epoch_step]
        filtered_theta_over_epochs = filtered_theta_over_epochs[::epoch_step]
    return filtered_epochs, filtered_theta_over_epochs

# --- Composite plotting functions (with z-scale set to linear) ---
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)

def plot_composite_loss_functions_for_condition(cond_results, condition_name, desired_theta, num_steps, dt, loss_list, save_path, plot_epoch_range, plot_epoch_step, swap_columns=False):
    total_rows = 1 + len(loss_list)  # one row for "constant" + one row per loss function
    fig = plt.figure(figsize=(12, 3 * total_rows))
    gs = gridspec.GridSpec(total_rows, 2)
    
    # Top row: "constant" spanning both columns
    ax_const = fig.add_subplot(gs[0, :], projection='3d')
    epochs_const, theta_const = cond_results["constant"]
    epochs_const, theta_const = filter_epoch_data(epochs_const, theta_const, plot_epoch_range, plot_epoch_step)
    plot_3d_epoch_evolution_on_axis(ax_const, epochs_const, theta_const, desired_theta, "constant", num_steps, dt)
    
    # For each loss in the provided list, plot the pair of original and mirrored
    for i, loss in enumerate(loss_list):
        if not swap_columns:
            # Left: original; Right: mirrored.
            ax_left = fig.add_subplot(gs[i+1, 0], projection='3d')
            if loss in cond_results:
                epochs_loss, theta_loss = cond_results[loss]
                epochs_loss, theta_loss = filter_epoch_data(epochs_loss, theta_loss, plot_epoch_range, plot_epoch_step)
                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}")
            
            ax_right = fig.add_subplot(gs[i+1, 1], projection='3d')
            mirrored_loss = loss + "_mirrored"
            if mirrored_loss in cond_results:
                epochs_mir, theta_mir = cond_results[mirrored_loss]
                epochs_mir, theta_mir = filter_epoch_data(epochs_mir, theta_mir, plot_epoch_range, plot_epoch_step)
                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}")
        else:
            # Swap: Left: mirrored; Right: original.
            mirrored_loss = loss + "_mirrored"
            ax_left = fig.add_subplot(gs[i+1, 0], projection='3d')
            if mirrored_loss in cond_results:
                epochs_mir, theta_mir = cond_results[mirrored_loss]
                epochs_mir, theta_mir = filter_epoch_data(epochs_mir, theta_mir, plot_epoch_range, plot_epoch_step)
                plot_3d_epoch_evolution_on_axis(ax_left, epochs_mir, theta_mir, desired_theta, mirrored_loss, num_steps, dt)
            else:
                ax_left.set_title(f"No data for {mirrored_loss}")
            
            ax_right = fig.add_subplot(gs[i+1, 1], projection='3d')
            if loss in cond_results:
                epochs_loss, theta_loss = cond_results[loss]
                epochs_loss, theta_loss = filter_epoch_data(epochs_loss, theta_loss, plot_epoch_range, plot_epoch_step)
                plot_3d_epoch_evolution_on_axis(ax_right, epochs_loss, theta_loss, desired_theta, loss, num_steps, dt)
            else:
                ax_right.set_title(f"No data for {loss}")
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=300)
    plt.close()
    print(f"Saved composite plot to {save_path}")

# --- Load simulation results from JSON files (one per condition) ---
output_dir = "/home/judson/Neural-Networks-in-GNC/inverted_pendulum/analysis/time_weighting2"
# Get all JSON files (each representing one condition)
condition_files = [f for f in os.listdir(output_dir) if f.endswith(".json")]

# Define simulation parameters (must match those used in run_tests.py)
dt = 0.02
num_steps = 500

# Specify the epoch range and step for plotting
plot_epoch_range = (0, 100)  # Only plot epochs between 0 and 100
plot_epoch_step = 1          # Use every epoch (change this if you want to sample)

# For each condition file, load its data and create the composite plots
for file in condition_files:
    condition = file.replace(".json", "")
    condition_file = os.path.join(output_dir, file)
    with open(condition_file, "r") as f:
        cond_results = json.load(f)
    
    # Determine desired_theta from constant loss data (last theta value of the last epoch)
    epochs_const, theta_const = cond_results["constant"]
    desired_theta = theta_const[-1][-1]
    
    # Create a directory for this condition's plots
    condition_path = os.path.join(output_dir, condition)
    os.makedirs(condition_path, exist_ok=True)
    
    # Composite figure for linear, quadratic, cubic (with constant at 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(cond_results, condition, desired_theta, num_steps, dt, loss_list1, composite_save_path1, plot_epoch_range, plot_epoch_step)
    
    # Composite figure for inverse, inverse_squared, inverse_cubed (mirrored version on the left)
    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(cond_results, condition, desired_theta, num_steps, dt, loss_list2, composite_save_path2, plot_epoch_range, plot_epoch_step, swap_columns=True)
    
    print(f"Completed plotting for condition: {condition}")