"""
Collection of functions for accelerator physics for PhD projects 

Author: Elias Waagaard, elias.walter.waagaard@cern.ch
"""
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import matplotlib as mpl
from scipy.fft import fft, fftfreq
import NAFFlib # Numerical Analysis of the Fundamental Frequencies 

# Initiate class containing all the necessary functions for plotting 
class plot_tools: 
        
    def plot_twiss(fig, twiss, twiss_from_madx=False, plot_magnets=False, also_closed_orbit=False):
        """
        Method to plot Twiss parameters
        As parameter input, can either use Xtrack, or from MAD-X generated Twiss tables 
        """
        if also_closed_orbit:
            spbet = fig.add_subplot(3,1,1)
            spco = fig.add_subplot(3,1,2, sharex=spbet)
            spdisp = fig.add_subplot(3,1,3, sharex=spbet)
        else:
            spbet = fig.add_subplot(2,1,1)
            spdisp = fig.add_subplot(2,1,2, sharex=spbet)
            
        spbet.plot(twiss['s'], twiss['betx'])
        spbet.plot(twiss['s'], twiss['bety'])
        spbet.yaxis.label.set_size(18)

        if also_closed_orbit:
            spco.plot(twiss['s'], twiss['x'])
            spco.plot(twiss['s'], twiss['y'])
            spco.yaxis.label.set_size(18)

        spdisp.plot(twiss['s'], twiss['dx'])
        spdisp.plot(twiss['s'], twiss['dy'])
        spdisp.xaxis.label.set_size(18)
        spdisp.yaxis.label.set_size(18)

        spbet.set_ylabel(r'$\beta_{x,y}$ [m]')
        if also_closed_orbit:
            spco.set_ylabel(r'(Closed orbit)$_{x,y}$ [m]')
        spdisp.set_ylabel(r'$D_{x,y}$ [m]')
        spdisp.set_xlabel('s [m]')

        if not twiss_from_madx:
            fig.suptitle(
                r'$q_x$ = ' f'{twiss["qx"]:.2f}' r' $q_y$ = ' f'{twiss["qy"]:.2f}' '\n'
                r"$Q'_x$ = " f'{twiss["dqx"]:.2f}' r" $Q'_y$ = " f'{twiss["dqy"]:.2f}'
                r' $\gamma_{tr}$ = '  f'{1/np.sqrt(twiss["momentum_compaction_factor"]):.2f}', fontsize=18
            )
        if twiss_from_madx:
            fig.suptitle(
                r'$q_x$ = ' f'{twiss.summary["q1"]:.2f}' r' $q_y$ = ' f'{twiss.summary["q2"]:.2f}' '\n'
                r"$Q'_x$ = " f'{twiss.summary["dq1"]:.2f}' r" $Q'_y$ = " f'{twiss.summary["dq2"]:.2f}'
                r' $\gamma_{tr}$ = '  f'{twiss.summary["gammatr"]:.2f}', fontsize=18
            )            
        
            
        # Plot quadrupoles and dipole magnets if desired  --> still not working...
        if plot_magnets:
            
            # Check if Twiss is dataframe 
            if not isinstance(twiss, pd.DataFrame): 
                twiss = twiss.dframe()
                
            for _, row in twiss.iterrows():
            
                if row['keyword'] == 'quadrupole':
                    _ = spbet.add_patch(
                        mpl.patches.Rectangle(
                            (row['s']-row['l'], 0), row['l'], np.sign(row['k1l']),
                            facecolor='k', edgecolor='k'))
                elif (row['keyword'] == 'rbend' or 
                      row['keyword'] == 'sbend'):
                    _ = spbet.add_patch(
                        mpl.patches.Rectangle(
                            (row['s']-row['l'], -1), row['l'], 2,
                            facecolor='None', edgecolor='k'))
            
        # Use share x ticks and set size
        plt.setp(spbet.get_xticklabels(), visible=False)
        plt.xticks(fontsize=14)
        plt.yticks(fontsize=14)


        fig.subplots_adjust(left=.15, right=.92, hspace=.27)
        fig.tight_layout()
        

    def plot_phase_space_ellipse(fig, tracker, axis='horizontal'):
        """
        Method to plot phase space ellipse from X-suite tracking 
        """
        fig.suptitle('Phase space ellipse',fontsize=18)
        ax = fig.add_subplot(1, 1, 1)  # create an axes object in the figure
        ax.yaxis.label.set_size(20)
        ax.xaxis.label.set_size(20)
        plt.xticks(fontsize=14)  
        plt.yticks(fontsize=14)
        if axis == 'horizontal':
            ax.plot(tracker.record_last_track.x, tracker.record_last_track.px, 'ro', markersize=0.2, alpha=0.3)
            ax.set_ylabel("$p_{x}$")
            ax.set_xlabel("$x$")
        if axis == 'vertical':
            ax.plot(tracker.record_last_track.y, tracker.record_last_track.py, 'bo', markersize=0.2, alpha=0.3)
            ax.set_ylabel("$p_{y}$")
            ax.set_xlabel("$y$")
        fig.tight_layout()
            
            
    def plot_centroid_motion(fig, tracker, axis='horizontal'):
        """
        Method to plot centroid from tracker, either in horizontal or vertical plane 
        """
        fig.suptitle('X-suite tracking - centroid motion',fontsize=18)
        ax = fig.add_subplot(1, 1, 1)  # create an axes object in the figure
        ax.yaxis.label.set_size(20)
        ax.xaxis.label.set_size(20)
        plt.xticks(fontsize=14)  
        plt.yticks(fontsize=14)
        if axis == 'horizontal':
            ax.plot(np.mean(tracker.record_last_track.x, axis=0), marker='o', color='r', markersize=5)
            ax.set_ylabel("Horizontal centroid $X$ [m]")
            ax.set_xlabel("#turns")
        if axis == 'vertical':
            ax.plot(np.mean(tracker.record_last_track.y, axis=0), marker='o', color='b', markersize=5)
            ax.set_ylabel("Vertical centroid $Y$ [m]")
            ax.set_xlabel("#turns")
        fig.tight_layout()
            
            
    def simple_FFT(fig, tracker, axis='horizontal'):
        """
        Method perform simple FFT to find tune 
        -  following Volker Ziemann's example from his book "Hands-on Accelerator Physics using Matlab" 
        Chapter 3: transverse optics 
        """
        ax = fig.add_subplot(1, 1, 1)  # create an axes object in the figure
        num_turns = len(tracker.record_last_track.x[0])
        x_fft = fftfreq(num_turns)[:num_turns//2]  # integer division
        if axis == 'horizontal':
            y_fft = 2*np.abs(fft(np.mean(tracker.record_last_track.x, axis=0)))/num_turns
            ax.set_xlabel("Fractional horizontal tune")
        if axis == 'vertical':
            y_fft = 2*np.abs(fft(np.mean(tracker.record_last_track.y, axis=0)))/num_turns
            ax.set_xlabel("Fractional vertical tune")
        fig.suptitle('FFT spectrum of tracking',fontsize=18)
        ax.yaxis.label.set_size(20)
        ax.xaxis.label.set_size(20)
        plt.xticks(fontsize=14)  
        plt.yticks(fontsize=14)
        ax.plot(x_fft, y_fft[:int(num_turns/2)])  # only plot positive frequencies, i.e. first half of vector   
        ax.set_ylabel("Amplitude [m]")
        fig.tight_layout()
        
        
    def get_tune_footprint(fig, tracker):
        """
        Method to find tune of all particles using NAFF (Numerical Analysis of the Fundamental Frequencies)
        also possiblity to plot working point
        """
        Q_x = np.zeros(len(tracker.record_last_track.x))    
        Q_y = np.zeros(len(tracker.record_last_track.y))
        
        # Iterate over turn-by-turn data to find horizontal and vertical tune of each particle
        for count, particle in enumerate(tracker.record_last_track.x):
            Q_x[count] = NAFFlib.get_tune(particle) 
            
        for count, particle in enumerate(tracker.record_last_track.y):
            Q_y[count] = NAFFlib.get_tune(particle) 
            
        fig.suptitle('Tune footprint',fontsize=18)
        ax = fig.add_subplot(1, 1, 1)  # create an axes object in the figure
        ax.yaxis.label.set_size(20)
        ax.xaxis.label.set_size(20)
        plt.xticks(fontsize=14)  
        plt.yticks(fontsize=14)
        ax.plot(Q_x, Q_y, 'go', markersize=1.5, alpha=0.3)
        ax.set_ylabel("$Q_{y}$")
        ax.set_xlabel("$Q_{x}$")
        #ax.set_xlim(0.15-1e-4, 0.15+1e-4)

        fig.tight_layout()
        
        return Q_x, Q_y, ax
    
# Descripte class of particles to find properties 
class particles:
    
    def print_particle(particle_object):
        """
        Method to print particle properties
        """
        df = particle_object.to_pandas()
        dash = '-' * 55
        print("PARTICLES:\n\n")
        print('{:<27} {:>12}'.format("Property", "Value"))
        print(dash)
        for column in df:
            print('{:<27} {:>12}'.format(df[column].name, df[column].values[0]))
        print(dash)
        print('\n')


# Tools specifically for CPymad and MAD-X
class madx_tools:
    
    def print_seq(madx, seq='SPS'):
        """
        Function to print elements in sequence 
        """
        #Print the elements in the reduced short sequence
        dash = '-' * 65
        print('{:<27} {:>12} {:>15} {:>8}'.format("Element", "Location", "Type", "Length"))
        print(dash)
        for ele in madx.sequence[seq].elements:
            print('{:<27} {:>12.6f} {:>15} {:>8.3}'.format(ele.name, ele.at, ele.base_type.name, ele.length))
            
        return
    

    def print_madx_error():
        """
        If context manager has been used, print the lines of the temporary error file 
        """
        with open('tempfile', 'r') as f:
            lines = f.readlines()
            for ele in lines:
                if '+=+=+= fatal' in ele:
                    print('{}'.format(ele))
                    
    
    def plot_envelope(fig, madx, twiss, seq_name='sps', axis='horizontal', nx=5, ny=5, hcolor="b"):
        """
        Function to plot beam envelope with aperture, can choose horizontal (default) or vertical  
        Returns axis object such that apertures can be plotted in the same plot 
        """
        ax = fig.add_subplot(1, 1, 1)  # create an axes object in the figure
        ax.yaxis.label.set_size(20)
        ax.xaxis.label.set_size(20)
        plt.xticks(fontsize=14)  
        plt.yticks(fontsize=14)

        # Extract beam parameters
        ex = madx.sequence[seq_name].beam.ex
        ey = madx.sequence[seq_name].beam.ey
        sige = madx.sequence[seq_name].beam.sige
    
        # Define some parameters for the beam
        one_sigma_x = np.sqrt(ex*twiss.betx + (sige*twiss.dx)**2) #beam size in x
        one_sigma_y = np.sqrt(ey*twiss.bety + (sige*twiss.dy)**2) #beam size in y
    
        # Choose whether to plot horizontal or vertical axis
        if axis=='horizontal':
            fig.suptitle('Beam envelope - horizontal',fontsize=18)
            ax.plot(twiss.s, twiss.x, color = hcolor)
            ax.set_ylabel("x [m]")
            ax.set_xlabel("s [m]")
            ax.fill_between(twiss.s, twiss.x+nx*one_sigma_x,twiss.x-nx*one_sigma_x, alpha = 0.4, color = hcolor, label='_nolegend_')
        elif axis=='vertical':
            fig.suptitle('Beam envelope - vertical',fontsize=18)
            ax.plot(twiss.s, twiss.y, color = "r")
            ax.set_ylabel("y [m]")
            ax.set_xlabel("s [m]")
            ax.fill_between(twiss.s, twiss.y+ny*one_sigma_y,twiss.y-ny*one_sigma_y, alpha = 0.4, color = "r", label='_nolegend_')        
        else:
            print("Unvalid vertical parameter!")
        fig.tight_layout()
        
        return ax
        
    # ---------------------------------------------- METHODS RELATED TO APERTURES --------------------------------------------------------------------
    
    def get_apertures_real(twiss, axis='horizontal'):
        """
        Method to extract real apertures with sequence element lengths
        """
        pos = list(twiss['s'])
        #Choose whether to plot horizontal or vertical axis:
        if axis=='horizontal':
            aper = list(twiss['aper_1'])
        elif axis=='vertical':
            aper = list(twiss['aper_2'])
        else:
            print("Unvalid axis parameter!")
            
        #Initiate arrays for new aperture 
        new_aper = aper[:] 
        new_pos = pos[:]
        indices = []
    
        #Search for all non-zero aperture elements 
        for i in range(len(aper) - 1, 0, -1):
            if aper[i] != 0:
                new_aper.insert(i, aper[i])
                indices.append(i)
    
        indices = list(reversed(indices))
    
        #Keep track of exact position in new array with counter 
        counter = 0
        for j in indices:
            new_pos.insert(j + counter, (pos[j] - twiss.l[j]))
            counter += 1
        
        #Replace all zeros with Nan
        for i in range(len(new_aper)):
            if new_aper[i] == 0:
                new_aper[i] = np.nan
        
        return np.array(new_pos), np.array(new_aper) 
    
    
    def search_next(i, list):
        """
        Return list without empty elements
        """
        for j in range(i, len(list)):
            if list[j] != 0:
                return list[j]
    
    
    def plot_apertures_real(ax, s, aper, unit="m", offset=None):
        """
        Plot the real aperture, with arrays containing nan values - possibility to add offset 
        """
        aper_flipped = aper.copy()
        for i in range(len(aper_flipped)):
            if aper_flipped[i] != np.nan:
                aper_flipped[i] = -1 * aper_flipped[i] 
        
        if offset is not None:  #Add offset if given
            for i in range(len(aper)):
                aper[i] = aper[i] + + offset[i]
                aper_flipped[i] = aper_flipped[i] + offset[i]
            
        if ax is None:
            ax = plt.gca()
        if not unit == "mm":
            ax.fill_between(s, aper, 0.2, color="black", alpha=0.4, label='_nolegend_')  #previously step="pre", but do not use if entry and exit aperture offset differs 
            ax.fill_between(s, aper_flipped, -0.2, color="black", alpha=0.4, label='_nolegend_') #previously step="pre"
            ax.plot(s, aper, "k", label='_nolegend_')   #previously drawstyle="steps", but do not use if entry and exit aperture offset differs
            ax.plot(s, aper_flipped, "k", label='_nolegend_') #drawstyle="steps"
        else:
            ax.fill_between(s, aper, 0.2 * 1e3, color="black", alpha=0.4, label='_nolegend_')   #previously step="pre"
            ax.fill_between(
                s,
                aper_flipped,
                -0.2 * 1e3,
                color="black",
                label='_nolegend_',
                alpha=0.4,
            )  #previously step="pre"
            ax.plot(s, aper, "k", label='_nolegend_')  #previously drawstyle="steps"
            ax.plot(s, aper_flipped, "k", label='_nolegend_')   #previously drawstyle="steps"
 