"""Analysis functions for BIFROST simulations.
    
This module includes functions for data reading, manipulation, plotting and analysis.

Functions:

"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse
import seaborn as sns
import scipy as sp
from tqdm import tqdm
import os
from .mcstas_functions import read_data_files
from time import time
from .neutron_functions import *
import pkg_resources
from . import constants



def load_constants(file):
    with pkg_resources.resource_stream(__name__,"constants/"+file) as f:
        data = np.load(f)
    return data

    

# Calculation of energy distribution on a tube
def e_cos(x, a, b, c):
    return a/np.cos(b*(x -c))

def fit_energy(x, y):
    # Initial guess for the parameters
    initial_guess = [min(y), 0.02, 50]
    params, covariance = sp.optimize.curve_fit(e_cos, x, y, p0=initial_guess)
    a, b, c = params
    return a, b, c

def weighted_cov(X,w):
    w /= np.sum(w)
    mean = np.average(X, axis=0, weights=w)
    X_centered = X - mean
    
    cov = (X_centered.T * w) @ X_centered / np.sum(w)
    return mean, cov

def flat_to_symmetric(flat,d=4):
    mat = np.zeros((d, d))
    idx = np.tril_indices(d)
    mat[idx] = flat
    mat[(idx[1], idx[0])] = flat
    return mat


# Plotting functions

def reduced_res_plot(data, file_name, folder):
    labels = [r"$\Delta Q_x$ [Å]",r"$\Delta Q_z$ [Å]",r"$\Delta \hbar \omega$ [meV]"]
    cols = ["dQ_x", "dQ_z", "dhbar_w"]
    fig, axs = plt.subplots(3, 3, figsize=(12, 12))
    fig.suptitle(f"Resolution Plot")
    for i in range(3):
        for j in range(3):
            if i == j:
                sns.histplot(data, x=cols[i], bins=20, ax=axs[i, j], weights="Intensity")
                
                axs[i, j].set_xlabel(labels[i])
                axs[i, j].set_ylabel("Intensity")
            else:
                ax = sns.histplot(data, x=cols[i], y=cols[j], bins=20, ax=axs[i, j], cmap="viridis", cbar=True, weights="Intensity")
                cbar = ax.collections[0].colorbar
                axs[i, j].set_xlabel(labels[i])
                axs[i, j].set_ylabel(labels[j])
                cbar.set_label("Intensity")
    plt.tight_layout()
    plt.savefig(os.path.join(folder, file_name))
    plt.close()

def plot_cov_ellipse(cov, pos, ax=None, **kwargs):
    if ax is None:
        ax = plt.gca()
    vals, vecs = np.linalg.eigh(cov)
    order = vals.argsort()[::-1]
    vals, vecs = vals[order], vecs[:, order]
    width, height = 2.355 * np.sqrt(vals)
    angle = np.degrees(np.arctan2(*vecs[:, 0][::-1]))
    ellip = Ellipse(xy=pos, width=width, height=height, angle=angle, lw=3,**kwargs)
    ax.add_patch(ellip)
    return ellip




# Classes for BIFROST data

class BIFROST_nD:
    """
    Class for BIFROST data generated by several Monitor_nD components.
    
    This class is used to read, manipulate, and plot data from the BIFROST instrument.
    This is specifically only for data that will be gotten from the actual BIFROST instrument
    
    Attributes:
        folder (str): Path to the folder containing the data files.
        
    """
    
    def read_data(self, remove=True, reread=False):
        if os.path.exists(os.path.join(self.folder, "data.csv")) and not reread:
            self.data = pd.read_csv(os.path.join(self.folder, "data.csv"))
            return
        # Get list of all files in the folder
        files = [f for f in os.listdir(self.folder) if "list" in f]
        if isinstance(files, bytes):
            files = [f.decode("utf-8") for f in files]
        data = pd.DataFrame()
        
        for file in files:
            data_i, parameters = read_data_files(self.folder, file, D="N")
            # Get data names:
            data_names = parameters["title"].split()
            data_col = data.columns.tolist()
            for i in range(data_i.shape[1]):
                if data_names[i] not in data_col:
                    data[data_names[i]] = data_i[:, i]
            
        # Save the   data to a CSV file
        data.to_csv(os.path.join(self.folder, "data.csv"), index=False)
        np.save(os.path.join(self.folder, "parameters.npy"), parameters)
        if remove:
            for file in files:
                os.remove(os.path.join(self.folder, file))
                
        self.data = data
        self.parameters = parameters
        return 

    
    def RT_2_CA(self):
        analyzer = self.data["TUBE"] // 3 + 1
        self.data["analyzer"] = analyzer.astype(int)
        channel = 3*(self.data["RING"] // 2) + self.data["TUBE"] % 3 + 1
        self.data["channel"] = channel.astype(int)
        
    def pixel_split(self):
        diff = self.data["event_charge_right"] - self.data["event_charge_left"]
        add = self.data["event_charge_left"] + self.data["event_charge_right"]
        L_x = diff / add
        
        ds = np.linspace(-1, 1, 301)
        
        L_x = [l if l< 1 else 0.999999 for l in L_x]
        bins = np.digitize(L_x, ds)
        
        self.data["D_tube"] = bins // 100 + 1
        self.data["d_i"] = bins % 100
        
        # Flip second tube
        self.data.loc[self.data["D_tube"] == 2, "d_i"] = 99 - self.data.loc[self.data["D_tube"] == 2, "d_i"]   
        
        return
    
    def E_f_k_f(self):
        Es = load_constants("Es_fit.npy")
        L = load_constants("dL.npy")
        phi = load_constants("phi.npy")
        
        a4 = float(self.parameters["a4"])
        
        C = self.data["channel"]
        A = self.data["analyzer"]
        T = self.data["D_tube"]
        d = self.data["d_i"]
        
        E = Es[C-1, A-1, T-1, d]
        L = L[C-1, A-1, T-1, d]
        phi = np.deg2rad(phi[C-1, A-1, T-1, d])
        
        
        self.data["E_f"] = E
        self.data["L_2"] = L
        self.data["t_2"] = L / E2v_(E)
        self.data["dphi"] = phi
        
        self.data["k_f"] = E2k_(E)
        self.data["k_fx"] = self.data["k_f"] * np.sin(phi)
        self.data["k_fy"] = np.zeros_like(self.data["k_f"])
        self.data["k_fz"] = self.data["k_f"] * np.cos(phi)
        
    def psc_timing(self):
        
        psc_v = float(self.parameters["psc1speed"])
        psc_offset_1 = float(self.parameters["psc1phase"])
        psc_offset_2 = float(self.parameters["psc2phase"])
        
        psc_off = (psc_offset_1 + psc_offset_2) / (720*psc_v)
        psc_opening = (psc_offset_2 - psc_offset_1+170) / (360*psc_v)
        
        return psc_off, psc_opening
    
    def E_i_k_i(self):
        t2 = self.data["t_2"]
        t = self.data["event_time"]
        
        t1 = t - t2
        L_MS = 162 # m
        L_PSC = 6.342 # m
        
        psc_off, psc_opening = self.psc_timing()
        
        self.data["dt"] = np.ones_like(t1)*psc_opening
        
        t1 = t1 - psc_off
        L1 = L_MS - L_PSC
        
        v = L1 / t1
        E = v2E_(v)
        k = v2k_(v)
        
        self.data["E_i"] = E
        self.data["k_i"] = k
        self.data["k_ix"] = np.zeros_like(k)
        self.data["k_iy"] = np.zeros_like(k)
        self.data["k_iz"] = k
        
    def Q_hw(self):
        final_data = pd.DataFrame()
        
        final_data["Q_x"] = self.data["k_ix"] - self.data["k_fx"]
        final_data["Q_y"] = self.data["k_iy"] - self.data["k_fy"]
        final_data["Q_z"] = self.data["k_iz"] - self.data["k_fz"]
        final_data["hbar_omega"] = self.data["E_i"] - self.data["E_f"]
        final_data["Intensity"] = self.data["Intensity"]
        
        self.data = final_data
        
    def convert(self):
        self.RT_2_CA()
        self.pixel_split()
        self.E_f_k_f()
        self.E_i_k_i()
        self.Q_hw()
        
    def __init__(self, folder, remove=True, reread=False):
        self.folder = folder
        
        self.read_data(remove=remove, reread=reread)
        self.convert()
        
    def add_data(self, folder):
        new = BIFROST_nD(folder)
        new.read_data(remove=True)
        new.convert()
        new_data = new.data
        self.data = pd.concat([self.data, new_data], ignore_index=True)
        self.data.reset_index(drop=True, inplace=True)
        return
        
class BIFROST_Res(BIFROST_nD):
    """
    Class for BIFROST data generated by several Monitor_nD components.
    
    This class is used to calculate the resolution of the BIFROST instrument.

    Args:
        BIFROST_nD (class): Class for BIFROST data generated by several Monitor_nD components.
    """
    
    def Q_hw(self):
        
        # Calculate the assumed Q, hw
        self.data["Q_x"] = self.data["k_ix"] - self.data["k_fx"]
        self.data["Q_y"] = self.data["k_iy"] - self.data["k_fy"]
        self.data["Q_z"] = self.data["k_iz"] - self.data["k_fz"]
        self.data["hbar_omega"] = self.data["E_i"] - self.data["E_f"]
        
        # Calculate the actual Q, hw
        self.data["Qx"] = self.data["kix"] - self.data["kfx"]
        self.data["Qy"] = self.data["kiy"] - self.data["kfy"]
        self.data["Qz"] = self.data["kiz"] - self.data["kfz"]
        ki = np.sqrt(self.data["kix"]**2 + self.data["kiy"]**2 + self.data["kiz"]**2)
        kf = np.sqrt(self.data["kfx"]**2 + self.data["kfy"]**2 + self.data["kfz"]**2)
        self.data["hw"] = k2E_(ki) - k2E_(kf)
        
        self.data["dQx"] = self.data["Q_x"] - self.data["Qx"]
        self.data["dQy"] = self.data["Q_y"] - self.data["Qy"]
        self.data["dQz"] = self.data["Q_z"] - self.data["Qz"]
        self.data["dhw"] = self.data["hbar_omega"] - self.data["hw"]
        
    def convert(self):
        self.RT_2_CA()
        self.pixel_split()
        self.E_f_k_f()
        self.E_i_k_i()
        self.Q_hw()
        
    
    def __init__(self, folder, remove=True, reread=False):
        self.folder = folder
        self.read_data(remove=remove, reread=reread)
        self.convert()

    def pixel_resolution(self, save_folder=None, save_file=None, threshold=1, plot=False, data_save=False):
        if save_folder == None:
            save_folder = self.folder+"/pixel_resolution"
        
        Res = pd.DataFrame(columns=["channel", "analyzer", "D_tube", "d_i", "I","Qx_m", "Qy_m", "Qz_m", "hw_m", 
                               "sigma_Qx_2", "sigma_Qx_Qy", "sigma_Qx_Qz", "sigma_Qx_hw", "sigma_Qy_2", 
                               "sigma_Qy_Qz", "sigma_Qy_hw", "sigma_Qz_2", "sigma_Qz_hw", "sigma_hw_2"])
        
        for c in range(1, 10):
            if len(self.data[self.data["channel"] == c]) == 0:
                continue
            for a in range(1, 6):
                if len(self.data[self.data["analyzer"] == a]) == 0:
                    continue
                for t in range(1, 4):
                    for d in range(100):
                        data_i = self.data[(self.data["channel"] == c) & (self.data["analyzer"] == a) & (self.data["TUBE"] == t) & (self.data["d_i"] == d)]
                        intensity = data_i["Intensity"].sum()
                        if intensity < threshold:
                            continue
                        
                        if plot:
                            if save_file == None:
                                save_file = f"pixel_{c}_{a}_{t}_{d}.png"
                            reduced_res_plot(data_i, save_folder=save_folder, save_file=save_file)
                            
                        mean, cov = weighted_cov(data_i[["dQx", "dQy", "dQz", "dhbar_w"]].to_numpy(), data_i["Intensity"].to_numpy())
                        
                        res_i = pd.DataFrame([{"channel": c, "analyzer": a, "D_tube": t, "d_i": d, "I": intensity,
                            "Qx_m": mean[0], "Qy_m": mean[1], "Qz_m": mean[2], "hw_m": mean[3], 
                            "sigma_Qx_2": cov[0, 0], "sigma_Qx_Qy": cov[0, 1], "sigma_Qy_2": cov[1, 1], 
                            "sigma_Qx_Qz": cov[0, 2], "sigma_Qy_Qz": cov[1, 2], "sigma_Qz_2": cov[2, 2],
                            "sigma_Qx_hw": cov[0, 3], "sigma_Qy_hw": cov[1, 3], "sigma_Qz_hw": cov[2, 3],
                            "sigma_hw_2": cov[3, 3]}])

                        Res = pd.concat([Res, res_i], ignore_index=True)

        self.Resolution = Res
        if data_save:
            Res.to_csv(os.path.join(self.folder, "Resolution.csv"), index=False)
            
            
    def resolution_ellipse(self, channel, analyzer, tube, pixel):
        Res_i = self.Resolution[(self.Resolution["channel"] == channel) & (self.Resolution["analyzer"] == analyzer) & (self.Resolution["D_tube"] == tube) & (self.Resolution["d_i"] == pixel)]
        
        if len(Res_i) == 0:
            return None
        
        mean = Res_i[["Qx_m", "Qz_m", "hw_m"]].to_numpy()[0]
        cov = Res_i[["sigma_Qx_2", "sigma_Qx_Qz", "sigma_Qz_2", "sigma_Qx_hw", "sigma_Qz_hw", "sigma_hw_2"]].to_numpy()
        cov = flat_to_symmetric(cov, 3)
        
        fig, axs = plt.subplots(1, 3, figsize=(12, 4))
        pairs = [(0, 1), (0, 2), (1, 2)]
        labels = [r"$\Delta Q_x$ [Å]", r"$\Delta Q_z$ [Å]", r"$\Delta \hbar \omega$ [meV]"]

        for k in range(3):
            ax = axs[k]
            ax.cla()
            i, j = pairs[k]
            sub_cov = cov[np.ix_([i, j], [i, j])]
            sub_mean = [mean[i], mean[j]]
            plot_cov_ellipse(sub_cov, sub_mean, ax=ax, edgecolor='blue', facecolor='none')
            ax.scatter(*sub_mean, color='blue')
            ax.set_xlabel(labels[i])
            ax.set_ylabel(labels[j])
            #ax.axis('equal')
            ax.grid(True)
        
        fig.suptitle(f"Channel {channel}, Analyzer {analyzer}, Tube {tube}, Pixel {pixel}")
        plt.savefig(os.path.join(self.folder, "resolution_viewer.png"))
        plt.close()        
        
class Theory_Res:
    """
    
    Class for BIFROST Theoretical Resolution function
    
    """
    
    def __init__(self, dt, hw, stt):
        self.dt = dt
        self.hw = hw
        self.stt = stt
    
    def resolution(self, channel, analyzer, tube, pixel):
        E_fs = [2.7, 3.2, 3.8, 4.4, 5]
        L2s = [1.189,1.314,1.418,1.52,1.622]
        delta_phi = [0.86,0.77,0.72,0.67, 0.62]
        Ef = E_fs[analyzer-1]
        Ei = Ef + self.hw
        
        ki = E2k_(Ei)
        kf = E2k_(Ef)
        
        L1 = 162
        L2 = L2s[analyzer-1]*2
        
        t1 = L1 / E2v_(Ei)
        t2 = L2 / E2v_(Ef)
        
        d = 3.355
        theta_a = theta_Bragg(E2lambda_(Ef),d)
        theta_s = np.deg2rad(self.stt)
        
        delta_phi = np.deg2rad(delta_phi[analyzer-1])
        
        J = np.array([
            [0, ki/L1, 2*Ei/L1],
            [0, ki*t2/(L2*t1), 2*Ei*t2/(L2*t1)],
            [0, -ki/t1, -2*Ei/t1],
            [kf*np.sin(theta_s)/d, t2/t1*ki/d + kf*np.cos(theta_s)/d, 2*Ei*t2/(d*t1)+2*Ef/d],
            [kf*np.cos(theta_s)/np.tan(theta_a), (t2/t1*ki+kf*np.cos(theta_s))/np.tan(theta_a), 2*Ei/np.tan(theta_a)*t2/t1+2*Ef/np.tan(theta_a)],
            [ki, 0, 0],
            [-kf*np.cos(theta_s), kf*np.sin(theta_s), 0]
        ])
        
        sigma = np.diag(np.array([0.01/2.355, 0.002/2.355, np.sqrt(self.dt**2+2e-5**2),0.0001/d, np.deg2rad(3/60)/2.355, 0.025/2.355, delta_phi/2.355]))
        
        cov = J.T @ sigma @ J
        
        self.cov = cov
        
        

def compare_resolution(folder, stt, channel, analyzer, tube, pixel, plot=True, save_file=None):
    """
    Compare the resolution of the BIFROST instrument with the theoretical resolution.
    
    Args:
        folder (str): Path to the folder containing the data files.
        channel (int): Channel number.
        analyzer (int): Analyzer number.
        tube (int): Tube number.
        pixel (int): Pixel number.
        plot (bool): If True, plot the resolution comparison.
        
    Returns:
        None
    """
    
    # Load the data
    data = BIFROST_Res(folder)
    
    # Get the resolution data
    Res = data.Resolution[(data.Resolution["channel"] == channel) & (data.Resolution["analyzer"] == analyzer) & (data.Resolution["D_tube"] == tube) & (data.Resolution["d_i"] == pixel)]
    
    sim_cov = Res[["sigma_Qx_2", "sigma_Qx_Qz", "sigma_Qz_2", "sigma_Qx_hw", "sigma_Qz_hw", "sigma_hw_2"]].to_numpy()
    sim_cov = flat_to_symmetric(sim_cov, 3)
    sim_mean = [Res["Qx_m"].values[0], Res["Qz_m"].values[0], Res["hw_m"].values[0]]
    
    if len(Res) == 0:
        print("No data found for the given parameters.")
        return
    
    # Get the theoretical resolution
    theory = Theory_Res(data.dt, Res["hw_m"].values[0], stt)
    theory.resolution(channel, analyzer, tube, pixel)
    
    theory_cov = theory.cov
    
    if plot:
        fig, axs = plt.subplots(1, 3, figsize=(12, 4), sharex=True, sharey=True)
        pairs = [(0, 1), (0, 2), (1, 2)]
        labels = [r"$\Delta Q_x$ [Å]", r"$\Delta Q_z$ [Å]", r"$\Delta \hbar \omega$ [meV]"]
        
        for k in range(3):
            ax = axs[k]
            ax.cla()
            i, j = pairs[k]
            
            plot_cov_ellipse(sim_cov[np.ix_([i, j], [i, j])], sim_mean, ax=ax, edgecolor='blue', facecolor='none')
            plot_cov_ellipse(theory_cov[np.ix_([i, j], [i, j])], sim_mean, ax=ax, edgecolor='red', facecolor='none')
            ax.scatter(sim_mean[i], sim_mean[j], color='blue', label="Simulated")
            ax.set_xlabel(labels[i])
            ax.set_ylabel(labels[j])
            
        plt.suptitle(f"Resolution hw={Res['hw_m'].values[0]:.2f} dt={data.dt:.2f} stt={stt} {channel}-{analyzer}-{tube}-{pixel}")
        axs[0].legend()
        plt.tight_layout()
        plt.savefig(os.path.join(folder, save_file))
        plt.close()
    
    # Bhattacharyya distance
    
    det_sim = np.linalg.det(sim_cov)
    det_theory = np.linalg.det(theory_cov)
    det_sum = np.linalg.det(sim_cov + theory_cov)/2
    
    D_b = np.log(det_sum/np.sqrt(det_sim*det_theory))
    
    print(f"Bhattacharyya distance: {D_b:.2f}")
    return D_b, sim_cov, sim_mean, theory_cov