import pickle
import os
import matplotlib.pyplot as plt
from  pprint import pprint
import numpy as np
from matplotlib.ticker import ScalarFormatter
import common
from datetime import date

"""
These functions plot stuff. One can but does not have to use them.
"""
def plot_state_densities_and_force_field_with_trajectories(Isotope, file_name, date_str = date.today().isoformat()):
    folder_name = "Data\\"+date_str
    sols_folder = folder_name+'\Solutions'
    
    other_data_folder = folder_name+"\OtherData\\"

    other_data_list = []
    for filename in os.listdir(other_data_folder):
        if Isotope in filename and file_name in filename and filename.endswith(".npz"):
            other_data_list.append(np.load(other_data_folder+filename))
            print(other_data_folder+filename)

    #other_data_list.append(np.load("Z:\Users\Lenny\Lenny\ (Master)\Simulations\ and\ other\ Code\AtomLightSimulations\CleanedUpSimulation\Data\\2025-04-07\OtherData\data_dir_161into2DMOT.npz"))
    if len(other_data_list) > 1:
        print("There were multiple 'Other Data' Files in the directory,"
        "that belong to you Isotope. The list contains:\n", other_data_list)
    elif len(other_data_list) == 0:
        print("Couldn't find a forcefield to plot. List seems to be empty.")
        
    other_data_dir = other_data_list[0]

    Sols = {}

    for filename in os.listdir(sols_folder):
        if Isotope in filename and file_name in filename and filename.endswith(".pkl"):
            key = os.path.splitext(filename)[0]  # Remove .pkl from filename
            full_path = os.path.join(sols_folder, filename)
            with open(full_path, 'rb') as f:
                Sols[key] = pickle.load(f)

    

    basis_g = other_data_dir['basis_g']
    basis_e = other_data_dir['basis_e']
    x0 = other_data_dir['x0']/100
    t0 = other_data_dir['t0']


    
    for names in Sols:
        sol = Sols[names]
        data = sol.N
        num_states = data.shape[0]
        num_timesteps = data.shape[1]

        v0 = sol.v[0,0]
        
        # Normalize (optional)
        #data = data / data.max()
        # Assume your data is in `data` with shape (timesteps, num_states)
        data_g = np.transpose(data[:len(basis_g), :])  # First g states
        data_e = np.transpose(data[len(basis_g):, :])  # Last e states
        fig, axes = plt.subplots(1, 2, figsize=(14, 6), sharey=True)

        # Plot System 1
        im1 = axes[0].imshow(data_g, aspect='auto', origin='lower',
                            cmap='viridis', interpolation='none')
        axes[0].set_title('J = 8')
        axes[0].set_xlabel('State Index')
        axes[0].set_ylabel('position x[m]')


        state_labels_str = [f"{int(s[0])},{int(s[1])}" for s in basis_g[::]]
        axes[0].set_xticks(np.arange(len(basis_g)))
        axes[0].set_xticklabels(state_labels_str, rotation=90, fontsize=8)
        
        
        indices = np.linspace(0, len(sol.t) - 1, 40, dtype=int) #write 40 labels
        position_label = sol.r[0]*x0
        pretty_position_label = [f"{x:.3f}" for x in position_label[indices]]

        axes[0].set_yticks(np.arange(len(sol.t))[indices])
        axes[0].set_yticklabels(pretty_position_label, rotation=0, fontsize=8)
        fig.colorbar(im1, ax=axes[0], label='State Density')

        # Plot System 2
        im2 = axes[1].imshow(data_e, aspect='auto', origin='lower', cmap='viridis', interpolation='none')
        axes[1].set_title('J = 9')
        axes[1].set_xlabel('State Index')        
        state_labels_str = [f"{int(s[0])},{int(s[1])}" for s in basis_e[::-1]]
        axes[1].set_xticks(np.arange(len(basis_e)))
        axes[1].set_xticklabels(state_labels_str, rotation=90, fontsize=8)
        fig.colorbar(im2, ax=axes[1], label='State Density')
        fig.suptitle("Density of mF states at v_0="+str(v0*(x0/t0))+"mps")
        plt.tight_layout(pad=2.0)

        plotname = "StateDensityDy"+Isotope+"at"+str(round(v0*(x0/t0)))+"mps"
        filedir = common.get_data_filepath("Plots", plotname, ".png")
        plt.show()
        #plt.savefig(filedir, dpi=200)

        plt.close()



    Fx_RE = other_data_dir['acceleration_profile']
    x = other_data_dir['x']
    v = other_data_dir['v']

    X,V = np.meshgrid(x, v)
    formatter = ScalarFormatter(useMathText=True)
    formatter.set_powerlimits((-2, 2))  # Forces scientific notation outside this range

    fig, ax0 = plt.subplots(nrows=1, ncols=1)

    im0 = ax0.pcolormesh(X*x0, V*x0/t0, np.nan_to_num(Fx_RE))#, cmap="hot")
    cb0 = plt.colorbar(im0)
    cb0.ax.yaxis.set_major_formatter(formatter)
    cb0.set_label('$a [m/s^2]$')
    ax0.set_xlabel('$x[m]$')
    ax0.set_ylabel('$v[m/s]$')
    ax0.set_title("Acceleration from rate equaitons")

    for names in Sols:
        sol = Sols[names]
        ax0.plot(sol.r[0]*x0, sol.v[0]*x0/t0, "w")
    ax0.set_xlim(x[0]*x0, x[-1]*x0)
    plt.tight_layout()

    plotname = "PhaseSpaceDy"+Isotope
    filedir = common.get_data_filepath("Plots", plotname, ".png")
    plt.show()
    #plt.savefig(filedir, dpi=200)    
    plt.close()

def plot_force_field_and_trajectories(Isotope, file_name, date_str = date.today().isoformat()):
    folder_name = "Data\\"+date_str
    sols_folder = folder_name+'\Solutions'
    
    other_data_folder = folder_name+"\OtherData\\"

    other_data_list = []
    for filename in os.listdir(other_data_folder):
        if Isotope in filename and file_name in filename and filename.endswith(".npz"):
            other_data_list.append(np.load(other_data_folder+filename))
            print(other_data_folder+filename)

    #other_data_list.append(np.load("Z:\Users\Lenny\Lenny\ (Master)\Simulations\ and\ other\ Code\AtomLightSimulations\CleanedUpSimulation\Data\\2025-04-07\OtherData\data_dir_161into2DMOT.npz"))
    if len(other_data_list) > 1:
        print("There were multiple 'Other Data' Files in the directory,"
        "that belong to you Isotope. The list contains:\n", other_data_list)
    elif len(other_data_list) == 0:
        print("Couldn't find a forcefield to plot. List seems to be empty.")
        
    other_data_dir = other_data_list[0]
    #other_data_dir = np.load("Data\\2025-04-14\OtherData\data_dir_Dy161RijIntoPushBeam0Gauss0DetPol-1.npz")
    Sols = {}

    for filename in os.listdir(sols_folder):
        if Isotope in filename and file_name in filename and filename.endswith(".pkl"):
            key = os.path.splitext(filename)[0]  # Remove .pkl from filename
            full_path = os.path.join(sols_folder, filename)
            with open(full_path, 'rb') as f:
                Sols[key] = pickle.load(f)

    

    basis_g = other_data_dir['basis_g']
    basis_e = other_data_dir['basis_e']
    x0 = other_data_dir['x0']/100
    t0 = other_data_dir['t0']


    Fx_RE = other_data_dir['acceleration_profile']
    x = other_data_dir['x']
    v = other_data_dir['v']

    X,V = np.meshgrid(x, v)
    formatter = ScalarFormatter(useMathText=True)
    formatter.set_powerlimits((-2, 2))  # Forces scientific notation outside this range

    fig, ax0 = plt.subplots(nrows=1, ncols=1)

    im0 = ax0.pcolormesh(X*x0, V*x0/t0, np.nan_to_num(Fx_RE))
    cb0 = plt.colorbar(im0)
    cb0.ax.yaxis.set_major_formatter(formatter)
    cb0.set_label('$a [m/s^2]$')
    ax0.set_xlabel('$x[m]$')
    ax0.set_ylabel('$v[m/s]$')
    ax0.set_title("Acceleration from rate equaitons")

    for names in Sols:
        sol = Sols[names]
        ax0.plot(sol.r[0]*x0, sol.v[0]*x0/t0, "w")
    ax0.set_xlim(x[0]*x0, x[-1]*x0)
    plt.tight_layout()

    plotname = "PhaseSpaceDy"+Isotope
    filedir = common.get_data_filepath("Plots", plotname, ".png")
    plt.show()
    #plt.savefig(filedir, dpi=200)    
    plt.close()

def plot_force_field(Isotope, file_name, date_str = date.today().isoformat()):
    folder_name = "Data\\"+date_str
    sols_folder = folder_name+'\Solutions'

    other_data_folder = folder_name+"\OtherData\\"

    other_data_list = []
    for filename in os.listdir(other_data_folder):
        if file_name in filename and Isotope in filename and filename.endswith(".npz"):
            other_data_list.append(np.load(other_data_folder+filename))
            print(other_data_folder+filename)

    #other_data_list.append(np.load("Z:\Users\Lenny\Lenny\ (Master)\Simulations\ and\ other\ Code\AtomLightSimulations\CleanedUpSimulation\Data\\2025-04-07\OtherData\data_dir_161into2DMOT.npz"))
    if len(other_data_list) > 1:
        print("There were multiple 'Other Data' Files in the directory,"
        "that belong to you Isotope. The first entry will be chosen. The list contains:\n", other_data_list)
    elif len(other_data_list) == 0:
        print("Couldnt read a file in 'Other Data'. Check that you chose the correct date, Isotope and filename.")
        
    other_data_dir = other_data_list[0]

    basis_g = other_data_dir['basis_g']
    basis_e = other_data_dir['basis_e']
    x0 = other_data_dir['x0']/100
    t0 = other_data_dir['t0']


    Fx_RE = other_data_dir['acceleration_profile']
    x = other_data_dir['x']
    v = other_data_dir['v']

    X,V = np.meshgrid(x, v)
    formatter = ScalarFormatter(useMathText=True)
    formatter.set_powerlimits((-2, 2))  # Forces scientific notation outside this range

    fig, ax0 = plt.subplots(nrows=1, ncols=1)

    im0 = ax0.pcolormesh(X*x0, V*x0/t0, np.nan_to_num(Fx_RE))
    cb0 = plt.colorbar(im0)
    cb0.ax.yaxis.set_major_formatter(formatter)
    cb0.set_label('$a [m/s^2]$')
    ax0.set_xlabel('$x[m]$')
    ax0.set_ylabel('$v[m/s]$')
    ax0.set_title("Acceleration from rate equaitons")
    plt.tight_layout()
    plotname = "PhaseSpaceDy"+Isotope
    filedir = common.get_data_filepath("Plots", plotname, ".png")
    plt.show()
    #plt.savefig(filedir, dpi=200)    
    plt.close()