LongitudinalImpedance() (mbtrack2.impedance.csr.FreeSpaceCSR method)

LongitudinalWakeFunction() (mbtrack2.impedance.csr.FreeSpaceCSR method)

LongitudinalImpedance() (mbtrack2.impedance.csr.ParallelPlatesCSR method)

LongitudinalWakeFunction() (mbtrack2.impedance.csr.ParallelPlatesCSR method)

add() (mbtrack2.impedance.impedance_model.ImpedanceModel method)

add_global() (mbtrack2.impedance.impedance_model.ImpedanceModel method)

compute_sum() (mbtrack2.impedance.impedance_model.ImpedanceModel method)

compute_sum_names() (mbtrack2.impedance.impedance_model.ImpedanceModel method)

effective_impedance() (mbtrack2.impedance.impedance_model.ImpedanceModel method)

energy_loss() (mbtrack2.impedance.impedance_model.ImpedanceModel method)

group_attributes() (mbtrack2.impedance.impedance_model.ImpedanceModel method)

load() (mbtrack2.impedance.impedance_model.ImpedanceModel method)

plot_area() (mbtrack2.impedance.impedance_model.ImpedanceModel method)

power_loss_spectrum() (mbtrack2.impedance.impedance_model.ImpedanceModel method)

rename_attribute() (mbtrack2.impedance.impedance_model.ImpedanceModel method)

save() (mbtrack2.impedance.impedance_model.ImpedanceModel method)

sum_beta() (mbtrack2.impedance.impedance_model.ImpedanceModel method)

CircularResistiveWall (class in mbtrack2.impedance.resistive_wall)

LongitudinalImpedance() (mbtrack2.impedance.resistive_wall.Coating method)

TransverseImpedance() (mbtrack2.impedance.resistive_wall.Coating method)

long_impedance() (mbtrack2.impedance.resonator.PureInductive method)

long_impedance() (mbtrack2.impedance.resonator.PureResistive method)

long_impedance() (mbtrack2.impedance.resonator.Resonator method)

long_wake_function() (mbtrack2.impedance.resonator.Resonator method)

transverse_impedance() (mbtrack2.impedance.resonator.Resonator method)

transverse_wake_function() (mbtrack2.impedance.resonator.Resonator method)

Z_over_n() (mbtrack2.impedance.tapers.StupakovCircularTaper method)

angle (mbtrack2.impedance.tapers.StupakovCircularTaper property)

radius_prime (mbtrack2.impedance.tapers.StupakovCircularTaper property)

Z0 (mbtrack2.impedance.tapers.StupakovRectangularTaper property)

Z_over_n() (mbtrack2.impedance.tapers.StupakovRectangularTaper method)

angle (mbtrack2.impedance.tapers.StupakovRectangularTaper property)

gap_prime (mbtrack2.impedance.tapers.StupakovRectangularTaper property)

long() (mbtrack2.impedance.tapers.StupakovRectangularTaper method)

quad() (mbtrack2.impedance.tapers.StupakovRectangularTaper method)

xdip() (mbtrack2.impedance.tapers.StupakovRectangularTaper method)

ydip() (mbtrack2.impedance.tapers.StupakovRectangularTaper method)

component_type (mbtrack2.impedance.wakefield.ComplexData property)

initialize_coefficient() (mbtrack2.impedance.wakefield.ComplexData method)

name_and_coefficients_table() (mbtrack2.impedance.wakefield.ComplexData method)

to_wakefunction() (mbtrack2.impedance.wakefield.Impedance method)

add_several_wakefields() (mbtrack2.impedance.wakefield.WakeField static method)

add_wakefields() (mbtrack2.impedance.wakefield.WakeField static method)

append_to_model() (mbtrack2.impedance.wakefield.WakeField method)

impedance_components (mbtrack2.impedance.wakefield.WakeField property)

wake_components (mbtrack2.impedance.wakefield.WakeField property)

deconvolution() (mbtrack2.impedance.wakefield.WakeFunction method)

loss_factor() (mbtrack2.impedance.wakefield.WakeFunction method)

to_impedance() (mbtrack2.impedance.wakefield.WakeFunction method)

lcbi_growth_rate() (in module mbtrack2.instability.instabilities)

lcbi_growth_rate_mode() (in module mbtrack2.instability.instabilities)

lcbi_stability_diagram() (in module mbtrack2.instability.instabilities)

rwmbi_growth_rate() (in module mbtrack2.instability.instabilities)

rwmbi_threshold() (in module mbtrack2.instability.instabilities)

track() (mbtrack2.tracking.aperture.LongitudinalAperture method)

transverse_map_sector_generator() (in module mbtrack2.tracking.element)

init_beam_no_mpi() (mbtrack2.tracking.feedback.FIRDamper method)

distance_between_bunches (mbtrack2.tracking.particles.Beam property)

update_distance_between_bunches() (mbtrack2.tracking.particles.Beam method)

update_filling_pattern() (mbtrack2.tracking.particles.Beam method)

cavity_phasor_record (mbtrack2.tracking.rf.CavityResonator property)

generator_phasor (mbtrack2.tracking.rf.CavityResonator property)

ig_phasor_record (mbtrack2.tracking.rf.CavityResonator property)

init_phasor_track() (mbtrack2.tracking.rf.CavityResonator method)

is_DC_Robinson_stable() (mbtrack2.tracking.rf.CavityResonator method)

plot_DC_Robinson_stability() (mbtrack2.tracking.rf.CavityResonator method)

set_optimal_coupling() (mbtrack2.tracking.rf.CavityResonator method)

set_optimal_detune() (mbtrack2.tracking.rf.CavityResonator method)

DFB_parameter_set() (mbtrack2.tracking.rf.DirectFeedback method)

IIR_init() (mbtrack2.tracking.rf.ProportionalIntegralLoop method)

IIRcutoff (mbtrack2.tracking.rf.ProportionalIntegralLoop property)

Ig2Vg_matrix() (mbtrack2.tracking.rf.ProportionalIntegralLoop method)

init_FFconst() (mbtrack2.tracking.rf.ProportionalIntegralLoop method)

init_Ig2Vg_matrix() (mbtrack2.tracking.rf.ProportionalIntegralLoop method)

get_longitudinal_twiss() (mbtrack2.tracking.synchrotron.Synchrotron method)

get_mcf_order() (mbtrack2.tracking.synchrotron.Synchrotron method)

synchrotron_tune() (mbtrack2.tracking.synchrotron.Synchrotron method)

LongRangeResistiveWall (class in mbtrack2.tracking.wakepotential)

Wdip() (mbtrack2.tracking.wakepotential.LongRangeResistiveWall method)

Wlong() (mbtrack2.tracking.wakepotential.LongRangeResistiveWall method)

get_kick() (mbtrack2.tracking.wakepotential.LongRangeResistiveWall method)

track() (mbtrack2.tracking.wakepotential.LongRangeResistiveWall method)

track_bunch() (mbtrack2.tracking.wakepotential.LongRangeResistiveWall method)

update_tables() (mbtrack2.tracking.wakepotential.LongRangeResistiveWall method)

charge_density() (mbtrack2.tracking.wakepotential.WakePotential method)

check_sampling() (mbtrack2.tracking.wakepotential.WakePotential method)

dipole_moment() (mbtrack2.tracking.wakepotential.WakePotential method)

get_wakepotential() (mbtrack2.tracking.wakepotential.WakePotential method)

plot_gaussian_wake() (mbtrack2.tracking.wakepotential.WakePotential method)

plot_last_wake() (mbtrack2.tracking.wakepotential.WakePotential method)

prepare_wakefunction() (mbtrack2.tracking.wakepotential.WakePotential method)

reduce_sampling() (mbtrack2.tracking.wakepotential.WakePotential method)

reference_loss() (mbtrack2.tracking.wakepotential.WakePotential method)

to_buffer() (mbtrack2.tracking.monitors.monitors.BeamMonitor method)

to_buffer_no_mpi() (mbtrack2.tracking.monitors.monitors.BeamMonitor method)

track() (mbtrack2.tracking.monitors.monitors.BeamMonitor method)

write() (mbtrack2.tracking.monitors.monitors.BeamMonitor method)

write_no_mpi() (mbtrack2.tracking.monitors.monitors.BeamMonitor method)

BeamSpectrumMonitor (class in mbtrack2.tracking.monitors.monitors)

to_buffer() (mbtrack2.tracking.monitors.monitors.BeamSpectrumMonitor method)

track() (mbtrack2.tracking.monitors.monitors.BeamSpectrumMonitor method)

write() (mbtrack2.tracking.monitors.monitors.BeamSpectrumMonitor method)

track() (mbtrack2.tracking.monitors.monitors.BunchMonitor method)

BunchSpectrumMonitor (class in mbtrack2.tracking.monitors.monitors)

to_buffer() (mbtrack2.tracking.monitors.monitors.BunchSpectrumMonitor method)

track() (mbtrack2.tracking.monitors.monitors.BunchSpectrumMonitor method)

write() (mbtrack2.tracking.monitors.monitors.BunchSpectrumMonitor method)

track() (mbtrack2.tracking.monitors.monitors.CavityMonitor method)

file_name (mbtrack2.tracking.monitors.monitors.Monitor property)

monitor_init() (mbtrack2.tracking.monitors.monitors.Monitor method)

to_buffer() (mbtrack2.tracking.monitors.monitors.Monitor method)

track_bunch_data() (mbtrack2.tracking.monitors.monitors.Monitor method)

PhaseSpaceMonitor (class in mbtrack2.tracking.monitors.monitors)

to_buffer() (mbtrack2.tracking.monitors.monitors.PhaseSpaceMonitor method)

track() (mbtrack2.tracking.monitors.monitors.PhaseSpaceMonitor method)

to_buffer() (mbtrack2.tracking.monitors.monitors.ProfileMonitor method)

track() (mbtrack2.tracking.monitors.monitors.ProfileMonitor method)

write() (mbtrack2.tracking.monitors.monitors.ProfileMonitor method)

WakePotentialMonitor (class in mbtrack2.tracking.monitors.monitors)

to_buffer() (mbtrack2.tracking.monitors.monitors.WakePotentialMonitor method)

track() (mbtrack2.tracking.monitors.monitors.WakePotentialMonitor method)

write() (mbtrack2.tracking.monitors.monitors.WakePotentialMonitor method)

plot_beamspectrum() (in module mbtrack2.tracking.monitors.plotting)

plot_bunchdata() (in module mbtrack2.tracking.monitors.plotting)

plot_bunchspectrum() (in module mbtrack2.tracking.monitors.plotting)

plot_cavitydata() (in module mbtrack2.tracking.monitors.plotting)

plot_phasespacedata() (in module mbtrack2.tracking.monitors.plotting)

plot_profiledata() (in module mbtrack2.tracking.monitors.plotting)

streak_beamdata() (in module mbtrack2.tracking.monitors.plotting)

streak_beamspectrum() (in module mbtrack2.tracking.monitors.plotting)

streak_bunchspectrum() (in module mbtrack2.tracking.monitors.plotting)

BeamLoadingEquilibrium (class in mbtrack2.utilities.beamloading)

PTBL_threshold() (mbtrack2.utilities.beamloading.BeamLoadingEquilibrium method)

center_of_mass() (mbtrack2.utilities.beamloading.BeamLoadingEquilibrium method)

dV() (mbtrack2.utilities.beamloading.BeamLoadingEquilibrium method)

ddV() (mbtrack2.utilities.beamloading.BeamLoadingEquilibrium method)

deltaVRF() (mbtrack2.utilities.beamloading.BeamLoadingEquilibrium method)

du_dz() (mbtrack2.utilities.beamloading.BeamLoadingEquilibrium method)

energy_balance() (mbtrack2.utilities.beamloading.BeamLoadingEquilibrium method)

integrate_func() (mbtrack2.utilities.beamloading.BeamLoadingEquilibrium method)

plot_dV() (mbtrack2.utilities.beamloading.BeamLoadingEquilibrium method)

plot_rho() (mbtrack2.utilities.beamloading.BeamLoadingEquilibrium method)

plot_voltage() (mbtrack2.utilities.beamloading.BeamLoadingEquilibrium method)

rho() (mbtrack2.utilities.beamloading.BeamLoadingEquilibrium method)

std_rho() (mbtrack2.utilities.beamloading.BeamLoadingEquilibrium method)

to_solve() (mbtrack2.utilities.beamloading.BeamLoadingEquilibrium method)

u() (mbtrack2.utilities.beamloading.BeamLoadingEquilibrium method)

uexp() (mbtrack2.utilities.beamloading.BeamLoadingEquilibrium method)

voltage() (mbtrack2.utilities.beamloading.BeamLoadingEquilibrium method)

tune_shift_from_effective_impedance() (in module mbtrack2.utilities.misc)

change_values() (mbtrack2.utilities.optics.PhysicalModel method)

plot_aperture() (mbtrack2.utilities.optics.PhysicalModel method)

plot_resistivity() (mbtrack2.utilities.optics.PhysicalModel method)

read_IW2D_folder() (in module mbtrack2.utilities.read_impedance)

gaussian_bunch_spectrum() (in module mbtrack2.utilities.spectrum)

# -*- coding: utf-8 -*-
"""
Define coherent synchrotron radiation (CSR) wakefields in various models.
"""
import mpmath as mp
import numpy as np
from scipy.constants import c, mu_0
from scipy.special import gamma

from mbtrack2.impedance.wakefield import Impedance, WakeField, WakeFunction

class FreeSpaceCSR(WakeField):
    """
    Free space steady-state coherent synchrotron radiation Wakefield element,
    based on [1]. 
    
    Impedance is computed using Eq. (A10) of [2].

    Parameters
    ----------
    time : array of float
        Time points where the wake function will be evaluated in [s].
    frequency : array of float
        Frequency points where the impedance will be evaluated in [Hz].
    length : float 
        Length of the impedacen to consider in [m].
    radius : float 
        Dipole radius of curvature in [m].
        
    References
    [1] : Faltens, A. N. D. R. I. S., & Laslett, L. J. (1973). Longitudinal 
    coupling impedance of a stationary electron ring in a cylindrical 
    geometry. part. Accel., 4, 151-157.
    [2] : Agoh, T., and K. Yokoya. "Calculation of coherent synchrotron 
    radiation using mesh." Physical Review Special Topics-Accelerators and 
    Beams 7.5 (2004): 054403.

    """

    def __init__(self, time, frequency, length, radius):
        super().__init__()

        self.length = length
        self.radius = radius
        self.Z0 = mu_0 * c

        Zl = self.LongitudinalImpedance(frequency)
        # Wl = self.LongitudinalWakeFunction(time)

        Zlong = Impedance(variable=frequency,
                          function=Zl,
                          component_type='long')
        # Wlong = WakeFunction(variable = time, function = Wl, component_type="long")

        super().append_to_model(Zlong)
        # super().append_to_model(Wlong)

    def LongitudinalImpedance(self, frequency):
        """
        Compute the free space steady-state CSR impedance.
        Based on Eq. (A10) of [1].
        
        This formula is valid only if omega << (3 * gamma^3 * c) / (2 * R).

        Parameters
        ----------
        frequency : float array
            Frequency in [Hz].

        Returns
        -------
        Zl : complex array
            Longitudinal impedance in [ohm].
            
        References
        [1] : Agoh, T., and K. Yokoya. "Calculation of coherent synchrotron 
        radiation using mesh." Physical Review Special Topics-Accelerators and 
        Beams 7.5 (2004): 054403.
        """

        Zl = (self.Z0 * self.length / (2 * np.pi) * gamma(2 / 3) *
              (-1j * 2 * np.pi * frequency /
               (3 * c * self.radius**2))**(1 / 3))
        return Zl

    def LongitudinalWakeFunction(self, time):
        raise NotImplementedError

class ParallelPlatesCSR(WakeField):
    """
    Perfectly conducting parallel plates steady-state coherent synchrotron 
    radiation Wakefield element, based on [1]. 

    distance : float
        Vertical distance between the parallel plates in [m].
        
    Attributes
    threshold : float
        Shielding threshold in the parallel plates model in [Hz].
        
    [1] : Agoh, T., and K. Yokoya. "Calculation of coherent synchrotron 

    """

    def __init__(self, time, frequency, length, radius, distance):

        self.distance = distance

    @property
    def threshold(self):
        """Shielding threshold in the parallel plates model in [Hz]."""
        return (3*c) / (2 * np.pi) * (self.radius / self.distance**3)**0.5

    def LongitudinalImpedance(self, frequency, tol=1e-5):
        """
        Compute the CSR impedance using the perfectly conducting parallel 
        plates steady-state model.
        
        Impedance is computed using Eq. (A1) of [1].

        tol : float, optinal
            Desired maximum final error on sum_func. 

        Returns
        -------
            
        """

        Zl = np.zeros(frequency.shape, dtype=complex)
        constant = (2 * np.pi * self.Z0 * self.length / self.distance *
                    (2 / self.radius)**(1 / 3))
        for i, f in enumerate(frequency):
            k = 2 * mp.pi * f / c

            sum_value = mp.nsum(lambda p: self.sum_func(p, k), [0, mp.inf],
                                tol=tol,
                                method='r+s')

            Zl[i] = constant * (1 / k)**(1 / 3) * complex(sum_value)

        return Zl

    def sum_func(self, p, k):
        """
        Utility function for LongitudinalImpedance.

        p : int
        k : float

        Returns
        -------
        sum_value : mpc

        """
        xp = (2*p + 1) * mp.pi / self.distance * (self.radius / 2 / k**2)**(1 /
                                                                            3)
        Ai = mp.airyai(xp**2)
        Bi = mp.airybi(xp**2)
        Aip = mp.airyai(xp**2, 1)
        Bip = mp.airybi(xp**2, 1)
        return Aip * (Aip + 1j*Bip) + xp**2 * Ai * (Ai + 1j*Bi)

"""
Module where the ImpedanceModel class is defined.
"""
import pickle
from copy import deepcopy

import matplotlib.pyplot as plt
import pandas as pd
from matplotlib import colormaps
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
from scipy.integrate import trapz
from scipy.interpolate import interp1d

from mbtrack2.impedance.wakefield import WakeField, WakeFunction
from mbtrack2.utilities.misc import (
    beam_loss_factor,
    double_sided_impedance,
    effective_impedance,
)
from mbtrack2.utilities.spectrum import (
    beam_spectrum,
    gaussian_bunch_spectrum,
    spectral_density,
)

class ImpedanceModel():
    """
    Define the impedance model of the machine.

    The model must be completed with successive add(...) and add_global(...)
    calls, then compute_sum() must be run.

    The transverse impedance and wake functions are beta weighted and divided
    by the beta at the tracking location (ring.optics.local_beta).

    ring : Synchrotron object

    wakefields : list of WakeField objects
        WakeFields in the model.
    positions : list of arrays
        Positions corresponding the different WakeFields.
    names : list of str
        Names of the WakeField objects.
    sum : WakeField
        Sum of every WakeField in the model weigthed by beta functions.
    sum_"name" : WakeField
        Sum of the "name" Wakefield weigthed by beta functions.
    sum_names : array
        Names of attributes where the WakeFields are summed by name.
    globals : list of WakeField objects
        Globals WakeFields in the model.
    globals_names : list of str
        Names of the global WakeFields objects.

    Methods
    -------
    add(wakefield, positions, name)
        Add the same WakeField object at different locations to the model.
    add_global(wakefield, name)
        Add a "global" WakeField object which will added to the sum WakeField
        but not weighted by beta functions.
    sum_beta(wake, beta)
        Weight a WakeField object by an array of beta functions.
    compute_sum_names()
        Compute the weighted WakeField for each WakeField object type.
    compute_sum()
        Compute the sum of all weighted WakeField into self.sum.
    plot_area(Z_type="Zlong", component="real", sigma=None, attr_list=None)
        Plot the contributions of different kind of WakeFields.
    save(file)
        Save impedance model to file.
    load(file)
        Load impedance model from file.
    """

    def __init__(self, ring):
        self.ring = ring
        self.optics = self.ring.optics
        self.wakefields = []
        self.positions = []
        self.names = []
        self.globals = []
        self.globals_names = []
        self.sum_names = []

    def add(self, wakefield, positions, name=None):
        """

        wakefield : WakeField
            WakeField object to add to the model.
        positions : array, float or int
            Array of longitudinal positions where the elements are loacted.
        name : str, optional
            Name of the element type. If None, the name of the WakeField object
            is used. The default is None.

        Returns
        -------
        None.

        """
        if name is None:
            name = wakefield.name
            raise ValueError("Please give a valid name.")
        if name not in self.names:
            self.names.append(name)
        else:
            raise ValueError("This name is already taken.")
        self.wakefields.append(wakefield)
        self.positions.append(positions)

    def add_global(self, wakefield, name=None):
        """

        To use with "distributed" elements, for example a resistive wall
        wakefield computed from an effective radius (so which has already been
        weighted by beta functions).

        Returns
        -------
        None.

        """
        self.globals.append(wakefield)
        if name not in self.globals_names:
            self.globals_names.append(name)
        else:
        # setattr(self, name, wakefield)

    def sum_beta(self, wake, beta):
        """

        wake : WakeField
            WakeField element object.
        beta : array of shape (2, N)
            Beta function at the locations of the elements.

        Returns
        -------
        wake_sum : WakeField
            WakeField object weighted by beta functions.

        """
        wake_sum = deepcopy(wake)
        local_beta = self.ring.optics.local_beta
        for component_name in wake_sum.components:
            comp = getattr(wake_sum, component_name)
            weight = ((beta[0, :]**comp.power_x) * (beta[1, :]**comp.power_y))
            if comp.plane == "x":
                weight = weight.sum() / local_beta[0]
            if comp.plane == "y":
                weight = weight.sum() / local_beta[1]
            else:
                weight = weight.sum()
            setattr(wake_sum, component_name, weight * comp)
        return wake_sum

    def compute_sum_names(self):
        """
        The new summed WakeField object are set to into self.sum_name.
        """
        for idx, wake in enumerate(self.wakefields):
            attribute_name = "sum_" + self.names[idx]
            beta = self.optics.beta(self.positions[idx])
            wake_sum = self.sum_beta(wake, beta)
            wake_sum.name = attribute_name
            setattr(self, attribute_name, wake_sum)
            self.sum_names.append(attribute_name)

    def compute_sum(self):
        """
        """
        self.compute_sum_names()
        for i, name in enumerate(self.sum_names):
            if i == 0:
                self.sum = deepcopy(getattr(self, name))
                self.sum.name = "sum"
            else:
                wake2 = getattr(self, name)
                for component_name2 in wake2.components:
                    comp2 = getattr(wake2, component_name2)
                    try:
                        comp1 = getattr(self.sum, component_name2)
                        setattr(self.sum, component_name2, comp1 + comp2)
                    except AttributeError:
                        setattr(self.sum, component_name2, comp2)
        for i, wake2 in enumerate(self.globals):
            name = self.globals_names[i]
            setattr(self, name, wake2)
            self.sum_names.append(name)
            if not hasattr(self, "sum"):
                self.sum = deepcopy(wake2)
            else:
                    try:

    def group_attributes(self,
                         string_in_name,
                         names_to_group=None,
                         property_list=['Zlong']):
        """Groups attributes in the ImpedanceModel based on a given string pattern.
        Args:
            string_in_name (str): The string pattern used to match attribute names for grouping. If names_to_group is given, this is a name of a new attribute instead.
            names_to_group (list, optional): List of attribute names to be explicitly grouped.
                                            If not provided, attributes matching the string pattern will be automatically grouped.
                                            Defaults to None.
            property_list (list, optional): List of property names to be grouped for each attribute.
                                            Defaults to ['Zlong'].
        Returns:
            int: Returns 0 indicating the successful grouping of attributes.
        Notes:
            - This method groups attributes in the ImpedanceModel class based on a given string pattern.
            - If names_to_group is not provided, it automatically selects attributes from the existing sum_names that match the given string pattern.
            - A new WakeField instance is created and assigned to the attribute named 'string_in_name'.
            - The specified properties from the first attribute in names_to_group are appended to the new WakeField instance.
            - The values of the specified properties from the remaining attributes in names_to_group are added to the corresponding properties in the new WakeField instance.
            - The names_to_group attributes are removed from the ImpedanceModel class.
            - The new grouped attribute is added to sum_names, and the first attribute in names_to_group is removed from sum_names."""
        attrs = self.sum_names
        if names_to_group == None:
            names_to_group = []
            for attr in attrs:
                if string_in_name in attr:
                    names_to_group.append(attr)
        setattr(self, string_in_name, WakeField())
        for prop in property_list:
            getattr(self, string_in_name).append_to_model(
                getattr(getattr(self, names_to_group[0]), prop))
        for attr in names_to_group[1:]:
            for prop in property_list:
                old_values = getattr(getattr(self, string_in_name), prop)
                new_values = getattr(getattr(self, attr), prop)
                setattr(getattr(self, string_in_name), prop,
                        old_values + new_values)
            self.sum_names.remove(attr)
            delattr(self, attr)
        self.sum_names.append(string_in_name)
        self.sum_names.remove(names_to_group[0])
        delattr(self, names_to_group[0])
        return 0

    def rename_attribute(self, old_name, new_name):
        """Renames an attribute in the ImpedanceModel.
        Args:
            old_name (str): The current name of the attribute to be renamed.
            new_name (str): The new name for the attribute.
        Raises:
            KeyError: If the old_name doesn't exist as an attribute in the ImpedanceModel.
        Notes:
            - This method renames an attribute in the ImpedanceModel class.
            - The attribute with the old_name is removed from the class's dictionary (__dict__) using the pop() method.
            - The attribute is then added back to the class's dictionary with the new_name using the __dict__ attribute.
            - If the old_name doesn't exist as an attribute in the ImpedanceModel, a KeyError is raised.
        """
        self.__dict__[new_name] = self.__dict__.pop(old_name)

    def plot_area(self,
                  Z_type="Zlong",
                  component="real",
                  sigma=None,
                  attr_list=None,
                  zoom=False):
        """

        Z_type : str, optional
            Type of impedance to plot.
        component : str, optional
            Component to plot, can be "real" or "imag". 
        sigma : float, optional
            RMS bunch length in [s] to use for the spectral density. If equal
            to None, the spectral density is not plotted.
        attr_list : list or array of str, optional
            Attributes to plot.
        zoom : bool
            If True, add a zoomed plot on top right corner.

        """
        if attr_list is None:
            attr_list = self.sum_names

        # manage legend
        Ztype_dict = {
            "Zlong": 0,
            "Zxdip": 1,
            "Zydip": 2,
            "Zxquad": 3,
            "Zyquad": 4
        }
        scale = [1e-3, 1e-6, 1e-6, 1e-6, 1e-6]
        label_list = [
            r"$Z_\mathrm{long} \; (\mathrm{k}\Omega)$",
            r"$\frac{1}{\beta_0} \sum_{j} \beta_{x,j} Z_{x,j}^\mathrm{Dip} \; (\mathrm{M}\Omega/\mathrm{m})$",
            r"$\frac{1}{\beta_0} \sum_{j} \beta_{y,j} Z_{y,j}^\mathrm{Dip} \; (\mathrm{M}\Omega/\mathrm{m})$",
            r"$\frac{1}{\beta_0} \sum_{j} \beta_{x,j} Z_{x,j}^\mathrm{Quad} \; (\mathrm{M}\Omega/\mathrm{m})$",
            r"$\frac{1}{\beta_0} \sum_{j} \beta_{y,j} Z_{y,j}^\mathrm{Quad} \; (\mathrm{M}\Omega/\mathrm{m})$"
        ]
        leg = Ztype_dict[Z_type]

        # sort plot by decresing area size
        area = np.zeros((len(attr_list), ))
        for index, attr in enumerate(attr_list):
            try:
                sum_imp = getattr(getattr(self, attr), Z_type)
                area[index] = trapz(sum_imp.data[component],
                                    sum_imp.data.index)
            except AttributeError:
                pass
        sorted_index = area.argsort()[::-1]

        # Init fig
        fig = plt.figure()
        ax = fig.gca()
        zero_impedance = getattr(self.sum, Z_type) * 0
        total_imp = 0
        legend = []

        if sigma is not None:
            legend.append("Spectral density for sigma = " + str(sigma) + " s")

        # Main plot
        colorblind = colormaps['tab10'].colors
        for index in sorted_index:
            attr = attr_list[index]
            # Set all impedances with common indexes using + zero_impedance
            try:
                sum_imp = getattr(getattr(self, attr), Z_type) + zero_impedance
                ax.fill_between(sum_imp.data.index * 1e-9,
                                total_imp,
                                total_imp +
                                sum_imp.data[component] * scale[leg],
                                edgecolor=colorblind[index % 10],
                                color=colorblind[index % 10])
                total_imp += sum_imp.data[component] * scale[leg]
                if attr[:4] == "sum_":
                    legend.append(attr[4:])
                else:
                    legend.append(attr)
                pass

            spect = spectral_density(zero_impedance.data.index, sigma)
            spect = spect / spect.max() * total_imp.max()
            ax.plot(sum_imp.data.index * 1e-9, spect, 'r', linewidth=2.5)

        ax.legend(legend, loc="upper left", ncol=2)
        ax.set_xlabel("Frequency (GHz)")
        ax.set_ylabel(label_list[leg] + " - " + component + " part")
        ax.set_title(label_list[leg] + " - " + component + " part")

        if zoom is True:
            in_ax = inset_axes(
                ax,
                width="30%",  # width = 30% of parent_bbox
                height=1.5,  # height : 1 inch
                loc=1)

            total_imp = 0
            for index in sorted_index:
                attr = attr_list[index]
                # Set all impedances with common indexes using + zero_impedance
                try:
                    sum_imp = getattr(getattr(self, attr),
                                      Z_type) + zero_impedance
                    in_ax.fill_between(sum_imp.data.index * 1e-3,
                                       total_imp,
                                       total_imp +
                                       sum_imp.data[component] * 1e-9,
                                       edgecolor=colorblind[index % 10],
                                       color=colorblind[index % 10])
                    total_imp += sum_imp.data[component] * 1e-9
                except AttributeError:
                    pass
            in_ax.set_xlim([0, 200])
            in_ax.set_xlabel("Frequency (kHz)")
            in_ax.set_ylabel(r"$[\mathrm{G}\Omega]$")

        return fig

    def effective_impedance(self,
                            m,
                            mu,
                            sigma,
                            M,
                            tuneS,
                            xi=None,
                            mode="Hermite"):
        """
        Compute the longitudinal and transverse effective impedance.

        mu : int
            coupled bunch mode number, goes from 0 to (M-1) where M is the
            number of bunches
        m : int
            head-tail (or azimutal/synchrotron) mode number
        sigma : float
            RMS bunch length in [s]
        M : int
            Number of bunches.
        tuneS : float
            Synchrotron tune.
        xi : float, optional
            (non-normalized) chromaticity
        mode: str, optional
            type of the mode taken into account for the computation:
            -"Hermite" modes for Gaussian bunches

        Returns
        -------
        summary : DataFrame
            Longitudinal and transverse effective impedance.

        """

        attr_list = self.sum_names

        eff_array = np.zeros((len(attr_list), 3), dtype=complex)

        for i, attr in enumerate(attr_list):
            try:
                impedance = getattr(getattr(self, attr), "Zlong")
                eff_array[i,
                          0] = effective_impedance(self.ring, impedance, m, mu,
                                                   sigma, M, tuneS, xi, mode)
                pass

            try:
                impedance = getattr(getattr(self, attr), "Zxdip")
                          1] = effective_impedance(self.ring, impedance, m, mu,
                pass

            try:
                impedance = getattr(getattr(self, attr), "Zydip")
                          2] = effective_impedance(self.ring, impedance, m, mu,
                pass

        eff_array[:, 0] = eff_array[:, 0] * self.ring.omega0 * 1e3
        eff_array[:, 1] = eff_array[:, 1] * 1e-3
        eff_array[:, 2] = eff_array[:, 2] * 1e-3

        summary = pd.DataFrame(eff_array,
                               index=attr_list,
                               columns=[
                                   "Z/n [mOhm]", "sum betax x Zxeff [kOhm]",
                                   "sum betay x Zyeff [kOhm]"
                               ])

        return summary

    def energy_loss(self, sigma, M, bunch_spacing, I, n_points=10e6):
        """
        Compute the beam and bunch loss factor and energy losses for each type 
        of element in the model assuming Gaussian bunches and constant spacing
        between bunches.

            RMS bunch length in [s].
        M : int
            Number of bunches in the beam.
        bunch_spacing : float
            Time between two bunches in [s].
        I : float
            Total beam current in [A].
        n_points : float, optional
            Number of points used in the frequency spectrums.

        Returns
        -------
        summary : Dataframe
            Contains the beam and bunch loss factor and energy loss for the 
            full model and for each type of different component.

        """

        fmax = self.sum.Zlong.data.index.max()
        fmin = self.sum.Zlong.data.index.min()

        Q = I * self.ring.T0 / M

        if fmin >= 0:
            fmin = -1 * fmax
        f = np.linspace(fmin, fmax, int(n_points))

        beam_spect = beam_spectrum(f, M, bunch_spacing, sigma=sigma)

        bunch_spect = gaussian_bunch_spectrum(f, sigma)

        loss_array = np.zeros((len(attr_list), 2))

            try:
                loss_array[i, 0] = beam_loss_factor(impedance, f, beam_spect,
                                                    self.ring)
                loss_array[i, 1] = beam_loss_factor(impedance, f, bunch_spect,
                pass

        loss_array = loss_array * 1e-12
        summary = pd.DataFrame(loss_array,
                               columns=[
                                   "loss factor (beam) [V/pC]",
                                   "loss factor (bunch) [V/pC]"
                               ])

        summary["P (beam) [W]"] = summary["loss factor (beam) [V/pC]"] * \
            1e12*Q**2/(self.ring.T0)
        summary["P (bunch) [W]"] = summary["loss factor (bunch) [V/pC]"] * \
            1e12*Q**2/(self.ring.T0)*M

    def power_loss_spectrum(self,
                            sigma,
                            M,
                            bunch_spacing,
                            I,
                            n_points=10e6,
                            max_overlap=False,
                            plot=False):
        """
        Compute the power loss spectrum of the summed longitudinal impedance 
        as in Eq. (4) of [1].

        Assume Gaussian bunches and constant spacing between bunches.

        M : int
        I : float
            Number of points used in the frequency spectrum.
        max_overlap : bool, optional
            If True, the bunch spectrum (scaled to the number of bunches) is 
            used instead of the beam spectrum to compute the maximum value of 
            the power loss spectrum at each frequency. Should only be used to 
            maximise the power loss at a given frequency (e.g. for HOMs) and 
            not for the full spectrum.
        plot : bool, optional
            If True, plots:
                - the overlap between the real part of the longitudinal impedance 
                and the beam spectrum.
                - the power loss spectrum.

        Returns
        -------
        pf0 : array
            Frequency points.
        power_loss : array
            Power loss spectrum in [W].

        [1] : L. Teofili, et al. "A Multi-Physics Approach to Simulate the RF 
        Heating 3D Power Map Induced by the Proton Beam in a Beam Intercepting 
        Device", in IPAC'18, 2018, doi:10.18429/JACoW-IPAC2018-THPAK093

        """

        impedance = self.sum.Zlong
        fmax = impedance.data.index.max()
        fmin = impedance.data.index.min()

            double_sided_impedance(impedance)

        frequency = np.linspace(fmin, fmax, int(n_points))
        if max_overlap is False:
            spectrum = beam_spectrum(frequency, M, bunch_spacing, sigma)
        else:
            spectrum = gaussian_bunch_spectrum(frequency, sigma) * M

        pmax = np.floor(fmax / self.ring.f0)
        pmin = np.floor(fmin / self.ring.f0)

        p = np.arange(pmin + 1, pmax)
        pf0 = p * self.ring.f0
        ReZ = np.real(impedance(pf0))
        spectral_density = np.abs(spectrum)**2
        # interpolation of the spectrum is needed to avoid problems liked to
        # division by 0
        # computing the spectrum directly to the frequency points gives
        # wrong results
        spect = interp1d(frequency, spectral_density)
        power_loss = (self.ring.f0 * Q)**2 * ReZ * spect(pf0)

        if plot is True:
            fig, ax = plt.subplots()
            twin = ax.twinx()
            p1, = ax.plot(pf0, ReZ, color="r", label="Re[Z] [Ohm]")
            p2, = twin.plot(pf0,
                            spect(pf0) * (self.ring.f0 * Q)**2,
                            color="b",
                            label="Beam spectrum [a.u.]")
            ax.set_xlabel("Frequency [Hz]")
            ax.set_ylabel("Re[Z] [Ohm]")
            twin.set_ylabel("Beam spectrum [a.u.]")
            plots = [p1, p2]
            ax.legend(handles=plots, loc="best")

            ax.plot(pf0, power_loss)
            ax.set_ylabel("Power loss [W]")

        return pf0, power_loss

    def save(self, file):
        """

        The same pandas version is needed on both saving and loading computer
        for the pickle to work.

        file : str
            File where the impedance model is saved.

        Returns
        -------
        None.

        """
        to_save = {
            "wakefields": self.wakefields,
            "positions": self.positions,
            "names": self.names,
            "globals": self.globals,
            "globals_names": self.globals_names
        }
        with open(file, "wb") as f:
            pickle.dump(to_save, f)

    def load(self, file):
        """

        file : str

        Returns
        -------
        None.

        """

        with open(file, 'rb') as f:
            to_load = pickle.load(f)

        self.wakefields = to_load["wakefields"]
        self.positions = to_load["positions"]
        self.names = to_load["names"]
        self.globals = to_load["globals"]
        self.globals_names = to_load["globals_names"]
        self.compute_sum()

"""
Define resistive wall elements based on the WakeField class.
"""

from scipy.constants import c, epsilon_0, mu_0
from scipy.integrate import quad

def skin_depth(frequency, rho, mu_r=1, epsilon_r=1):
    """
    General formula for the skin depth.
    
        Frequency points in [Hz].
    rho : float
        Resistivity in [ohm.m].
    mu_r : float, optional
        Relative magnetic permeability.
    epsilon_r : float, optional
        Relative electric permittivity.
    
    Returns
    -------
    delta : array of float
        Skin depth in [m].
    
    """

    delta = (np.sqrt(2 * rho / (np.abs(2 * np.pi * frequency) * mu_r * mu_0)) *
             np.sqrt(
                 np.sqrt(1 + (rho * np.abs(2 * np.pi * frequency) * epsilon_r *
                              epsilon_0)**2) +
                 rho * np.abs(2 * np.pi * frequency) * epsilon_r * epsilon_0))
    return delta

class CircularResistiveWall(WakeField):
    """
    Resistive wall WakeField element for a circular beam pipe.
    
    Impedance from approximated formulas from Eq. (2.77) of Chao book [1].
    Wake function formulas from [2].
        
        Beam pipe length in [m].
    rho : float
        Beam pipe radius in [m].
    exact : bool, optional
        If False, approxmiated formulas are used for the wake function 
        computations.
    atol : float, optional
        Absolute tolerance used to enforce fundamental theorem of beam loading
        for the exact expression of the longitudinal wake function.
        Default is 1e-20.
    
    [1] : Chao, A. W. (1993). Physics of collective beam instabilities in high 
    energy accelerators. Wiley.
    [2] : Skripka, Galina, et al. "Simultaneous computation of intrabunch and 
    interbunch collective beam motions in storage rings." Nuclear Instruments 
    and Methods in Physics Research Section A: Accelerators, Spectrometers, 
    Detectors and Associated Equipment 806 (2016): 221-230.

    """

    def __init__(self,
                 time,
                 frequency,
                 length,
                 rho,
                 radius,
                 exact=False,
                 atol=1e-20):

        self.rho = rho
        self.t0 = (2 * self.rho * self.radius**2 / self.Z0)**(1 / 3) / c

        omega = 2 * np.pi * frequency
        Z1 = length * (1 + np.sign(frequency) * 1j) * rho / (
            2 * np.pi * radius * skin_depth(frequency, rho))
        Z2 = c / omega * length * (1 + np.sign(frequency) * 1j) * rho / (
            np.pi * radius**3 * skin_depth(frequency, rho))

        Wl = self.LongitudinalWakeFunction(time, exact, atol)
        Wt = self.TransverseWakeFunction(time, exact)

                          function=Z1,
        Zxdip = Impedance(variable=frequency,
                          function=Z2,
                          component_type='xdip')
        Zydip = Impedance(variable=frequency,
                          component_type='ydip')
        Wlong = WakeFunction(variable=time, function=Wl, component_type="long")
        Wxdip = WakeFunction(variable=time, function=Wt, component_type="xdip")
        Wydip = WakeFunction(variable=time, function=Wt, component_type="ydip")

        super().append_to_model(Zxdip)
        super().append_to_model(Zydip)
        super().append_to_model(Wlong)
        super().append_to_model(Wxdip)
        super().append_to_model(Wydip)

    def LongitudinalWakeFunction(self, time, exact=False, atol=1e-20):
        """
        Compute the longitudinal wake function of a circular resistive wall 
        using Eq. (22), or approxmiated expression Eq. (24), of [1]. The 
        approxmiated expression is valid if the time is large compared to the 
        characteristic time t0.
        
        If some time value is smaller than atol, then the fundamental theorem 
        of beam loading is applied: Wl(0) = Wl(0+)/2.

        time : array of float
            Time points where the wake function is evaluated in [s].
        exact : bool, optional
            If True, the exact expression is used. The default is False.
        atol : float, optional
            Absolute tolerance used to enforce fundamental theorem of beam loading
            for the exact expression of the longitudinal wake function.
            Default is 1e-20.

        Returns
        -------
        wl : array of float
            Longitudinal wake function in [V/C].
            
        [1] : Skripka, Galina, et al. "Simultaneous computation of intrabunch and 
        interbunch collective beam motions in storage rings." Nuclear Instruments 
        and Methods in Physics Research Section A: Accelerators, Spectrometers, 
        Detectors and Associated Equipment 806 (2016): 221-230.
        """
        wl = np.zeros_like(time)
        idx1 = time < 0
        wl[idx1] = 0
        if exact == True:
            idx2 = time > 20 * self.t0
            idx3 = np.logical_not(np.logical_or(idx1, idx2))
            wl[idx3] = self.__LongWakeExact(time[idx3], atol)
        else:
            idx2 = np.logical_not(idx1)
        wl[idx2] = self.__LongWakeApprox(time[idx2])
        return wl

    def TransverseWakeFunction(self, time, exact=False):
        """
        Compute the transverse wake function of a circular resistive wall 
        using Eq. (25), or approxmiated expression Eq. (26), of [1]. The 
        
        Exact expression (Eq. (25) from [1]) is corrected by factor (c * t0).

        Returns
        -------
        wt : array of float
            Transverse wake function in [V/C].
            
        """
        wt = np.zeros_like(time)
        wt[idx1] = 0
            wt[idx3] = self.__TransWakeExact(time[idx3])
        else:
        wt[idx2] = self.__TransWakeApprox(time[idx2])
        return wt

    def __LongWakeExact(self, time, atol):
        factor = 4 * self.Z0 * c / (np.pi * self.radius**2) * self.length
        for i, t in enumerate(time):
            val, err = quad(lambda z: self.__function(t, z), 0, np.inf)
            wl[i] = factor * (
                np.exp(-t / self.t0) / 3 * np.cos(np.sqrt(3) * t / self.t0) -
                np.sqrt(2) / np.pi * val)
            if np.isclose(0, t, atol=atol):
                wl[i] = wl[i] / 2
        return wl

    def __TransWakeExact(self, time):
        factor = ((8 * self.Z0 * c**2 * self.t0) / (np.pi * self.radius**4) *
                  self.length)
            val, err = quad(lambda z: self.__function2(t, z), 0, np.inf)
            wt[i] = factor * (
                1 / 12 *
                (-1 * np.exp(-t / self.t0) * np.cos(np.sqrt(3) * t / self.t0) +
                 np.sqrt(3) * np.exp(-t / self.t0) *
                 np.sin(np.sqrt(3) * t / self.t0)) - np.sqrt(2) / np.pi * val)
        return wt

    def __LongWakeApprox(self, t):
        wl = -1 * (1 / (4 * np.pi * self.radius) * np.sqrt(self.Z0 * self.rho /
                                                           (c * np.pi)) /
                   t**(3 / 2)) * self.length
        return wl

    def __TransWakeApprox(self, t):
        wt = (1 / (np.pi * self.radius**3) *
              np.sqrt(self.Z0 * c * self.rho / np.pi) / t**(1 / 2) *
              self.length)
        return wt

    def __function(self, t, x):
        return ((x**2 * np.exp(-1 * (x**2) * t / self.t0)) / (x**6 + 8))

    def __function2(self, t, x):
        return ((-1 * np.exp(-1 * (x**2) * t / self.t0)) / (x**6 + 8))

class Coating(WakeField):

                 length,
                 rho1,
                 rho2,
                 radius,
                 thickness,
                 approx=False):
        """
        WakeField element for a coated circular beam pipe.
        
        The longitudinal and tranverse impedances are computed using formulas
        from [1].

        f : array of float
            Frequency points where the impedance is evaluated in [Hz].
        length : float
            Length of the beam pipe to consider in [m].
        rho1 : float
            Resistivity of the coating in [ohm.m].
        rho2 : float
            Resistivity of the bulk material in [ohm.m].
        radius : float
            Radius of the beam pipe to consier in [m].
        thickness : float
            Thickness of the coating in [m].
        approx : bool, optional
            If True, used approxmiated formula. The default is False.

        [1] : Migliorati, M., E. Belli, and M. Zobov. "Impact of the resistive 
        wall impedance on beam dynamics in the Future Circular e+ e− Collider." 
        Physical Review Accelerators and Beams 21.4 (2018): 041001.

        """

        self.rho1 = rho1
        self.rho2 = rho2
        self.thickness = thickness

        Zl = self.LongitudinalImpedance(frequency, approx)
        Zt = self.TransverseImpedance(frequency, approx)

                          function=Zt,

    def LongitudinalImpedance(self, f, approx):
        """
        Compute the longitudinal impedance of a coating using Eq. (5), or 
        approxmiated expression Eq. (8), of [1]. The approxmiated expression 
        is valid if the skin depth of the coating is large compared to the 
        coating thickness. 

        approx : bool
            If True, used approxmiated formula.

        Returns
        -------
        Zl : array
            Longitudinal impedance values in [ohm].
            
        """

        Z0 = mu_0 * c
        factor = Z0 * f / (2 * c * self.radius) * self.length
        skin1 = skin_depth(f, self.rho1)
        skin2 = skin_depth(f, self.rho2)

        if approx == False:
            alpha = skin1 / skin2
            tanh = np.tanh((1 + 1j * np.sign(f)) * self.thickness / skin1)
            bracket = ((np.sign(f) + 1j) * skin1 * (alpha*tanh + 1) /
                       (alpha+tanh))
        else:
            valid_approx = self.thickness / np.min(skin1)
            if valid_approx < 0.01:
                print(
                    "Approximation is not valid. Returning impedance anyway.")
            bracket = ((np.sign(f) + 1j) * skin2 + 2 * 1j * self.thickness *
                       (1 - self.rho2 / self.rho1))

        Zl = factor * bracket

        return Zl

    def TransverseImpedance(self, f, approx):
        """
        Compute the transverse impedance of a coating using Eq. (6), or 
        approxmiated expression Eq. (9), of [1]. The approxmiated expression 

        Returns
        -------
        Zt : array
            Transverse impedance values in [ohm].
            
        """

        Z0 = mu_0 * c
        factor = Z0 / (2 * np.pi * self.radius**3) * self.length

            bracket = ((1 + 1j * np.sign(f)) * skin1 * (alpha*tanh + 1) /
        else:
                print(
            bracket = ((1 + 1j * np.sign(f)) * skin2 +
                       2 * 1j * self.thickness * np.sign(f) *

        Zt = factor * bracket

        return Zt

"""
This module defines the impedances and wake functions from the resonator model 
based on the WakeField class.
"""

class Resonator(WakeField):

    def __init__(self, time, frequency, Rs, fr, Q, plane, atol=1e-20):
        """
        Resonator model WakeField element which computes the impedance and the 
        wake function in both longitudinal and transverse case.
        
        ! For transverse case, if Q < 2, there is a kick factor mismatch 
        between time and frequency domain.

            Time points where the wake function will be evaluated in [s].
        frequency : array of float
            Frequency points where the impedance will be evaluated in [Hz].
        Rs : float
            Shunt impedance in [ohm] in longitudinal or [ohm/m] in transverse.
        fr : float
            Resonance frequency in [Hz].
        Q : float
            Quality factor.
        plane : str or list
            Plane on which the resonator is used: "long", "x" or "y".
            Absolute tolerance used to enforce fundamental theorem of beam 
            loading for the exact expression of the longitudinal wake function.
            
        [1]  B. W. Zotter and S. A. Kheifets, "Impedances and Wakes in High-Energy 
        Particle Ac-celerators", Eq. (3.10) and (3.15), pp.51-53.

        """
        self.Rs = Rs
        self.fr = fr
        self.wr = 2 * np.pi * self.fr
        self.Q = Q
        if isinstance(plane, str):
            self.plane = [plane]
        elif isinstance(plane, list):
            self.plane = plane

        if self.Q >= 0.5:
            self.Q_p = np.sqrt(self.Q**2 - 0.25)
        else:
            self.Q_p = np.sqrt(0.25 - self.Q**2)
        self.wr_p = (self.wr * self.Q_p) / self.Q

        for dim in self.plane:
            if dim == "long":
                Zlong = Impedance(variable=frequency,
                                  function=self.long_impedance(frequency),
                                  component_type="long")
                super().append_to_model(Zlong)
                Wlong = WakeFunction(variable=time,
                                     function=self.long_wake_function(
                                         time, atol),
                                     component_type="long")
                super().append_to_model(Wlong)

            elif dim == "x" or dim == "y":
                Zdip = Impedance(variable=frequency,
                                 function=self.transverse_impedance(frequency),
                                 component_type=dim + "dip")
                super().append_to_model(Zdip)
                Wdip = WakeFunction(
                    variable=time,
                    function=self.transverse_wake_function(time),
                    component_type=dim + "dip")
                super().append_to_model(Wdip)
            else:
                raise ValueError("Plane must be: long, x or y")

    def long_wake_function(self, t, atol):
            wl = ((self.wr * self.Rs / self.Q) * np.exp(-1 * self.wr * t /
                                                        (2 * self.Q)) *
                  (np.cos(self.wr_p * t) - np.sin(self.wr_p * t) /
                   (2 * self.Q_p)))
        elif self.Q < 0.5:
                  (np.cosh(self.wr_p * t) - np.sinh(self.wr_p * t) /
        if np.any(np.abs(t) < atol):
            wl[np.abs(t) < atol] = wl[np.abs(t) < atol] / 2
        if np.any(t < -atol):
            wl[t < -atol] = 0
        return wl

    def long_impedance(self, f):
        return self.Rs / (1 + 1j * self.Q * (f / self.fr - self.fr / f))

    def transverse_impedance(self, f):
        return self.Rs * self.fr / f / (1 + 1j * self.Q *
                                        (f / self.fr - self.fr / f))

    def transverse_wake_function(self, t):
            wt = (self.wr * self.Rs / self.Q_p *
                  np.exp(-1 * t * self.wr / 2 / self.Q_p) *
                  np.sin(self.wr_p * t))
        else:
                  np.sinh(self.wr_p * t))
        if np.any(t < 0):
            wt[t < 0] = 0
        return wt

class PureInductive(WakeField):
    """
    Pure inductive Wakefield element which computes associated longitudinal 
    impedance and wake function.
    
    L : float
        Inductance value in [Ohm/Hz].
    n_wake : int or float, optional
        Number of points used in the wake function.
    n_imp : int or float, optional
        Number of points used in the impedance.
    imp_freq_lim : float, optional
        Maximum frequency used in the impedance. 
    nout, trim : see Impedance.to_wakefunction
    """

                 L,
                 n_wake=1e6,
                 n_imp=1e6,
                 imp_freq_lim=1e11,
                 nout=None,
                 trim=False):
        self.L = L
        self.n_wake = int(n_wake)
        self.n_imp = int(n_imp)
        self.imp_freq_lim = imp_freq_lim

        freq = np.linspace(start=1, stop=self.imp_freq_lim, num=self.n_imp)
        imp = Impedance(variable=freq,
                        function=self.long_impedance(freq),
                        component_type="long")
        super().append_to_model(imp)

        wf = imp.to_wakefunction(nout=nout, trim=trim)
        super().append_to_model(wf)

        return 1j * self.L * f

class PureResistive(WakeField):
    """
    Pure resistive Wakefield element which computes associated longitudinal 
    
    R : float
        Resistance value in [Ohm].
    """

                 R,
        self.R = R

        return self.R

"""
Module where taper elements are defined.
"""

from scipy.constants import c, mu_0, pi

from mbtrack2.impedance.wakefield import Impedance, WakeField

class StupakovRectangularTaper(WakeField):
    """
    Rectangular vertical taper WakeField element, using the low frequency 
    approxmiation. Assume constant taper angle. Formulas from [1].
    
    ! Valid for low w
    
    frequency: frequency points where the impedance will be evaluated in [Hz]
    gap_entrance : full vertical gap at taper entrance in [m]
    gap_exit: full vertical gap at taper exit in [m]
    length: taper length in [m]
    width : full horizontal width of the taper in [m]
    """

                 gap_entrance,
                 gap_exit,
                 length,
                 width,
                 m_max=100,
                 n_points=int(1e4)):

        self.frequency = frequency
        self.gap_entrance = gap_entrance
        self.gap_exit = gap_exit
        self.width = width
        self.m_max = m_max
        self.n_points = n_points

                          function=self.long(),
                          function=self.xdip(),
                          function=self.ydip(),
        Zxquad = Impedance(variable=frequency,
                           function=-1 * self.quad(),
                           component_type='xquad')
        Zyquad = Impedance(variable=frequency,
                           function=self.quad(),
                           component_type='yquad')

        super().append_to_model(Zxquad)
        super().append_to_model(Zyquad)

    @property
    def gap_prime(self):
        return (self.gap_entrance - self.gap_exit) / self.length

    @property
    def angle(self):
        return np.arctan(
            (self.gap_entrance / 2 - self.gap_exit / 2) / self.length)

    @property
    def Z0(self):
        return mu_0 * c

    def long(self, frequency=None):

        if frequency is None:
            frequency = self.frequency

        def F(x, m_max):
            m = np.arange(0, m_max)
            phi = np.outer(pi * x / 2, 2*m + 1)
            val = 1 / (2*m + 1) / (np.cosh(phi)**2) * np.tanh(phi)
            return val.sum(1)

        z = np.linspace(0, self.length, self.n_points)
        g = np.linspace(self.gap_entrance, self.gap_exit, self.n_points)

        to_integrate = self.gap_prime**2 * F(g / self.width, self.m_max)
        integral = trapz(to_integrate, x=z)

        return -1j * frequency * self.Z0 / (2*c) * integral

    def Z_over_n(self, f0):
        return np.imag(self.long(1)) * f0

    def ydip(self):

        def G1(x, m_max):
            val = (2*m + 1) / (np.sinh(phi)**2) / np.tanh(phi)
            val = x[:, None]**3 * val

        to_integrate = self.gap_prime**2 / (g**3) * G1(g / self.width,
                                                       self.m_max)

        return -1j * pi * self.width * self.Z0 / 4 * integral

    def xdip(self):

        def G3(x, m_max):
            phi = np.outer(pi * x, m)
            val = 2 * m / (np.cosh(phi)**2) * np.tanh(phi)
            val = x[:, None]**2 * val

        to_integrate = self.gap_prime**2 / (g**2) * G3(g / self.width,

        return -1j * pi * self.Z0 / 4 * integral

    def quad(self):

        def G2(x, m_max):
            val = (2*m + 1) / (np.cosh(phi)**2) * np.tanh(phi)

        to_integrate = self.gap_prime**2 / (g**2) * G2(g / self.width,

class StupakovCircularTaper(WakeField):
    """
    Circular taper WakeField element, using the low frequency 
    
    radius_entrance : radius at taper entrance in [m]
    radius_exit : radius at taper exit in [m]
    length : taper length in [m]
    """

                 radius_entrance,
                 radius_exit,
                 length,

        self.radius_entrance = radius_entrance
        self.radius_exit = radius_exit

                          function=self.dip(),

    @property
            (self.radius_entrance - self.radius_exit) / self.length)

    @property
    def radius_prime(self):
        return (self.radius_entrance - self.radius_exit) / self.length

    @property

        return (self.Z0 /
                (2*pi) * np.log(self.radius_entrance / self.radius_exit) -
                1j * self.Z0 * frequency /
                (2*c) * self.radius_prime**2 * self.length)

    def dip(self, frequency=None):

        r = np.linspace(self.radius_entrance, self.radius_exit, self.n_points)

        to_integrate = self.radius_prime**2 / (r**2)

        return (self.Z0 * c / (4 * pi**2 * frequency) *
                (1 / (self.radius_exit**2) - 1 /
                 (self.radius_entrance**2)) - 1j * self.Z0 / (2*pi) * integral)

"""
This module defines general classes to describes wakefields, impedances and 
wake functions.
"""

import warnings

import scipy as sc
from scipy.constants import c

class ComplexData:
    """
    Define a general data structure for a complex function based on a pandas 
    DataFrame.

    variable : list, numpy array
       contains the function variable values

    function : list or numpy array of comp lex numbers
        contains the values taken by the complex function
    """

                 variable=np.array([-1e15, 1e15]),
                 function=np.array([0, 0])):
        self.data = pd.DataFrame(
            {
                'real': np.real(function),
                'imag': np.imag(function)
            },
            index=variable)
        self.data.index.name = 'variable'

    def add(self,
            structure_to_add,
            method='zero',
            interp_kind='cubic',
            index_name="variable"):
        """
        Method to add two structures. If the data don't have the same length,
        different cases are handled.
        
        structure_to_add : ComplexData object, int, float or complex.
            structure to be added.
        method : str, optional
            manage how the addition is done, possibilties are: 
            -common: the common range of the index (variable) from the two
                structures is used. The input data are cross-interpolated
                so that the result data contains the points at all initial
                points in the common domain.
            -extrapolate: extrapolate the value of both ComplexData.
            -zero: outside of the common range, the missing values are zero. In 
                the common range, behave as "common" method.
        interp_kind : str, optional
            interpolation method which is passed to pandas and to
            scipy.interpolate.interp1d.
        index_name : str, optional
            name of the dataframe index passed to the method
            
        Returns
        -------
        ComplexData 
            Contains the sum of the two inputs.
        """

        # Create first two new DataFrames merging the variable
        # from the two input data

        if isinstance(structure_to_add, (int, float, complex)):
            structure_to_add = ComplexData(
                variable=self.data.index,
                function=(structure_to_add * np.ones(len(self.data.index))))

        data_to_concat = structure_to_add.data.index.to_frame().set_index(
            index_name)

        initial_data = pd.concat([self.data, data_to_concat], sort=True)
        initial_data = initial_data[~initial_data.index.duplicated(
            keep='first')]
        initial_data = initial_data.sort_index()

        data_to_add = pd.concat([
            structure_to_add.data,
            self.data.index.to_frame().set_index(index_name)
        ],
                                sort=True)
        data_to_add = data_to_add[~data_to_add.index.duplicated(keep='first')]
        data_to_add = data_to_add.sort_index()

        if method == 'common':
            max_variable = min(structure_to_add.data.index.max(),
                               self.data.index.max())

            min_variable = max(structure_to_add.data.index.min(),
                               self.data.index.min())

            initial_data = initial_data.interpolate(method=interp_kind)
            mask = ((initial_data.index <= max_variable)
                    & (initial_data.index >= min_variable))
            initial_data = initial_data[mask]

            data_to_add = data_to_add.interpolate(method=interp_kind)
            mask = ((data_to_add.index <= max_variable)
                    & (data_to_add.index >= min_variable))
            data_to_add = data_to_add[mask]

            result_structure = ComplexData()
            result_structure.data = initial_data + data_to_add
            return result_structure

        if method == 'extrapolate':
            print('Not there yet')
            return self

        if method == 'zero':

            initial_data[mask] = initial_data[mask].interpolate(
                method=interp_kind)

            data_to_add[mask] = data_to_add[mask].interpolate(

            initial_data.replace(to_replace=np.nan, value=0, inplace=True)
            data_to_add.replace(to_replace=np.nan, value=0, inplace=True)

    def __radd__(self, structure_to_add):
        return self.add(structure_to_add, method='zero')

    def __add__(self, structure_to_add):

    def multiply(self, factor):
        """
        Multiply a data strucure with a float or an int.
        If the multiplication is done with something else, throw a warning.
        """
        if isinstance(factor, (int, float)):
            result_structure.data = self.data * factor
        else:
            warnings.warn(('The multiplication factor is not a float '
                           'or an int.'), UserWarning)

    def __mul__(self, factor):
        return self.multiply(factor)

    def __rmul__(self, factor):

    def __call__(self, values, interp_kind="cubic"):
        """
        Interpolation of the data by calling the class to have a function-like
        behaviour.
        
        values : list or numpy array of complex, int or float
            values to be interpolated.
            interpolation method which is passed to scipy.interpolate.interp1d.
            
        Returns
        -------
        numpy array 
            Contains the interpolated data.
        """
        real_func = interp1d(x=self.data.index,
                             y=self.data["real"],
                             kind=interp_kind)
        imag_func = interp1d(x=self.data.index,
                             y=self.data["imag"],
        return real_func(values) + 1j * imag_func(values)

    def __repr__(self):
        """Return representation of data"""
        return f'{(self.data)!r}'

    def initialize_coefficient(self):
        """
        Define the impedance coefficients and the plane of the impedance that
        are presents in attributes of the class.
        """
        table = self.name_and_coefficients_table()

        try:
            component_coefficients = table[self.component_type].to_dict()
        except KeyError:
            print('Component type {} does not exist'.format(
                self.component_type))
            raise

        self.a = component_coefficients["a"]
        self.b = component_coefficients["b"]
        self.c = component_coefficients["c"]
        self.d = component_coefficients["d"]
        self.plane = component_coefficients["plane"]

    def name_and_coefficients_table(self):
        """
        Return a table associating the human readbale names of an impedance
        component and its associated coefficients and plane.
        """

        component_dict = {
            'long': {
                'a': 0,
                'b': 0,
                'c': 0,
                'd': 0,
                'plane': 'z'
            },
            'xcst': {
                'a': 0,
                'b': 0,
                'c': 0,
                'd': 0,
                'plane': 'x'
            },
            'ycst': {
                'a': 0,
                'b': 0,
                'c': 0,
                'd': 0,
                'plane': 'y'
            },
            'xdip': {
                'a': 1,
                'b': 0,
                'c': 0,
                'd': 0,
            },
            'ydip': {
                'a': 0,
                'b': 1,
                'c': 0,
                'd': 0,
            },
            'xydip': {
                'a': 0,
                'b': 1,
                'c': 0,
                'd': 0,
            },
            'yxdip': {
                'a': 1,
                'b': 0,
                'c': 0,
                'd': 0,
            },
            'xquad': {
                'a': 0,
                'b': 0,
                'c': 1,
                'd': 0,
            },
            'yquad': {
                'a': 0,
                'b': 0,
                'c': 0,
                'd': 1,
            },
            'xyquad': {
                'a': 0,
                'b': 0,
                'c': 0,
                'd': 1,
            },
            'yxquad': {
                'a': 0,
                'b': 0,
                'c': 1,
                'd': 0,
            },
        }

        return pd.DataFrame(component_dict)

    @property
    def power_x(self):
        power_x = self.a / 2 + self.c / 2.
        if self.plane == 'x':
            power_x += 1. / 2.
        return power_x

    @property
    def power_y(self):
        power_y = self.b / 2. + self.d / 2.
        if self.plane == 'y':
            power_y += 1. / 2.
        return power_y

    @property
    def component_type(self):
        return self._component_type

    @component_type.setter
    def component_type(self, value):
        self._component_type = value
        self.initialize_coefficient()

class WakeFunction(ComplexData):
    """
    Define a WakeFunction object based on a ComplexData object.
    
    variable : array-like
        Time domain structure of the wake function in [s].
    function : array-like
        Wake function values in [V/C] for "long" and in [V/C/m] for others.
    component_type : str, optinal
        Type of the wake function: "long", "xdip", "xquad", ...
    
    data : DataFrame
    wake_type : str
    
    Methods
    -------
    to_impedance()
        Return an Impedance object from the WakeFunction data.
    deconvolution(freq_lim, sigma, mu)
        Compute a wake function from wake potential data.
    plot()
        Plot the wake function data.
    loss_factor(sigma)
        Compute the loss factor or the kick factor for a Gaussian bunch.
    """

                 function=np.array([0, 0]),
                 component_type='long'):
        super().__init__(variable, function)
        self._component_type = component_type
        self.data.index.name = "time [s]"

        return f'WakeFunction of component_type {self.component_type}:\n {(self.data)!r}'

    def add(self, structure_to_add, method='zero'):
        """
        Method to add two WakeFunction objects. The two structures are
        compared so that the addition is self-consistent.
        """

            result = super().add(structure_to_add,
                                 method=method,
                                 index_name="time [s]")
        elif isinstance(structure_to_add, WakeFunction):
            if (self.component_type == structure_to_add.component_type):
                result = super().add(structure_to_add,
                                     method=method,
                                     index_name="time [s]")
            else:
                warnings.warn(('The two WakeFunction objects do not have the '
                               'same coordinates or plane or type. '
                               'Returning initial WakeFunction object.'),
                              UserWarning)
                result = self

        wake_to_return = WakeFunction(result.data.index,
                                      result.data.real.values,
                                      self.component_type)
        return wake_to_return

        return self.add(structure_to_add)

        """
        Multiply a WakeFunction object with a float or an int.
        """
        result = super().multiply(factor)

    def to_impedance(self, freq_lim, nout=None, sigma=None, mu=None):
        """
        The WakeFunction data is assumed to be sampled equally.
        
        If both sigma and mu are given, the impedance data is divided by the
        Fourier transform of a Gaussian distribution. It then can be used to go 
        from wake potential data to impedance.

        freq_lim : float
            The maximum frequency for calculating the impedance [Hz].
        nout : int or float, optional
            Length of the transformed axis of the output. If None, it is taken
            to be 2*(n-1) where n is the length of the input. If nout > n, the
            input is padded with zeros. If nout < n, the input it cropped.
            Note that, for any nout value, nout//2+1 input points are required.
            The default is None.
            RMS of the Gaussian distribution in [s].
        mu : float
            Offset of the Gaussian distribution in [s].

        """
        tau = np.array(self.data.index)
        wp = np.array(self.data["real"])

        if nout is None:
            nout = len(tau)
        else:
            nout = int(nout)

        # FT of wake potential
        sampling = tau[1] - tau[0]
        freq = sc.fft.rfftfreq(nout, sampling)
        dtau = (tau[-1] - tau[0]) / len(tau)
        dft_wake = sc.fft.rfft(wp, n=nout, axis=0) * dtau

        i_limit = freq < freq_lim
        freq_trun = freq[i_limit]
        dft_wake_trun = dft_wake[i_limit]

        if (sigma is not None) and (mu is not None):
            dft_rho_trun = np.exp(-0.5*(sigma*2*np.pi*freq_trun)**2 + \
                                      1j*mu*2*np.pi*freq_trun)
        else:
            dft_rho_trun = -1

        if self.component_type == "long":
            imp = dft_wake_trun / dft_rho_trun * -1
        elif (self.component_type == "xdip") or (self.component_type
                                                 == "ydip"):
            imp = dft_wake_trun / dft_rho_trun * -1j
        else:
            raise NotImplementedError(self.component_type + " is not correct.")

        imp = Impedance(variable=freq_trun,
                        function=imp,
                        component_type=self.component_type)

        return imp

    def deconvolution(self, freq_lim, sigma, mu, nout=None):
        """
                
        The present data loaded into the WakeFunction object is assumed to be 
        an uniformly sampled wake potential from a Gaussian bunch distribution 
        with RMS bunch length sigma and offset mu.
        
        The offset mu depends on the code used to compute the wake potential:
            - for CST, it is the first time sample, so self.data.index[0].
            - for ABCI, it is "-1 * ISIG * SIG / c".

        mu : float

        """
        imp = self.to_impedance(freq_lim, sigma=sigma, mu=mu)
        wf = imp.to_wakefunction(nout=nout)
        return wf

    def plot(self):
        """

        Returns
        -------
        ax : class:`matplotlib.axes.Axes`
            Wake function data.

        """
        ax = self.data["real"].plot()
            unit = " [V/C]"
        else:
            unit = " [V/C/m]"
        label = r"$W_{" + self.component_type + r"}$" + unit
        ax.set_ylabel(label)
        return ax

    def loss_factor(self, sigma):
        """
        Compute the loss factor or the kick factor for a Gaussian bunch. 
        Formulas from Eq. (5.6), Eq. (5.12) and Eq. (5.17) of [1].
        
        Returns
        -------
        kloss: float
            Loss factor in [V/C] or kick factor in [V/C/m] depanding on the 
            impedance type.
            
        [1] : Zotter, Bruno W., and Semyon A. Kheifets. Impedances and wakes 
        in high-energy particle accelerators. World Scientific, 1998.
        
        """
        time = self.data.index
        S = 1 / (2 * np.sqrt(np.pi) * sigma) * np.exp(-1 * time**2 /
                                                      (4 * sigma**2))
        kloss = trapz(x=time, y=self.data["real"] * S)

        return kloss

class Impedance(ComplexData):
    """
    Define an Impedance object based on a ComplexData object.
    
        Frequency domain structure of the impedance in [Hz].
        Impedance values in [Ohm].
        Type of the impedance: "long", "xdip", "xquad", ...
    
    impedance_type : str
    
    Methods
    -------
    to_wakefunction()
        Return a WakeFunction object from the impedance data.
    plot()
        Plot the impedance data.
    """

        self.data.index.name = 'frequency [Hz]'

        return f'Impedance of component_type {self.component_type}:\n {(self.data)!r}'

    def add(self, structure_to_add, beta_x=1, beta_y=1, method='zero'):
        """
        Method to add two Impedance objects. The two structures are
        Beta functions can be precised as well.
        """

                                 index_name="frequency [Hz]")
        elif isinstance(structure_to_add, Impedance):
                weight = (beta_x**self.power_x) * (beta_y**self.power_y)
                result = super().add(structure_to_add * weight,
                                     index_name="frequency [Hz]")
            else:
                warnings.warn(('The two Impedance objects do not have the '
                               'Returning initial Impedance object.'),

        impedance_to_return = Impedance(
            result.data.index,
            result.data.real.values + 1j * result.data.imag.values,
            self.component_type)
        return impedance_to_return

        """
        Multiply a Impedance object with a float or an int.
        """

        """
        Formulas from Eq. (2) p239 and Eq.(7) p240 of [1].
        
        Returns
        -------
            
        [1] : Handbook of accelerator physics and engineering, 3rd printing.
        """

        positive_index = self.data.index > 0
        frequency = self.data.index[positive_index]

        # Import here to avoid circular import
        from mbtrack2.utilities import spectral_density
        sd = spectral_density(frequency, sigma, m=0)

        if (self.component_type == "long"):
            kloss = trapz(x=frequency,
                          y=2 * self.data["real"][positive_index] * sd)
        elif (self.component_type == "xdip" or self.component_type == "ydip"):
                          y=2 * self.data["imag"][positive_index] * sd)
        elif (self.component_type == "xquad"
              or self.component_type == "yquad"):
        else:
            raise TypeError("Impedance type not recognized.")

    def to_wakefunction(self, nout=None, trim=False):
        """
        The impedance data is assumed to be sampled equally.
    
        trim : bool or float, optional
            If True, the pseudo wake function is trimmed at time=0 and is forced 
            to zero where time<0. 
            If False, the original result is preserved.
            If a float is given, the pseudo wake function is trimmed from 
            time <= trim to 0. 
        """

        Z0 = (self.data['real'] + self.data['imag'] * 1j)
        Z = Z0[~np.isnan(Z0)]

        if self.component_type != "long":
            Z = Z / 1j

        freq = Z.index
        fs = (freq[-1] - freq[0]) / len(freq)
        sampling = freq[1] - freq[0]

            nout = len(Z)
        else:

        time_array = sc.fft.fftfreq(nout, sampling)
        Wlong_raw = sc.fft.irfft(np.array(Z), n=nout, axis=0) * nout * fs

        time = sc.fft.fftshift(time_array)
        Wlong = sc.fft.fftshift(Wlong_raw)

        if trim is not False:
            i_neg = np.where(time < trim)[0]
            Wlong[i_neg] = 0

        wf = WakeFunction(variable=time,
                          function=Wlong,
                          component_type=self.component_type)
        return wf

        """

        Returns
        -------
            Impedance data.

        """
        ax = self.data.plot()
            unit = " [ohm]"
        else:
            unit = " [ohm/m]"
        label = r"$Z_{" + self.component_type + r"}$" + unit
        return ax

class WakeField:
    """
    Defines a WakeField which corresponds to a single physical element which 
    produces different types of wakes, represented by their Impedance or 
    WakeFunction objects.
    
    structure_list : list, optional
        list of Impedance/WakeFunction objects to add to the WakeField.
    name : str, optional
    
    impedance_components : array of str
        Impedance components present for this element.
    wake_components : array of str
        WakeFunction components present for this element.
    components : array of str
        Impedance or WakeFunction components present for this element.
        
    Methods
    -------
    append_to_model(structure_to_add)
        Add Impedance/WakeFunction to WakeField.
    list_to_attr(structure_list)
        Add list of Impedance/WakeFunction to WakeField.
    add_wakefields(wake1, beta1, wake2, beta2)
        Add two WakeFields taking into account beta functions.
    add_several_wakefields(wakefields, beta)
        Add a list of WakeFields taking into account beta functions.
    drop(component)
        Delete a component or a list of component from the WakeField.
        Save WakeField element to file.
        Load WakeField element from file.
    """

    def __init__(self, structure_list=None, name=None):
        self.list_to_attr(structure_list)
        self.name = name

        """Return representation of WakeField components."""
        return f'WakeField {self.name} with components:\n {(list(self.components))!r}'

        return self._add(structure_to_add)

    def __iter__(self):
        """Iterate over components"""
        comp = [getattr(self, comp) for comp in self.components]
        return comp.__iter__()

    def _add(self, structure_to_add):
        """Allow to add two WakeField element with different components."""
        for comp in structure_to_add.components:
            self.append_to_model(getattr(structure_to_add, comp))
        return self

    def append_to_model(self, structure_to_add):
        """
        Add Impedance/WakeFunction component to WakeField.

        structure_to_add : Impedance or WakeFunction object
        """
        list_of_attributes = dir(self)
        if isinstance(structure_to_add, Impedance):
            attribute_name = "Z" + structure_to_add.component_type
            if attribute_name in list_of_attributes:
                raise ValueError("There is already a component of the type "
                                 "{} in this element.".format(attribute_name))
            else:
                self.__setattr__(attribute_name, structure_to_add)
            attribute_name = "W" + structure_to_add.component_type
            else:
        else:
            raise ValueError(
                "{} is not an Impedance nor a WakeFunction.".format(
                    structure_to_add))

    def list_to_attr(self, structure_list):
        """
        Add list of Impedance/WakeFunction components to WakeField.

        structure_list : list of Impedance or WakeFunction objects.
        """
        if structure_list is not None:
            for component in structure_list:
                self.append_to_model(component)

    @property
    def impedance_components(self):
        """
        Return an array of the impedance component names for the element.
        """
        valid = [
            "Zlong", "Zxdip", "Zydip", "Zxquad", "Zyquad", "Zxcst", "Zycst"
        ]
        return np.array([comp for comp in dir(self) if comp in valid])

    @property
    def wake_components(self):
        """
        Return an array of the wake function component names for the element.
        """
        valid = [
            "Wlong", "Wxdip", "Wydip", "Wxquad", "Wyquad", "Wxcst", "Wycst"
        ]

    @property
    def components(self):
        """
        Return an array of the component names for the element.
        """
        valid = [
            "Wlong", "Wxdip", "Wydip", "Wxquad", "Wyquad", "Wxcst", "Wycst",
        ]

    def drop(self, component):
        """

        component : str or list of str
            Component or list of components to drop.
            If "Z", drop all impedance components.
            If "W"", drop all wake function components.

        Returns
        -------
        None.

        """
        if component == "Z":
            component = self.impedance_components
        elif component == "W":
            component = self.wake_components

        if isinstance(component, str):
            delattr(self, component)
        elif isinstance(component, list) or isinstance(component, np.ndarray):
            for comp in component:
                delattr(self, comp)
        else:
            raise TypeError("component should be a str or list of str.")

        """
        
        file : str
            File where the WakeField element is saved.

        Returns
        -------
        None.

        """
            pickle.dump(self, f)

    @staticmethod
    def load(file):
        """
        
        file : str

        Returns
        -------
            Loaded wakefield.

        """
            wakefield = pickle.load(f)

        return wakefield

    def add_wakefields(wake1, beta1, wake2, beta2):
        """

        wake1 : WakeField
            WakeField to add.
        beta1 : array of shape (2,)
            Beta functions at wake1 postion.
        wake2 : WakeField
        beta2 : array of shape (2,)
            Beta functions at wake2 postion.

        Returns
        -------
            WakeField with summed components.

        """
        wake_sum = deepcopy(wake1)
        for component_name1 in wake1.components:
            comp1 = getattr(wake_sum, component_name1)
            weight1 = ((beta1[0]**comp1.power_x) * (beta1[1]**comp1.power_y))
            setattr(wake_sum, component_name1, weight1 * comp1)

        for component_name2 in wake2.components:
            comp2 = getattr(wake2, component_name2)
            weight2 = ((beta2[0]**comp2.power_x) * (beta2[1]**comp2.power_y))
            try:
                comp1 = getattr(wake_sum, component_name2)
                setattr(wake_sum, component_name2, comp1 + weight2*comp2)
                setattr(wake_sum, component_name2, weight2 * comp2)

    def add_several_wakefields(wakefields, beta):
        """
        
        wakefields : list of WakeField
            WakeFields to add.
        beta : array of shape (len(wakefields), 2)
            Beta functions at the WakeField postions.

        Returns
        -------

        """
        if len(wakefields) == 1:
            return wakefields[0]
        elif len(wakefields) > 1:
            wake_sum = WakeField.add_wakefields(wakefields[0], beta[:, 0],
                                                wakefields[1], beta[:, 1])
            for i in range(len(wakefields) - 2):
                wake_sum = WakeField.add_wakefields(wake_sum, [1, 1],
                                                    wakefields[i + 2],
                                                    beta[:, i + 2])
            return wake_sum
        else:
            raise ValueError("Error in input.")

"""
General calculations about instability thresholds.
"""

import math

from scipy.constants import c, e, epsilon_0, m_e, mu_0, pi

def mbi_threshold(ring, sigma, R, b):
    """
    Compute the microbunching instability (MBI) threshold for a bunched beam
    considering the steady-state parallel plate model [1][2].
    
    sigma : float
        RMS bunch length in [s]
    R : float
        dipole bending radius in [m]
    b : float
        vertical distance between the conducting parallel plates in [m]
        
    Returns
    -------
    I : float
        MBI current threshold in [A]
        
    [1] : Y. Cai, "Theory of microwave instability and coherent synchrotron 
    radiation in electron storage rings", SLAC-PUB-14561
    [2] : D. Zhou, "Coherent synchrotron radiation and microwave instability in 
    electron storage rings", PhD thesis, p112
    """

    sigma = sigma * c
    Ia = 4 * pi * epsilon_0 * m_e * c**3 / e  # Alfven current
    chi = sigma * (R / b**3)**(1 / 2)  # Shielding paramter
    xi = 0.5 + 0.34*chi
    N = (ring.L * Ia * ring.ac * ring.gamma * ring.sigma_delta**2 * xi *
         sigma**(1 / 3) / (c * e * R**(1 / 3)))
    I = N * e / ring.T0

    return I

def cbi_threshold(ring, I, Vrf, f, beta, Ncav=1):
    """
    Compute the longitudinal and transverse coupled bunch instability 
    thresolds driven by HOMs [1].
    
    Approximate formula, does not take into account variation with Q.
    For better estimate use lcbi_growth_rate.

    I : float
        Total beam current in [A].
    Vrf : float
        Total RF voltage in [V].
    f : float
        Frequency of the HOM in [Hz].
    beta : array-like of shape (2,)
        Horizontal and vertical beta function at the HOM position in [m].
    Ncav : int, optional
        Number of RF cavity.

    Returns
    -------
    Zlong : float
        Maximum longitudinal impedance of the HOM in [ohm].
    Zxdip : float
        Maximum horizontal dipolar impedance of the HOM in [ohm/m].
    Zydip : float
        Maximum vertical dipolar impedance of the HOM in [ohm/m].
        
    [1] : Ruprecht, Martin, et al. "Calculation of Transverse Coupled Bunch 
    Instabilities in Electron Storage Rings Driven By Quadrupole Higher Order 
    Modes." 7th Int. Particle Accelerator Conf.(IPAC'16), Busan, Korea. 
    """

    fs = ring.synchrotron_tune(Vrf) * ring.f0
    Zlong = fs / (f * ring.ac * ring.tau[2]) * (2 * ring.E0) / (ring.f0 * I *
                                                                Ncav)
    Zxdip = 1 / (ring.tau[0] * beta[0]) * (2 * ring.E0) / (ring.f0 * I * Ncav)
    Zydip = 1 / (ring.tau[1] * beta[1]) * (2 * ring.E0) / (ring.f0 * I * Ncav)

    return (Zlong, Zxdip, Zydip)

def lcbi_growth_rate_mode(ring,
                          I,
                          M,
                          mu,
                          synchrotron_tune=None,
                          Vrf=None,
                          fr=None,
                          RL=None,
                          QL=None,
                          Z=None):
    """
    Compute the longitudinal coupled bunch instability growth rate driven by
    an impedance for a given coupled bunch mode mu [1].
    
    Use either a list of resonators (fr, RL, QL) or an Impedance object (Z).

    I : float
    M : int
        Nomber of bunches in the beam.
    mu : int
        Coupled bunch mode number (= 0, ..., M-1).
    synchrotron_tune : float, optional
        Synchrotron tune.
    Vrf : float, optinal
        Total RF voltage in [V] used to compute synchrotron tune if not given.
    fr : float or list, optional
        Frequency of the resonator in [Hz].
    RL : float or list, optional
        Loaded shunt impedance of the resonator in [Ohm].
    QL : float or list, optional
        Loaded quality factor of the resonator.
    Z : Impedance, optional
        Longitunial impedance to consider.

    Returns
    -------
    float
        Coupled bunch instability growth rate for the mode mu in [s-1].
        
    [1] : Eq. 51 p139 of Akai, Kazunori. "RF System for Electron Storage 
    Rings." Physics And Engineering Of High Performance Electron Storage Rings 
    And Application Of Superconducting Technology. 2002. 118-149.

    """

    if synchrotron_tune is None and Vrf is None:
        raise ValueError("Either synchrotron_tune or Vrf is needed.")
    if synchrotron_tune is None:
        nu_s = ring.synchrotron_tune(Vrf)
    else:
        nu_s = synchrotron_tune
    factor = ring.eta() * I / (4 * np.pi * ring.E0 * nu_s)

    if isinstance(fr, (float, int)):
        fr = np.array([fr])
    elif isinstance(fr, list):
        fr = np.array(fr)
    if isinstance(RL, (float, int)):
        RL = np.array([RL])
    elif isinstance(RL, list):
        RL = np.array(RL)
    if isinstance(QL, (float, int)):
        QL = np.array([QL])
    elif isinstance(QL, list):
        QL = np.array(QL)

    if Z is None:
        omega_r = 2 * np.pi * fr
        n_max = int(10 * omega_r.max() / (ring.omega0 * M))

        def Zr(omega):
            Z = 0
            for i in range(len(fr)):
                Z += np.real(RL[i] /
                             (1 + 1j * QL[i] *
                              (omega_r[i] / omega - omega / omega_r[i])))
            return Z
    else:
        fmax = Z.data.index.max()
        n_max = int(2 * np.pi * fmax / (ring.omega0 * M))

            return np.real(Z(omega / (2 * np.pi)))

    n0 = np.arange(n_max)
    n1 = np.arange(1, n_max)
    omega_p = ring.omega0 * (n0*M + mu + nu_s)
    omega_m = ring.omega0 * (n1*M - mu - nu_s)

    sum_val = np.sum(omega_p * Zr(omega_p)) - np.sum(omega_m * Zr(omega_m))

    return factor * sum_val

def lcbi_growth_rate(ring,
                     I,
                     M,
                     synchrotron_tune=None,
                     Vrf=None,
                     fr=None,
                     RL=None,
                     QL=None,
                     Z=None):
    """
    Compute the maximum growth rate for longitudinal coupled bunch instability 
    driven an impedance [1].
    
    I : float
    M : int
        Loaded shunt impedance of the HOM in [Ohm].
        Loaded quality factor of the HOM.

    Returns
    -------
    growth_rate : float
        Maximum coupled bunch instability growth rate in [s-1].
    mu : int
        Coupled bunch mode number corresponding to the maximum coupled bunch 
        instability growth rate.
    growth_rates : array
        Coupled bunch instability growth rates for the different mode numbers 
        in [s-1].
        
    """
    growth_rates = np.zeros(M)
    for i in range(M):
        growth_rates[i] = lcbi_growth_rate_mode(
            ring,
            I,
            M,
            i,
            synchrotron_tune=synchrotron_tune,
            Vrf=Vrf,
            fr=fr,
            RL=RL,
            QL=QL,
            Z=Z)

    growth_rate = np.max(growth_rates)
    mu = np.argmax(growth_rates)

    return growth_rate, mu, growth_rates

def lcbi_stability_diagram(ring,
                           I,
                           M,
                           modes,
                           cavity_list,
                           detune_range,
                           synchrotron_tune=None,
                           Vrf=None):
    """
    Plot longitudinal coupled bunch instability stability diagram for a 
    arbitrary list of CavityResonator objects around a detuning range.
    
    Last object in the cavity_list is assumed to be the one with the variable 
    detuning.

        Ring parameters.
    I : float
    M : int
    modes : list
        Coupled bunch mode numbers to consider.
    cavity_list : list
        List of CavityResonator objects to consider, which can be:
            - active cavities
            - passive cavities
            - HOMs
            - mode dampers
    detune_range : array
        Detuning range (list of points) of the last CavityResonator object.

    Returns
    -------
    fig : Figure
        Show the shunt impedance threshold for different coupled bunch modes.

    """

    Rth = np.zeros_like(detune_range)
    fig, ax = plt.subplots()

    for mu in modes:
        fixed_gr = 0
        for cav in cavity_list[:-1]:
            fixed_gr += lcbi_growth_rate_mode(
                ring,
                I=I,
                mu=mu,
                fr=cav.fr,
                RL=cav.RL,
                QL=cav.QL,
                M=M,
                synchrotron_tune=synchrotron_tune,
                Vrf=Vrf)

        cav = cavity_list[-1]
        for i, det in enumerate(detune_range):
            gr = lcbi_growth_rate_mode(ring,
                                       I=I,
                                       mu=mu,
                                       fr=cav.m * ring.f1 + det,
                                       RL=cav.RL,
                                       QL=cav.QL,
                                       M=M,
                                       synchrotron_tune=synchrotron_tune,
                                       Vrf=Vrf)
            Rth[i] = (1 / ring.tau[2] - fixed_gr) * cav.Rs / gr

        ax.plot(detune_range * 1e-3,
                Rth * 1e-6,
                label="$\mu$ = " + str(int(mu)))

    plt.xlabel(r"$\Delta f$ [kHz]")
    plt.ylabel(r"$R_{s,max}$ $[M\Omega]$")
    plt.yscale("log")
    plt.legend()

    return fig

def rwmbi_growth_rate(ring, current, beff, rho_material, plane='x'):
    """
    Compute the growth rate of the transverse coupled-bunch instability induced
    by resistive wall impedance [1].

    current : float
    beff : float
        Effective radius of the vacuum chamber in [m].
    rho_material : float
        Resistivity of the chamber's wall material in [Ohm.m].
    plane : str, optional
        The plane in which the instability will be computed. Use 'x' for the 
        horizontal plane, and 'y' for the vertical.

    Reference
    ---------
    [1] Eq. (31) in R. Nagaoka and K. L. F. Bane, "Collective effects in a
    diffraction-limited storage ring", J. Synchrotron Rad. Vol 21, 2014. pp.937-960 

    """
    plane_dict = {'x': 0, 'y': 1}
    index = plane_dict[plane]
    beta0 = ring.optics.local_beta[index]
    omega0 = ring.omega0
    E0 = ring.E0
    R = ring.L / (2 * np.pi)
    frac_tune, int_tune = math.modf(ring.tune[index])
    Z0 = mu_0 * c

    gr = (beta0*omega0*current*R) / (4 * np.pi * E0 * beff**3) * (
        (2*c*Z0*rho_material) / ((1-frac_tune) * omega0))**0.5

    return gr

def rwmbi_threshold(ring, beff, rho_material, plane='x'):
    """
    Compute the threshold current of the transverse coupled-bunch instability 
    induced by resistive wall impedance [1].

    Reference
    ---------
    [1] Eq. (32) in R. Nagaoka and K. L. F. Bane, "Collective effects in a

    """
    tau_rad = ring.tau[index]
    Z0 = mu_0 * c

    Ith = (4 * np.pi * E0 * beff**3) / (c*beta0*tau_rad) * ((
        (1-frac_tune) * omega0) / (2*c*Z0*rho_material))**0.5

    return Ith

"""
Various calculations about ion trapping and instabilities in electron storage 
rings.
"""

from scipy.constants import (
    Boltzmann,
    c,
    e,
    epsilon_0,
    hbar,
    m_e,
    m_p,
    physical_constants,
    pi,
)

rp = 1 / (4*pi*epsilon_0) * e**2 / (m_p * c**2)
re = physical_constants["classical electron radius"][0]

def ion_cross_section(ring, ion):
    """
    Compute the collisional ionization cross section.
    
    Compute the inelastic collision cross section between a molecule or an atom
    by a relativistic electron using the relativistic Bethe asymptotic formula 
    [1].
    
    Values of M02 and C02 from [2].

    ion : str
        Ion type.

    Returns
    -------
        Cross section in [m**2].
        
    [1] : M. Inokuti, "Inelastic collisions of fast charged particles with 
    atoms and molecules-the bethe theory revisited", Reviews of modern physics 
    43 (1971).
    [2] : P. F. Tavares, "Bremsstrahlung detection of ions trapped in the EPA 
    electron beam", Part. Accel. 43 (1993).

    """
    if ion == "CO":
        M02 = 3.7
        C0 = 35.1
    elif ion == "H2":
        M02 = 0.7
        C0 = 8.1
    else:

    sigma = 4 * pi * (hbar / m_e / c)**2 * (
        M02 * (1 / ring.beta**2 * np.log(ring.beta**2 /
                                         (1 - ring.beta**2)) - 1) +
        C0 / ring.beta**2)

    return sigma

def ion_frequency(N,
                  Lsep,
                  sigmax,
                  sigmay,
                  ion="CO",
                  dim="y",
                  express="coupling"):
    """
    Compute the ion oscillation frequnecy.

    N : float
        Number of electrons per bunch.
    Lsep : float
        Bunch spacing in [m].
    sigmax : float or array
        Horizontal beam size in [m].
    sigmay : float or array
        Vertical beam size in [m].
    ion : str, optional
        Ion type. The default is "CO".
    dim : "y" o "x", optional
        Dimension to consider. The default is "y".
    express : str, optional
        Expression to use to compute the ion oscillation frequency. 
        The default is "coupling" corresponding to Gaussian electron and ion 
        distributions with coupling [1].
        Also possible is "no_coupling" corresponding to Gaussian electron and 
        ion distributions without coupling [2].

    Returns
    -------
    f : float or array
        Ion oscillation frequencies in [Hz].
        
    [1] : T. O. Raubenheimer and F. Zimmermann, "Fast beam-ion instability. I. 
    linear theory and simulations", Physical Review E 52 (1995).
    [2] : G. V. Stupakov, T. O. Raubenheimer, and F. Zimmermann, "Fast beam-ion 
    instability. II. effect of ion decoherence", Physical Review E 52 (1995).

    """

        A = 28
        A = 2
    elif ion == "CH4":
        A = 18
    elif ion == "H2O":
        A = 16
    elif ion == "CO2":
        A = 44

    if dim == "y":
        pass
    elif dim == "x":
        sigmay, sigmax = sigmax, sigmay
    else:
        raise ValueError

    if express == "coupling":
        k = 3 / 2
    elif express == "no_coupling":
        k = 1

    f = c * np.sqrt(2 * rp * N / (A * k * Lsep * sigmay *
                                  (sigmax+sigmay))) / (2*pi)

    return f

def fast_beam_ion(ring,
                  Nb,
                  nb,
                  Lsep,
                  sigmax,
                  sigmay,
                  P,
                  T,
                  beta,
                  model="linear",
                  delta_omega=0,
                  ion="CO",
                  dim="y"):
    """
    Compute fast beam ion instability rise time [1].
    
    Warning ! 
    If model="linear", the rise time is an assymptotic grow time 
    (i.e. y ~ exp(sqrt(t/tau))) [1].
    If model="decoherence", the rise time is an e-folding time 
    (i.e. y ~ exp(t/tau)) [2].
    If model="non-linear", the rise time is a linear growth time
    (i.e. y ~ t/tau) [3].
    
    The linear model assumes that [1]:
        x,y << sigmax,sigmay
    
    The decoherence model assumes that [2]:
        Lsep << c / (2 * pi * ion_frequency) 
        Lsep << c / (2 * pi * betatron_frequency)
        
    The non-linear model assumes that [3]:
        x,y >> sigmax,sigmay

    Nb : float
        Number of electron per bunch.
    nb : float
        Number of bunches.
    sigmax : float
    sigmay : float
    P : float
        Partial pressure of the molecular ion in [Pa].
    T : float
        Tempertature in [K].
    beta : float
        Average betatron function around the ring in [m].
    model : str, optional
        If "linear", use [1].
        If "decoherence", use [2].
        If "non-linear", use [3].
    delta_omega : float, optional
        RMS variation of the ion oscillation angular frequnecy around the ring
        in [Hz].

    Returns
    -------
    tau : float
        Instability rise time in [s].
        
    [3] : Chao, A. W., & Mess, K. H. (Eds.). (2013). Handbook of accelerator 
    physics and engineering. World scientific. 3rd Printing. p417.
    
    """
        pass
    else:

        A = 28
        A = 2

    sigma_i = ion_cross_section(ring, ion)

    d_gas = P / (Boltzmann*T)

    num = 4 * d_gas * sigma_i * beta * Nb**(3 / 2) * nb**2 * re * rp**(
        1 / 2) * Lsep**(1 / 2) * c
    den = 3 * np.sqrt(3) * ring.gamma * sigmay**(3 / 2) * (sigmay + sigmax)**(
        3 / 2) * A**(1 / 2)

    tau = den / num

    if model == "decoherence":
        tau = tau * 2 * np.sqrt(2) * nb * Lsep * delta_omega / c
    elif model == "non-linear":
        fi = ion_frequency(Nb, Lsep, sigmax, sigmay, ion, dim)
        tau = tau * 2 * pi * fi * ring.T1 * nb**(3 / 2)
    elif model == "linear":
        pass
    else:
        raise ValueError("model unknown")

    return tau

def plot_critical_mass(ring, bunch_charge, bunch_spacing, n_points=1e4):
    """
    Plot ion critical mass, using Eq. (7.70) p147 of [1]

    bunch_charge : float
        Bunch charge in [C].
    bunch_spacing : float
        Time in between two adjacent bunches in [s].
    n_points : float or int, optional
        Number of point used in the plot. The default is 1e4.

    Returns
    -------
    fig : figure
        
    [1] : Gamelin, A. (2018). Collective effects in a transient microbunching 
    regime and ion cloud mitigation in ThomX (Doctoral dissertation, 
    Université Paris-Saclay).

    """

    n_points = int(n_points)
    s = np.linspace(0, ring.L, n_points)
    sigma = ring.sigma(s)
    N = np.abs(bunch_charge / e)

    Ay = N * rp * bunch_spacing * c / (2 * sigma[2, :] *
                                       (sigma[2, :] + sigma[0, :]))
    Ax = N * rp * bunch_spacing * c / (2 * sigma[0, :] *

    fig = plt.figure()
    ax = plt.gca()
    ax.plot(s, Ax, label=r"$A_x^c$")
    ax.plot(s, Ay, label=r"$A_y^c$")
    ax.set_yscale("log")
    ax.plot(s, np.ones_like(s) * 2, label=r"$H_2^+$")
    ax.plot(s, np.ones_like(s) * 16, label=r"$H_2O^+$")
    ax.plot(s, np.ones_like(s) * 18, label=r"$CH_4^+$")
    ax.plot(s, np.ones_like(s) * 28, label=r"$CO^+$")
    ax.plot(s, np.ones_like(s) * 44, label=r"$CO_2^+$")
    ax.legend()
    ax.set_ylabel("Critical mass")
    ax.set_xlabel("Longitudinal position [m]")

    return fig

"""
This module defines aperture elements for tracking.
"""

from mbtrack2.tracking.element import Element

class CircularAperture(Element):
    """
    Circular aperture element. The particles which are outside of the circle 
    are 'lost' and not used in the tracking any more.
    
    radius : float
        radius of the circle in [m]
    """

    def __init__(self, radius):
        self.radius_squared = radius**2

    @Element.parallel
    def track(self, bunch):
        """
        Tracking method for the element.
        No bunch to bunch interaction, so written for Bunch objects and
        @Element.parallel is used to handle Beam objects.
        
        bunch : Bunch or Beam object
        """
        alive = bunch.particles["x"]**2 + bunch.particles[
            "y"]**2 < self.radius_squared
        bunch.alive[~alive] = False

class ElipticalAperture(Element):
    """
    Eliptical aperture element. The particles which are outside of the elipse 
    
    X_radius : float
        horizontal radius of the elipse in [m]
    Y_radius : float
        vertical radius of the elipse in [m]
    """

    def __init__(self, X_radius, Y_radius):
        self.X_radius = X_radius
        self.X_radius_squared = X_radius**2
        self.Y_radius = Y_radius
        self.Y_radius_squared = Y_radius**2

        """
        
        """
        alive = (bunch.particles["x"]**2 / self.X_radius_squared +
                 bunch.particles["y"]**2 / self.Y_radius_squared < 1)

class RectangularAperture(Element):
    """
    Rectangular aperture element. The particles which are outside of the 
    rectangle are 'lost' and not used in the tracking any more.
    
    X_right : float
        right horizontal aperture of the rectangle in [m]
    Y_top : float
        top vertical aperture of the rectangle in [m]
    X_left : float, optional
        left horizontal aperture of the rectangle in [m]
    Y_bottom : float, optional
        bottom vertical aperture of the rectangle in [m]
    """

    def __init__(self, X_right, Y_top, X_left=None, Y_bottom=None):
        self.X_right = X_right
        self.X_left = X_left
        self.Y_top = Y_top
        self.Y_bottom = Y_bottom

        """
        
        """

        if (self.X_left is None):
            alive_X = np.abs(bunch.particles["x"]) < self.X_right
        else:
            alive_X = ((bunch.particles["x"] < self.X_right) &
                       (bunch.particles["x"] > self.X_left))

        if (self.Y_bottom is None):
            alive_Y = np.abs(bunch.particles["y"]) < self.Y_top
        else:
            alive_Y = ((bunch.particles["y"] < self.Y_top) &
                       (bunch.particles["y"] > self.Y_bottom))

        alive = alive_X & alive_Y

class LongitudinalAperture(Element):
    """
    Longitudinal aperture element. The particles which are outside of the 
    longitudinal bounds are 'lost' and not used in the tracking any more.
    
    tau_up : float
        Upper longitudinal bound in [s].
    tau_low : float, optional
        Lower longitudinal bound in [s].
    """

    def __init__(self, tau_up, tau_low=None):
        self.tau_up = tau_up
        if tau_low is None:
            self.tau_low = tau_up * -1
        else:
            self.tau_low = tau_low

        """
        
        """

        alive = ((bunch.particles["tau"] < self.tau_up) &
                 (bunch.particles["tau"] > self.tau_low))

"""
This module defines the most basic elements for tracking, including Element,
an abstract base class which is to be used as mother class to every elements
included in the tracking.
"""

from abc import ABCMeta, abstractmethod
from functools import wraps

from mbtrack2.tracking.particles import Beam

class Element(metaclass=ABCMeta):
    """
    Abstract Element class used for subclass inheritance to define all kinds 
    of objects which intervene in the tracking.
    """

    @abstractmethod
    def track(self, beam):
        """
        Track a beam object through this Element.
        This method needs to be overloaded in each Element subclass.
        
        beam : Beam object
        """

    def parallel(track):
        """
        Defines the decorator @parallel which handle the embarrassingly 
        parallel case which happens when there is no bunch to bunch 
        interaction in the tracking routine.
        
        Adding @Element.parallel allows to write the track method of the 
        Element subclass for a Bunch object instead of a Beam object.
        
        track : function, method of an Element subclass
            track method of an Element subclass which takes a Bunch object as
            input
            
        Returns
        -------
        track_wrapper: function, method of an Element subclass
            track method of an Element subclass which takes a Beam object or a
            Bunch object as input
        """

        @wraps(track)
        def track_wrapper(*args, **kwargs):
            if isinstance(args[1], Beam):
                self = args[0]
                beam = args[1]
                if (beam.mpi_switch == True):
                    track(self, beam[beam.mpi.bunch_num], *args[2:], **kwargs)
                else:
                    for bunch in beam.not_empty:
                        track(self, bunch, *args[2:], **kwargs)
            else:
                bunch = args[1]
                track(self, bunch, *args[2:], **kwargs)

        return track_wrapper

class LongitudinalMap(Element):
    """
    Longitudinal map for a single turn in the synchrotron.
    
    """

        """
        
        """
        bunch["delta"] -= self.ring.U0 / self.ring.E0
        bunch["tau"] += self.ring.eta(
            bunch["delta"]) * self.ring.T0 * bunch["delta"]

class SynchrotronRadiation(Element):
    """
    Element to handle synchrotron radiation, radiation damping and quantum 
    excitation, for a single turn in the synchrotron.
    
    switch : bool array of shape (3,), optional
        allow to choose on which plane the synchrotron radiation is active
    """

    def __init__(self, ring, switch=np.ones((3, ), dtype=bool)):
        self.switch = switch

        """
        
        """
        if (self.switch[0] == True):
            rand = np.random.normal(size=len(bunch))
            bunch["delta"] = (
                (1 - 2 * self.ring.T0 / self.ring.tau[2]) * bunch["delta"] +
                2 * self.ring.sigma_delta *
                (self.ring.T0 / self.ring.tau[2])**0.5 * rand)

        if (self.switch[1] == True):
            bunch["xp"] = (
                (1 - 2 * self.ring.T0 / self.ring.tau[0]) * bunch["xp"] +
                2 * self.ring.sigma()[1] *
                (self.ring.T0 / self.ring.tau[0])**0.5 * rand)

        if (self.switch[2] == True):
            bunch["yp"] = (
                (1 - 2 * self.ring.T0 / self.ring.tau[1]) * bunch["yp"] +
                2 * self.ring.sigma()[3] *
                (self.ring.T0 / self.ring.tau[1])**0.5 * rand)

class TransverseMap(Element):
    """
    Transverse map for a single turn in the synchrotron.
    
    """

        self.alpha = self.ring.optics.local_alpha
        self.beta = self.ring.optics.local_beta
        self.gamma = self.ring.optics.local_gamma
        self.dispersion = self.ring.optics.local_dispersion
        if self.ring.adts is not None:
            self.adts_poly = [
                np.poly1d(self.ring.adts[0]),
                np.poly1d(self.ring.adts[1]),
                np.poly1d(self.ring.adts[2]),
                np.poly1d(self.ring.adts[3])
            ]

        """
        
        """

        # Compute phase advance which depends on energy via chromaticity and ADTS
        if self.ring.adts is None:
            phase_advance_x = 2 * np.pi * (self.ring.tune[0] +
                                           self.ring.chro[0] * bunch["delta"])
            phase_advance_y = 2 * np.pi * (self.ring.tune[1] +
                                           self.ring.chro[1] * bunch["delta"])
        else:
            Jx = (self.ring.optics.local_gamma[0] * bunch['x']**2) + \
                  (2*self.ring.optics.local_alpha[0] * bunch['x']*bunch['xp']) + \
                  (self.ring.optics.local_beta[0] * bunch['xp']**2)
            Jy = (self.ring.optics.local_gamma[1] * bunch['y']**2) + \
                  (2*self.ring.optics.local_alpha[1] * bunch['y']*bunch['yp']) + \
                  (self.ring.optics.local_beta[1] * bunch['yp']**2)
            phase_advance_x = 2 * np.pi * (
                self.ring.tune[0] + self.ring.chro[0] * bunch["delta"] +
                self.adts_poly[0](Jx) + self.adts_poly[2](Jy))
            phase_advance_y = 2 * np.pi * (
                self.ring.tune[1] + self.ring.chro[1] * bunch["delta"] +
                self.adts_poly[1](Jx) + self.adts_poly[3](Jy))

        # 6x6 matrix corresponding to (x, xp, delta, y, yp, delta)
        matrix = np.zeros((6, 6, len(bunch)), dtype=np.float64)

        # Horizontal
        c_x = np.cos(phase_advance_x)
        s_x = np.sin(phase_advance_x)

        matrix[0, 0, :] = c_x + self.alpha[0] * s_x
        matrix[0, 1, :] = self.beta[0] * s_x
        matrix[0, 2, :] = self.dispersion[0]
        matrix[1, 0, :] = -1 * self.gamma[0] * s_x
        matrix[1, 1, :] = c_x - self.alpha[0] * s_x
        matrix[1, 2, :] = self.dispersion[1]
        matrix[2, 2, :] = 1

        # Vertical
        c_y = np.cos(phase_advance_y)
        s_y = np.sin(phase_advance_y)

        matrix[3, 3, :] = c_y + self.alpha[1] * s_y
        matrix[3, 4, :] = self.beta[1] * s_y
        matrix[3, 5, :] = self.dispersion[2]
        matrix[4, 3, :] = -1 * self.gamma[1] * s_y
        matrix[4, 4, :] = c_y - self.alpha[1] * s_y
        matrix[4, 5, :] = self.dispersion[3]
        matrix[5, 5, :] = 1

        x = matrix[0, 0, :] * bunch["x"] + matrix[
            0, 1, :] * bunch["xp"] + matrix[0, 2, :] * bunch["delta"]
        xp = matrix[1, 0, :] * bunch["x"] + matrix[
            1, 1, :] * bunch["xp"] + matrix[1, 2, :] * bunch["delta"]
        y = matrix[3, 3, :] * bunch["y"] + matrix[
            3, 4, :] * bunch["yp"] + matrix[3, 5, :] * bunch["delta"]
        yp = matrix[4, 3, :] * bunch["y"] + matrix[
            4, 4, :] * bunch["yp"] + matrix[4, 5, :] * bunch["delta"]

        bunch["x"] = x
        bunch["xp"] = xp
        bunch["y"] = y
        bunch["yp"] = yp

class SkewQuadrupole:
    """
    Thin skew quadrupole element used to introduce betatron coupling (the 
    length of the quadrupole is neglected).
    
    strength : float
        Integrated strength of the skew quadrupole [m].
        
    """

    def __init__(self, strength):
        self.strength = strength

        """
        
        """

        bunch['xp'] = bunch['xp'] - self.strength * bunch['y']
        bunch['yp'] = bunch['yp'] - self.strength * bunch['x']

class TransverseMapSector(Element):
    """
    Transverse map for a sector of the synchrotron, from an initial 
    position s0 to a final position s1.

    alpha0 : array of shape (2,)
        Alpha Twiss function at the initial location of the sector.
    beta0 : array of shape (2,)
        Beta Twiss function at the initial location of the sector.
    dispersion0 : array of shape (4,)
        Dispersion function at the initial location of the sector.
    alpha1 : array of shape (2,)
        Alpha Twiss function at the final location of the sector.
    beta1 : array of shape (2,)
        Beta Twiss function at the final location of the sector.
    dispersion1 : array of shape (4,)
        Dispersion function at the final location of the sector.
    phase_diff : array of shape (2,)
        Phase difference between the initial and final location of the 
        sector.
    chro_diff : array of shape (2,)
        Chromaticity difference between the initial and final location of 
        the sector.
    adts : array of shape (4,), optional
        Amplitude-dependent tune shift of the sector, see Synchrotron class 
        for details. The default is None.

    """

                 ring,
                 alpha0,
                 beta0,
                 dispersion0,
                 alpha1,
                 beta1,
                 dispersion1,
                 phase_diff,
                 chro_diff,
                 adts=None):
        self.alpha0 = alpha0
        self.beta0 = beta0
        self.gamma0 = (1 + self.alpha0**2) / self.beta0
        self.dispersion0 = dispersion0
        self.alpha1 = alpha1
        self.beta1 = beta1
        self.gamma1 = (1 + self.alpha1**2) / self.beta1
        self.dispersion1 = dispersion1
        self.tune_diff = phase_diff / (2 * np.pi)
        self.chro_diff = chro_diff
        if adts is not None:
                np.poly1d(adts[0]),
                np.poly1d(adts[1]),
                np.poly1d(adts[2]),
                np.poly1d(adts[3])
            ]
        else:
            self.adts_poly = None

        """
        
        """

        if self.adts_poly is None:
            phase_advance_x = 2 * np.pi * (self.tune_diff[0] +
                                           self.chro_diff[0] * bunch["delta"])
            phase_advance_y = 2 * np.pi * (self.tune_diff[1] +
                                           self.chro_diff[1] * bunch["delta"])
        else:
            Jx = (self.gamma0[0] * bunch['x']**2) + \
                  (2*self.alpha0[0] * bunch['x']*self['xp']) + \
                  (self.beta0[0] * bunch['xp']**2)
            Jy = (self.gamma0[1] * bunch['y']**2) + \
                  (2*self.alpha0[1] * bunch['y']*bunch['yp']) + \
                  (self.beta0[1] * bunch['yp']**2)
                self.tune_diff[0] + self.chro_diff[0] * bunch["delta"] +
                self.tune_diff[1] + self.chro_diff[1] * bunch["delta"] +

        matrix = np.zeros((6, 6, len(bunch)))

        matrix[0, 0, :] = np.sqrt(self.beta1[0] / self.beta0[0]) * (
            np.cos(phase_advance_x) + self.alpha0[0] * np.sin(phase_advance_x))
        matrix[0, 1, :] = np.sqrt(
            self.beta0[0] * self.beta1[0]) * np.sin(phase_advance_x)
        matrix[0, 2, :] = self.dispersion1[0] - matrix[
            0, 0, :] * self.dispersion0[0] - matrix[0,
                                                    1, :] * self.dispersion0[1]
        matrix[1, 0, :] = (
            (self.alpha0[0] - self.alpha1[0]) * np.cos(phase_advance_x) -
            (1 + self.alpha0[0] * self.alpha1[0]) *
            np.sin(phase_advance_x)) / np.sqrt(self.beta0[0] * self.beta1[0])
        matrix[1, 1, :] = np.sqrt(self.beta0[0] / self.beta1[0]) * (
            np.cos(phase_advance_x) - self.alpha1[0] * np.sin(phase_advance_x))
        matrix[1, 2, :] = self.dispersion1[1] - matrix[
            1, 0, :] * self.dispersion0[0] - matrix[1,

        # Vertical
        matrix[3, 3, :] = np.sqrt(self.beta1[1] / self.beta0[1]) * (
            np.cos(phase_advance_y) + self.alpha0[1] * np.sin(phase_advance_y))
        matrix[3, 4, :] = np.sqrt(
            self.beta0[1] * self.beta1[1]) * np.sin(phase_advance_y)
        matrix[3, 5, :] = self.dispersion1[2] - matrix[
            3, 3, :] * self.dispersion0[2] - matrix[3,
                                                    4, :] * self.dispersion0[3]
        matrix[4, 3, :] = (
            (self.alpha0[1] - self.alpha1[1]) * np.cos(phase_advance_y) -
            (1 + self.alpha0[1] * self.alpha1[1]) *
            np.sin(phase_advance_y)) / np.sqrt(self.beta0[1] * self.beta1[1])
        matrix[4, 4, :] = np.sqrt(self.beta0[1] / self.beta1[1]) * (
            np.cos(phase_advance_y) - self.alpha1[1] * np.sin(phase_advance_y))
        matrix[4, 5, :] = self.dispersion1[3] - matrix[
            4, 3, :] * self.dispersion0[2] - matrix[4,

def transverse_map_sector_generator(ring, positions):
    """
    Convenience function which generate a list of TransverseMapSector elements
    from a ring with an AT lattice.
    
    Tracking through all the sectors is equivalent to a full turn (and thus to 
    the TransverseMap object).

        Ring parameters, must .
    positions : array
        List of longitudinal positions in [m] to use as starting and end points
        of the TransverseMapSector elements.
        The array should contain the initial position (s=0) but not the end 
        position (s=ring.L), so like position = np.array([0, pos1, pos2, ...]).

    Returns
    -------
    sectors : list
        List of TransverseMapSector elements.

    """
    import at

    def _compute_chro(ring, pos, dp=1e-4):
        lat = deepcopy(ring.optics.lattice)
        lat.append(at.Marker("END"))
        N = len(lat)
        refpts = np.arange(N)
        *elem_neg_dp, = at.linopt2(lat, refpts=refpts, dp=-dp)
        *elem_pos_dp, = at.linopt2(lat, refpts=refpts, dp=dp)

        s = elem_neg_dp[2]["s_pos"]
        mux0 = elem_neg_dp[2]['mu'][:, 0]
        mux1 = elem_pos_dp[2]['mu'][:, 0]
        muy0 = elem_neg_dp[2]['mu'][:, 1]
        muy1 = elem_pos_dp[2]['mu'][:, 1]

        Chrox = (mux1-mux0) / (2*dp) / 2 / np.pi
        Chroy = (muy1-muy0) / (2*dp) / 2 / np.pi
        chrox = np.interp(pos, s, Chrox)
        chroy = np.interp(pos, s, Chroy)

        return np.array([chrox, chroy])

    if ring.optics.use_local_values:
        raise ValueError(
            "The Synchrotron object must be loaded from an AT lattice")

    N_sec = len(positions)
    sectors = []
    for i in range(N_sec):
        alpha0 = ring.optics.alpha(positions[i])
        beta0 = ring.optics.beta(positions[i])
        dispersion0 = ring.optics.dispersion(positions[i])
        mu0 = ring.optics.mu(positions[i])
        chro0 = _compute_chro(ring, positions[i])
        if i != (N_sec - 1):
            alpha1 = ring.optics.alpha(positions[i + 1])
            beta1 = ring.optics.beta(positions[i + 1])
            dispersion1 = ring.optics.dispersion(positions[i + 1])
            mu1 = ring.optics.mu(positions[i + 1])
            chro1 = _compute_chro(ring, positions[i + 1])
        else:
            alpha1 = ring.optics.alpha(positions[0])
            beta1 = ring.optics.beta(positions[0])
            dispersion1 = ring.optics.dispersion(positions[0])
            mu1 = ring.optics.mu(ring.L)
            chro1 = _compute_chro(ring, ring.L)
        phase_diff = mu1 - mu0
        chro_diff = chro1 - chro0
        sectors.append(
            TransverseMapSector(ring, alpha0, beta0, dispersion0, alpha1,
                                beta1, dispersion1, phase_diff, chro_diff))
    return sectors

"""
This module defines both exponential and FIR based bunch by bunch damper 
feedback for tracking.
"""

from mbtrack2.tracking.particles import Beam, Bunch

class ExponentialDamper(Element):
    """ 
    Simple bunch by bunch damper feedback system based on exponential damping.
    
        Synchrotron to use.
    plane : {"x","y","s"}
        Allow to choose on which plane the damper is active.
    damping_time : float
        Damping time in [s].
    phase_diff : array of shape (3,)
        Phase difference between the damper and the position monitor in [rad].
        
    """

    def __init__(self, ring, plane, damping_time, phase_diff):
        self.damping_time = damping_time
        self.phase_diff = phase_diff
        self.plane = plane
        if self.plane == "x":
            self.action = "xp"
            self.mean_idx = 1
        elif self.plane == "y":
            self.action = "yp"
            self.mean_idx = 3
        elif self.plane == "s":
            self.action = "delta"
            self.mean_idx = 5
        else:
            raise ValueError(f"plane should be x, y or s, not {self.plane}")

        """
        Tracking method for the feedback system
        No bunch to bunch interaction, so written for Bunch object and
        @Element.parallel is used to handle Beam object.
        
        """
        bunch[self.action] -= (2 * self.ring.T0 / self.damping_time *
                               np.sin(self.phase_diff) *
                               bunch.mean[self.mean_idx])

class FIRDamper(Element):
    """ 
    Bunch by bunch damper feedback system based on FIR filters.
    
    FIR computation is based on [1].
    
    tune : float
        Reference (betatron or synchrotron) tune for which the damper system 
        is set.
    turn_delay : int
        Number of turn delay before the kick is applied.
    tap_number : int
        Number of tap for the FIR filter.
    gain : float
        Gain of the FIR filter.
    phase : float
        Phase of the FIR filter in [degree].
    bpm_error : float, optional
        RMS measurement error applied to the computed mean position.
        Unit is [m] if the plane is "x" or "y" and [s] if the plane is "s".
        The default is None.
    max_kick : float, optional
        Maximum kick strength limitation.
        Unit is [rad] if the plane is "x" or "y" and no unit (delta) if the 
        plane is "s".

    pos : array
        Stored beam postions.
    kick : array
        Stored damper kicks.
    coef : array
        Coefficients of the FIR filter.
         
    Methods
    -------
    get_fir(tap_number, tune, phase, turn_delay, gain)
        Initialize the FIR filter and return an array containing the FIR 
        coefficients.
    plot_fir()
        Plot the gain and the phase of the FIR filter.
    track(beam_or_bunch)
        Tracking method.
      
    [1] T.Nakamura, S.Daté, K. Kobayashi, T. Ohshima. Proceedings of EPAC 
    2004. Transverse bunch by bunch feedback system for the Spring-8 
    storage ring.
    """

                 ring,
                 plane,
                 tune,
                 turn_delay,
                 tap_number,
                 gain,
                 phase,
                 bpm_error=None,
                 max_kick=None):

        self.tune = tune
        self.turn_delay = turn_delay
        self.tap_number = tap_number
        self.gain = gain
        self.phase = phase
        self.bpm_error = bpm_error
        self.max_kick = max_kick

            self.damp_idx = 0
            self.mean_idx = 0
            self.damp_idx = 1
            self.mean_idx = 2
            self.damp_idx = 2
            self.mean_idx = 4

        self.beam_no_mpi = False

        self.pos = np.zeros((self.tap_number, 1))
        self.kick = np.zeros((self.turn_delay + 1, 1))
        self.coef = self.get_fir(self.tap_number, self.tune, self.phase,
                                 self.turn_delay, self.gain)

    def get_fir(self, tap_number, tune, phase, turn_delay, gain):
        """
        Compute the FIR coefficients.
        
        FIR computation is based on [1].
        
        Returns
        -------
        FIR_coef : array
            Array containing the FIR coefficients.
        """
        it = np.zeros((tap_number, ))
        CC = np.zeros((
            5,
            tap_number,
        ))
        zeta = (phase * 2 * np.pi) / 360
        for k in range(tap_number):
            it[k] = (-k - turn_delay)

        phi = 2 * np.pi * tune
        cs = np.cos(phi * it)
        sn = np.sin(phi * it)

        CC[0][:] = 1
        CC[1][:] = cs
        CC[2][:] = sn
        CC[3][:] = it * sn
        CC[4][:] = it * cs

        TCC = np.transpose(CC)
        W = np.linalg.inv(CC.dot(TCC))
        D = W.dot(CC)

        FIR_coef = gain * (D[1][:] * np.cos(zeta) + D[2][:] * np.sin(zeta))
        return FIR_coef

    def plot_fir(self):
        """
        
        Returns
        -------
        fig : Figure
            Plot of the gain and phase.
            
        """
        tune = np.arange(0, 1, 0.0001)

        H_FIR = 0
        for k in range(len(self.coef)):
            H_FIR += self.coef[k] * np.exp(-1j * 2 * np.pi * (k) * tune)
        latency = np.exp(-1j * 2 * np.pi * tune * self.turn_delay)
        H_tot = H_FIR * latency

        gain = np.abs(H_tot)
        phase = np.angle(H_tot, deg=True)

        fig, [ax1, ax2] = plt.subplots(2, 1)
        ax1.plot(tune, gain)
        ax1.set_ylabel("Gain")

        ax2.plot(tune, phase)
        ax2.set_xlabel("Tune")
        ax2.set_ylabel("Phase in degree")

        return fig

    def track(self, beam_or_bunch):
        """

        beam_or_bunch : Beam or Bunch
            Data to track.
            
        """
        if isinstance(beam_or_bunch, Bunch):
            self.track_sb(beam_or_bunch)
        elif isinstance(beam_or_bunch, Beam):
            beam = beam_or_bunch
            if (beam.mpi_switch == True):
                self.track_sb(beam[beam.mpi.bunch_num])
            else:
                if self.beam_no_mpi is False:
                    self.init_beam_no_mpi(beam)
                for i, bunch in enumerate(beam.not_empty):
                    self.track_sb(bunch, i)
        else:
            TypeError("beam_or_bunch must be a Beam or Bunch")

    def init_beam_no_mpi(self, beam):
        """
        Change array sizes if Beam is used without mpi.

        beam : Beam
            Beam to track.

        """
        n_bunch = len(beam)
        self.pos = np.zeros((self.tap_number, n_bunch))
        self.kick = np.zeros((self.turn_delay + 1, n_bunch))
        self.beam_no_mpi = True

    def track_sb(self, bunch, bunch_number=0):
        """
        Core of the tracking method.

        bunch : Bunch
            Bunch to track.
        bunch_number : int, optional
            Number of bunch in beam.not_empty. 
            The default is 0.
            
        """
        self.pos[0, bunch_number] = bunch.mean[self.mean_idx]
        if self.bpm_error is not None:
            self.pos[0, bunch_number] += np.random.normal(0, self.bpm_error)

        kick = 0
        for k in range(self.tap_number):
            kick += self.coef[k] * self.pos[k, bunch_number]

        if self.max_kick is not None:
            if kick > self.max_kick:
                kick = self.max_kick
            elif kick < -1 * self.max_kick:
                kick = -1 * self.max_kick

        self.kick[-1, bunch_number] = kick
        bunch[self.action] += self.kick[0, bunch_number]

        self.pos[:, bunch_number] = np.roll(self.pos[:, bunch_number], 1)
        self.kick[:, bunch_number] = np.roll(self.kick[:, bunch_number], -1)

"""
Module to handle parallel computation
"""

class Mpi:
    """
    Class which handle parallel computation via the mpi4py module [1].
    
    filling_pattern : bool array of shape (h,)
        Filling pattern of the beam, like Beam.filling_pattern
        
    comm : MPI.Intracomm object
        MPI intra-comminicator of the processor group, used to manage 
        communication between processors. 
    rank : int
        Rank of the processor which run the program
    size : int
        Number of processor within the processor group (in fact in the 
        intra-comminicator group)
    table : int array of shape (size, 2)
        Table of correspondance between the rank of the processor and its 
        associated bunch number
    bunch_num : int
        Return the bunch number corresponding to the current processor
    next_bunch : int
        Return the rank of the next tracked bunch
    previous_bunch : int
        Return the rank of the previous tracked bunch
        
    Methods
    -------
    write_table(filling_pattern)
        Write a table with the rank and the corresponding bunch number for each
        bunch of the filling pattern
    rank_to_bunch(rank)
        Return the bunch number corresponding to rank
    bunch_to_rank(bunch_num)
        Return the rank corresponding to the bunch number bunch_num
    share_distributions(beam)
        Compute the bunch profiles and share it between the different bunches.
    share_means(beam)
        Compute the bunch means and share it between the different bunches.
    share_stds(beam)
        Compute the bunch standard deviations and share it between the 
        different bunches.
        
    [1] L. Dalcin, P. Kler, R. Paz, and A. Cosimo, Parallel Distributed 
    Computing using Python, Advances in Water Resources, 34(9):1124-1139, 2011.
    
    """

    def __init__(self, filling_pattern):
        from mpi4py import MPI
        self.MPI = MPI
        self.comm = MPI.COMM_WORLD
        self.rank = self.comm.Get_rank()
        self.size = self.comm.Get_size()
        self.write_table(filling_pattern)

    def write_table(self, filling_pattern):
        """
        
        filling_pattern : bool array of shape (h,)
            Filling pattern of the beam, like Beam.filling_pattern
        """
        if (filling_pattern.sum() != self.size):
            raise ValueError("The number of processors must be equal to the"
                             "number of (non-empty) bunches.")
        table = np.zeros((self.size, 2), dtype=int)
        table[:, 0] = np.arange(0, self.size)
        table[:, 1] = np.where(filling_pattern)[0]
        self.table = table

    def rank_to_bunch(self, rank):
        """
        
        rank : int
            Rank of a processor
            
        Returns
        -------
        bunch_num : int
            Bunch number corresponding to the input rank
        """
        return self.table[rank, 1]

    def bunch_to_rank(self, bunch_num):
        """
        
            Bunch number
            
        Returns
        -------
        rank : int
            Rank of the processor which tracks the input bunch number
        """
        try:
            rank = np.where(self.table[:, 1] == bunch_num)[0][0]
        except IndexError:
            print("The bunch " + str(bunch_num) +
                  " is not tracked on any processor.")
            rank = None
        return rank

    @property
    def bunch_num(self):
        """Return the bunch number corresponding to the current processor"""
        return self.rank_to_bunch(self.rank)

    @property
    def next_bunch(self):
        """Return the rank of the next tracked bunch"""
        if self.rank + 1 in self.table[:, 0]:
            return self.rank + 1
        else:
            return 0

    @property
    def previous_bunch(self):
        """Return the rank of the previous tracked bunch"""
        if self.rank - 1 in self.table[:, 0]:
            return self.rank - 1
        else:
            return max(self.table[:, 0])

    def share_distributions(self, beam, dimensions="tau", n_bin=75):
        """

        dimension : str or list of str, optional
            Dimensions in which the binning is done. The default is "tau".
        n_bin : int or list of int, optional
            Number of bins. The default is 75.

        """

        if (beam.mpi_switch == False):
            print("Error, mpi is not initialised.")

        if isinstance(dimensions, str):
            dimensions = [dimensions]

        if isinstance(n_bin, int):
            n_bin = np.ones((len(dimensions), ), dtype=int) * n_bin

        bunch = beam[self.bunch_num]

        charge_per_mp_all = self.comm.allgather(bunch.charge_per_mp)
        self.charge_per_mp_all = charge_per_mp_all

        for i in range(len(dimensions)):

            dim = dimensions[i]
            n = n_bin[i]

            if len(bunch) != 0:
                bins, sorted_index, profile, center = bunch.binning(
                    dimension=dim, n_bin=n)
            else:
                sorted_index = None
                profile = np.zeros((n - 1, ), dtype=np.int64)
                center = np.zeros((n - 1, ), dtype=np.float64)
                if beam.filling_pattern[self.bunch_num] is True:
                    beam.update_filling_pattern()
                    beam.update_distance_between_bunches()

            self.__setattr__(dim + "_center",
                             np.empty((self.size, n - 1), dtype=np.float64))
            self.comm.Allgather(
                [center, self.MPI.DOUBLE],
                [self.__getattribute__(dim + "_center"), self.MPI.DOUBLE])

            self.__setattr__(dim + "_profile",
                             np.empty((self.size, n - 1), dtype=np.int64))
                [profile, self.MPI.INT64_T],
                [self.__getattribute__(dim + "_profile"), self.MPI.INT64_T])

            self.__setattr__(dim + "_sorted_index", sorted_index)

    def share_means(self, beam):
        """

        """

        charge_all = self.comm.allgather(bunch.charge)
        self.charge_all = charge_all

        self.mean_all = np.empty((self.size, 6), dtype=np.float64)
        if len(bunch) != 0:
            mean = bunch.mean
        else:
            mean = np.zeros((6, ), dtype=np.float64)
        self.comm.Allgather([mean, self.MPI.DOUBLE],
                            [self.mean_all, self.MPI.DOUBLE])

    def share_stds(self, beam):
        """

        """

        self.std_all = np.empty((self.size, 6), dtype=np.float64)
            std = bunch.std
        else:
            std = np.zeros((6, ), dtype=np.float64)
        self.comm.Allgather([std, self.MPI.DOUBLE],
                            [self.std_all, self.MPI.DOUBLE])

"""
Module where particles, bunches and beams are described as objects.
"""

import h5py as hp
import seaborn as sns
from scipy.constants import c, e, m_e, m_p

class Particle:
    """
    Define a particle object.

    mass : float
        total particle mass in [kg]
    charge : float
        electrical charge in [C]
    E_rest : float
        particle rest energy in [eV]
    """

    def __init__(self, mass, charge):
        self.mass = mass
        self.charge = charge

    @property
    def E_rest(self):
        return self.mass * c**2 / e

class Electron(Particle):
    """ Define an electron"""

    def __init__(self):
        super().__init__(m_e, -1 * e)

class Proton(Particle):
    """ Define a proton"""

        super().__init__(m_p, e)

class Bunch:
    """
    Define a bunch object.
    
    mp_number : float, optional
        Macro-particle number.
    current : float, optional
        Bunch current in [A].
    track_alive : bool, optional
        If False, the code no longer take into account alive/dead particles.
        Should be set to True if element such as apertures are used.
        Can be set to False to gain a speed increase.
    alive : bool, optional
        If False, the bunch is defined as empty.
    load_from_file : str, optional
        Name of bunch save file generated by save method.
        If None, the bunch is initialized using the input parameters.
        Otherwise, the bunch from the file is loaded and the other inputs are 
        ignored.
        Default is None.
    load_suffix : str or int, optional
        Suffix to the group name used to load the data from the HDF5 file.
        
    mp_number : int
    alive : array of bool of shape (mp_number,)
        Array used to monitor which particle is dead or alive.
    charge_per_mp : float
        Charge per macro-particle in [C].
    particle_number : int
        Number of particles in the bunch.
    is_empty : bool
        Return True if the bunch is empty.
    mean : array of shape (6,)
        Mean position of alive particles for each coordinates.
    std : array of shape (6,)
        Standard deviation of the position of alive particles for each 
        coordinates.
    emit : array of shape (3,)
        Bunch emittance for each plane [1]. !!! -> Correct for long ?
        
    Methods
    -------
    init_gaussian(cov=None, mean=None, **kwargs)
        Initialize bunch particles with 6D gaussian phase space.
    plot_phasespace(x_var="tau", y_var="delta", plot_type="j")
        Plot phase space.
    save(file_name)
        Save bunch object data (6D phase space, current, and state) in an HDF5 
        file format.
    load(file_name)
        Load data from a HDF5 file recorded by Bunch save method.
    
    [1] Wiedemann, H. (2015). Particle accelerator physics. 4th edition. 
    Springer, Eq.(8.39) of p224.
    """

                 ring,
                 mp_number=1e3,
                 current=1e-3,
                 track_alive=True,
                 alive=True,
                 load_from_file=None,
                 load_suffix=None):

        if not alive:
            mp_number = 1
            current = 0
        self._mp_number = int(mp_number)

        self.particles = {
            "x": np.zeros(self.mp_number, dtype=np.float64),
            "xp": np.zeros(self.mp_number, dtype=np.float64),
            "y": np.zeros(self.mp_number, dtype=np.float64),
            "yp": np.zeros(self.mp_number, dtype=np.float64),
            "tau": np.zeros(self.mp_number, dtype=np.float64),
            "delta": np.zeros(self.mp_number, dtype=np.float64),
        }
        self.track_alive = track_alive
        self.alive = np.ones((self.mp_number, ), dtype=bool)
        self.current = current
            self.alive = np.zeros((self.mp_number, ), dtype=bool)

        if load_from_file is not None:
            self.load(load_from_file, load_suffix, track_alive)

    def __len__(self):
        """Return the number of alive particles"""
        return self.alive.sum()

    def __getitem__(self, label):
        """Return the columns label for alive particles"""
        if self.track_alive is True:
            return self.particles[label][self.alive]
        else:
            return self.particles[label]

    def __setitem__(self, label, value):
        """Set value to the columns label for alive particles"""
            self.particles[label][self.alive] = value
        else:
            self.particles[label] = value

        """Iterate over labels"""
        return self.particles.keys().__iter__()

        """Return representation of alive particles"""
        rep = pd.DataFrame(np.array([self[l] for l in self.__iter__()]).T,
                           columns=list(self.__iter__()))
        return f'Bunch with macro-particles: \n {rep!r}'

    @property
    def mp_number(self):
        """Macro-particle number"""
        return self._mp_number

    @mp_number.setter
    def mp_number(self, value):
        self._mp_number = int(value)
        self.__init__(self.ring, value, self.charge)

    @property
    def charge_per_mp(self):
        """Charge per macro-particle [C]"""
        return self._charge_per_mp

    @charge_per_mp.setter
    def charge_per_mp(self, value):
        self._charge_per_mp = value

    @property
    def charge(self):
        """Bunch charge in [C]"""
        return self.__len__() * self.charge_per_mp

    @charge.setter
    def charge(self, value):
        self.charge_per_mp = value / self.__len__()

    @property
    def particle_number(self):
        """Particle number"""
        return int(self.charge / np.abs(self.ring.particle.charge))

    @particle_number.setter
    def particle_number(self, value):
        self.charge_per_mp = value * self.ring.particle.charge / self.__len__()

    @property
    def current(self):
        """Bunch current [A]"""
        return self.charge / self.ring.T0

    @current.setter
    def current(self, value):
        self.charge_per_mp = value * self.ring.T0 / self.__len__()

    @property
    def is_empty(self):
        """Return True if the bunch is empty."""
        return ~np.any(self.alive)

    @property
    def mean(self):
        """
        Return the mean position of alive particles for each coordinates.
        """
        mean = [[self[name].mean()] for name in self]
        return np.squeeze(np.array(mean))

    @property
    def std(self):
        """
        Return the standard deviation of the position of alive 
        particles for each coordinates.
        """
        std = [[self[name].std()] for name in self]
        return np.squeeze(np.array(std))

    @property
    def emit(self):
        """
        Return the bunch emittance for each plane.
        """
        cor = np.squeeze([[self[name] - self[name].mean()] for name in self])
        emitX = np.sqrt(
            np.mean(cor[0]**2) * np.mean(cor[1]**2) -
            np.mean(cor[0] * cor[1])**2)
        emitY = np.sqrt(
            np.mean(cor[2]**2) * np.mean(cor[3]**2) -
            np.mean(cor[2] * cor[3])**2)
        emitS = np.sqrt(
            np.mean(cor[4]**2) * np.mean(cor[5]**2) -
            np.mean(cor[4] * cor[5])**2)
        return np.array([emitX, emitY, emitS])

    @property
    def cs_invariant(self):
        """
        Return the average Courant-Snyder invariant of each plane.

        """
        Jx = (self.ring.optics.local_gamma[0] * self['x']**2) + \
              (2*self.ring.optics.local_alpha[0] * self['x'])*self['xp'] + \
              (self.ring.optics.local_beta[0] * self['xp']**2)
        Jy = (self.ring.optics.local_gamma[1] * self['y']**2) + \
              (2*self.ring.optics.local_alpha[1] * self['y']*self['yp']) + \
              (self.ring.optics.local_beta[1] * self['yp']**2)
        Js = (self.ring.long_gamma * self['tau']**2) + \
              (2*self.ring.long_alpha * self['tau']*self['delta']) + \
              (self.ring.long_beta * self['delta']**2)
        return np.array((np.mean(Jx), np.mean(Jy), np.mean(Js)))

    def init_gaussian(self, cov=None, mean=None, **kwargs):
        """
        Covariance matrix is taken from [1] and dispersion is added following
        the method explained in [2].
                
        cov : (6,6) array, optional
            Covariance matrix of the bunch distribution
        mean : (6,) array, optional
            Mean of the bunch distribution
        
        [1] Wiedemann, H. (2015). Particle accelerator physics. 4th 
        edition. Springer, Eq.(8.38) of p223.
        [2] http://www.pp.rhul.ac.uk/bdsim/manual-develop/dev_beamgeneration.html

        """
        if mean is None:
            mean = np.zeros((6, ))

        if cov is None:
            sigma_0 = kwargs.get("sigma_0", self.ring.sigma_0)
            sigma_delta = kwargs.get("sigma_delta", self.ring.sigma_delta)
            optics = kwargs.get("optics", self.ring.optics)

            cov = np.zeros((6, 6))
            cov[0, 0] = self.ring.emit[0] * optics.local_beta[0] + (
                optics.local_dispersion[0] * self.ring.sigma_delta)**2
            cov[1, 1] = self.ring.emit[0] * optics.local_gamma[0] + (
                optics.local_dispersion[1] * self.ring.sigma_delta)**2
            cov[0, 1] = -1 * self.ring.emit[0] * optics.local_alpha[0] + (
                optics.local_dispersion[0] * optics.local_dispersion[1] *
                self.ring.sigma_delta**2)
            cov[1, 0] = -1 * self.ring.emit[0] * optics.local_alpha[0] + (
            cov[0, 5] = optics.local_dispersion[0] * self.ring.sigma_delta**2
            cov[5, 0] = optics.local_dispersion[0] * self.ring.sigma_delta**2
            cov[1, 5] = optics.local_dispersion[1] * self.ring.sigma_delta**2
            cov[5, 1] = optics.local_dispersion[1] * self.ring.sigma_delta**2
            cov[2, 2] = self.ring.emit[1] * optics.local_beta[1] + (
                optics.local_dispersion[2] * self.ring.sigma_delta)**2
            cov[3, 3] = self.ring.emit[1] * optics.local_gamma[1] + (
                optics.local_dispersion[3] * self.ring.sigma_delta)**2
            cov[2, 3] = -1 * self.ring.emit[1] * optics.local_alpha[1] + (
                optics.local_dispersion[2] * optics.local_dispersion[3] *
            cov[3, 2] = -1 * self.ring.emit[1] * optics.local_alpha[1] + (
            cov[2, 5] = optics.local_dispersion[2] * self.ring.sigma_delta**2
            cov[5, 2] = optics.local_dispersion[2] * self.ring.sigma_delta**2
            cov[3, 5] = optics.local_dispersion[3] * self.ring.sigma_delta**2
            cov[5, 3] = optics.local_dispersion[3] * self.ring.sigma_delta**2
            cov[4, 4] = sigma_0**2
            cov[5, 5] = sigma_delta**2

        values = np.random.multivariate_normal(mean, cov, size=self.mp_number)
        self.particles["x"] = values[:, 0]
        self.particles["xp"] = values[:, 1]
        self.particles["y"] = values[:, 2]
        self.particles["yp"] = values[:, 3]
        self.particles["tau"] = values[:, 4]
        self.particles["delta"] = values[:, 5]

    def binning(self, dimension="tau", n_bin=75):
        """
        Bin macro-particles.

        dimension : str, optional
            Dimension in which the binning is done. The default is "tau".
        n_bin : int, optional

        Returns
        -------
        bins : array of shape (n_bin,)
            Bins where the particles are sorted.
        sorted_index : array of shape (self.mp_number,)
            Bin number of each macro-particles.
        profile : array of shape (n_bin - 1,)
            Number of marco-particles in each bin.
        center : array of shape (n_bin - 1,)
            Center of each bin.

        """
        bin_min = self[dimension].min()
        bin_min = min(bin_min * 0.99, bin_min * 1.01)
        bin_max = self[dimension].max()
        bin_max = max(bin_max * 0.99, bin_max * 1.01)

        bins = np.linspace(bin_min, bin_max, n_bin)
        center = (bins[1:] + bins[:-1]) / 2
        sorted_index = np.searchsorted(bins, self[dimension], side='left')
        sorted_index -= 1
        profile = np.bincount(sorted_index, minlength=n_bin - 1)

        return (bins, sorted_index, profile, center)

    def plot_profile(self, dimension="tau", n_bin=75):
        """
        Plot bunch profile.

            Dimension to plot. The default is "tau".

        """
        bins, sorted_index, profile, center = self.binning(dimension, n_bin)
        ax.plot(center, profile)

    def plot_phasespace(self, x_var="tau", y_var="delta", kind="scatter"):
        """
        
        x_var : str 
            Dimension to plot on horizontal axis.
        y_var : str 
            Dimension to plot on vertical axis.
        kind : str, {"scatter", "kde", "hist", "hex", "reg", "resid"}
            Kind of plot to draw. See seeborn documentation.
            The defualt value is "scatter".
            
        Return
        ------
            Figure object with the plot on it.
        """

        label_dict = {
            "x": "x (mm)",
            "xp": "x' (mrad)",
            "y": "y (mm)",
            "yp": "y' (mrad)",
            "tau": "$\\tau$ (ps)",
            "delta": "$\\delta$"
        }
        scale = {
            "x": 1e3,
            "xp": 1e3,
            "y": 1e3,
            "yp": 1e3,
            "tau": 1e12,
            "delta": 1
        }

        fig = sns.jointplot(x=self.particles[x_var] * scale[x_var],
                            y=self.particles[y_var] * scale[y_var],
                            kind=kind)
        plt.xlabel(label_dict[x_var])
        plt.ylabel(label_dict[y_var])

        return fig

    def save(self, file_name, suffix=None, mpi_comm=None):
        """
        
        The output file is named as "<file_name>.hdf5".

        file_name : str
            Name of the output file.
        suffix : str or int, optional
            Suffix to the group name used to save the data into the HDF5 file.
            If None, it is not used.
        mpi_comm : MPI communicator, optional
            For internal use if mpi is used in Beam objects.
            Default is None.
        """

        if mpi_comm is None:
            f = hp.File(file_name + ".hdf5", "a", libver="earliest")
        else:
            f = hp.File(file_name + ".hdf5",
                        "a",
                        libver='earliest',
                        driver='mpio',
                        comm=mpi_comm)

        if suffix is None:
            group_name = "Bunch"
        else:
            group_name = "Bunch_" + str(suffix)

        g = f.create_group(group_name)
        g.create_dataset("alive", (self.mp_number, ), dtype=bool)
        g.create_dataset("phasespace", (self.mp_number, 6), dtype=float)
        g.create_dataset("current", (1, ), dtype=float)

        f[group_name]["alive"][:] = self.alive
        f[group_name]["current"][:] = self.current
        for i, dim in enumerate(self):
            f[group_name]["phasespace"][:, i] = self.particles[dim]

        f.close()

    def load(self, file_name, suffix=None, track_alive=True):
        """

            Path to the HDF5 file where the Bunch data is stored.
            Suffix to the group name used to load the data from the HDF5 file.
        track_alive : bool, optional
            If False, the code no longer take into account alive/dead particles.
            Should be set to True if element such as apertures are used.
            Can be set to False to gain a speed increase.
        """

        f = hp.File(file_name, "r")

        else:

        self.mp_number = len(f[group_name]['alive'][:])

            self.particles[dim] = f[group_name]["phasespace"][:, i]

        self.alive = f[group_name]['alive'][:]
        if f[group_name]['current'][:][0] != 0:
            self.current = f[group_name]['current'][:][0]
        else:
            self.charge_per_mp = 0

        f.close()

class Beam:
    """
    Define a Beam object composed of several Bunch objects. 
    
    bunch_list : list of Bunch object, optional

        Total beam current in [A]
        Total bunch charge in [C]
        Total number of particle in the beam
    filling_pattern : bool array of shape (ring.h,)
        Filling pattern of the beam
    bunch_current : array of shape (ring.h,)
        Current in each bunch in [A]
    bunch_charge : array of shape (ring.h,)
        Charge in each bunch in [C]
    bunch_particle : array of shape (ring.h,)
        Particle number in each bunch
    bunch_mean : array of shape (6, ring.h)
        Mean position of alive particles for each bunch
    bunch_std : array of shape (6, ring.h)
        Standard deviation of the position of alive particles for each bunch        
    bunch_emit : array of shape (6, ring.h)
        Bunch emittance of alive particles for each bunch
    mpi : Mpi object
    mpi_switch : bool
        Status of MPI parallelisation, should not be changed directly but with
        mpi_init() and mpi_close()
    bunch_index : array of shape (len(self,))
        Return an array with the positions of the non-empty bunches
        
    Methods
    ------
    init_beam(filling_pattern, current_per_bunch=1e-3, mp_per_bunch=1e3)
        Initialize beam with a given filling pattern and marco-particle number 
        per bunch. Then initialize the different bunches with a 6D gaussian
        phase space.
    mpi_init()
        Switch on MPI parallelisation and initialise a Mpi object
    mpi_gather()
        Gather beam, all bunches of the different processors are sent to 
        all processors. Rather slow
    mpi_share_distributions()
        Compute the bunch profile and share it between the different bunches.
    mpi_close()
        Call mpi_gather and switch off MPI parallelisation
    plot(var, option=None)
        Plot variables with respect to bunch number.
        Save beam object data in an HDF5 file format.
    load(file_name, mpi)
        Load data from a HDF5 file recorded by Beam save method.
    """

    def __init__(self, ring, bunch_list=None):
        self.mpi_switch = False
        if bunch_list is None:
            self.init_beam(np.zeros((self.ring.h, 1), dtype=bool))
        else:
            if (len(bunch_list) != self.ring.h):
                raise ValueError(("The length of the bunch list is {} ".format(
                    len(bunch_list)) + "but should be {}".format(self.ring.h)))
            self.bunch_list = bunch_list

        """Return the number of (not empty) bunches"""
        length = 0
        for bunch in self.not_empty:
            length += 1
        return length

    def __getitem__(self, i):
        """Return the bunch number i"""
        return self.bunch_list.__getitem__(i)

    def __setitem__(self, i, value):
        """Set value to the bunch number i"""
        self.bunch_list.__setitem__(i, value)

        """Iterate over all bunches"""
        return self.bunch_list.__iter__()

        """Return representation of the beam filling pattern"""
        return f'Beam with bunch current:\n {list((self.bunch_current))!r}'

    @property
    def not_empty(self):
        """Return a generator to iterate over not empty bunches."""
        for index, value in enumerate(self.filling_pattern):
            if value == True:
                yield self[index]
            else:
                pass

    @property
    def distance_between_bunches(self):
        """Return an array which contains the distance to the next bunch in 
        units of the RF period (ring.T1)"""
        return self._distance_between_bunches

    def update_distance_between_bunches(self):
        """Update the distance_between_bunches array"""
        filling_pattern = self.filling_pattern
        distance = np.zeros(filling_pattern.shape)
        last_value = 0

        # All bunches
        for index, value in enumerate(filling_pattern):
            if value == False:
                pass
            elif value == True:
                last_value = index
                count = 1
                for value2 in filling_pattern[index + 1:]:
                    if value2 == False:
                        count += 1
                    elif value2 == True:
                        break
                distance[index] = count

        # Last bunch case
        count2 = 0
        for index2, value2 in enumerate(filling_pattern):
            if value2 == True:
                break
            if value2 == False:
                count2 += 1
        distance[last_value] += count2

        self._distance_between_bunches = distance

    def init_beam(self,
                  filling_pattern,
                  current_per_bunch=1e-3,
                  mp_per_bunch=1e3,
                  track_alive=True,
                  mpi=False):
        """
        
        If the filling pattern is an array of bool then the current per bunch 
        is uniform, else the filling pattern can be an array with the current
        in each bunch.
        
        filling_pattern : numpy array or list of length ring.h
            Filling pattern of the beam, can be a list or an array of bool, 
            then current_per_bunch is used. Or can be an array with the current
            in each bunch.
        current_per_bunch : float, optional
            Current per bunch in [A]
        mp_per_bunch : float, optional
            Macro-particle number per bunch
        mpi : bool, optional
            If True, only a single bunch is fully initialized on each core, the
            other bunches are initialized with a single marco-particle.
        """

        if (len(filling_pattern) != self.ring.h):
            raise ValueError(("The length of filling pattern is {} ".format(
                len(filling_pattern)) +
                              "but should be {}".format(self.ring.h)))

        if mpi is True:
            mp_per_bunch_mpi = mp_per_bunch
            mp_per_bunch = 1

        filling_pattern = np.array(filling_pattern)
        bunch_list = []
        if filling_pattern.dtype == np.dtype("bool"):
            for value in filling_pattern:
                if value == True:
                    bunch_list.append(
                        Bunch(self.ring, mp_per_bunch, current_per_bunch,
                              track_alive))
                elif value == False:
                    bunch_list.append(Bunch(self.ring, alive=False))
        elif filling_pattern.dtype == np.dtype("float64"):
            for current in filling_pattern:
                if current != 0:
                        Bunch(self.ring, mp_per_bunch, current, track_alive))
                elif current == 0:
        else:
            raise TypeError("{} should be bool or float64".format(
                filling_pattern.dtype))

        self.bunch_list = bunch_list
        self.update_filling_pattern()
        self.update_distance_between_bunches()

            self.mpi_init()
            current = self[self.mpi.rank_to_bunch(self.mpi.rank)].current
            bunch = Bunch(self.ring, mp_per_bunch_mpi, current, track_alive)
            bunch.init_gaussian()
            self[self.mpi.rank_to_bunch(self.mpi.rank)] = bunch
        else:
            for bunch in self.not_empty:
                bunch.init_gaussian()

    def update_filling_pattern(self):
        """Update the beam filling pattern."""
        filling_pattern = []
        for bunch in self:
            if bunch.current != 0:
                filling_pattern.append(True)
            else:
                filling_pattern.append(False)
        self._filling_pattern = np.array(filling_pattern)

    @property
    def filling_pattern(self):
        """Return an array with the filling pattern of the beam as bool"""
        return self._filling_pattern

    @property
    def bunch_index(self):
        """Return an array with the positions of the non-empty bunches."""
        return np.where(self.filling_pattern == True)[0]

    @property
    def bunch_current(self):
        """Return an array with the current in each bunch in [A]"""
        bunch_current = [bunch.current for bunch in self]
        return np.array(bunch_current)

    @property
    def bunch_charge(self):
        """Return an array with the charge in each bunch in [C]"""
        bunch_charge = [bunch.charge for bunch in self]
        return np.array(bunch_charge)

    @property
    def bunch_particle(self):
        """Return an array with the particle number in each bunch"""
        bunch_particle = [bunch.particle_number for bunch in self]
        return np.array(bunch_particle)

    @property
        """Total beam current in [A]"""
        return np.sum(self.bunch_current)

    @property
        """Total beam charge in [C]"""
        return np.sum(self.bunch_charge)

    @property
        """Total number of particles in the beam"""
        return np.sum(self.bunch_particle)

    @property
    def bunch_mean(self):
        """Return an array with the mean position of alive particles for each
        bunches"""
        bunch_mean = np.zeros((6, self.ring.h))
        for idx, bunch in enumerate(self.not_empty):
            index = self.bunch_index[idx]
            bunch_mean[:, index] = bunch.mean
        return bunch_mean

    @property
    def bunch_std(self):
        """Return an array with the standard deviation of the position of alive 
        particles for each bunches"""
        bunch_std = np.zeros((6, self.ring.h))
            bunch_std[:, index] = bunch.std
        return bunch_std

    @property
    def bunch_emit(self):
        """Return an array with the bunch emittance of alive particles for each
        bunches and each plane"""
        bunch_emit = np.zeros((3, self.ring.h))
            bunch_emit[:, index] = bunch.emit
        return bunch_emit

    @property
    def bunch_cs(self):
        """Return an array with the average Courant-Snyder invariant for each 
        bunch"""
        bunch_cs = np.zeros((3, self.ring.h))
            bunch_cs[:, index] = bunch.cs_invariant
        return bunch_cs

    def mpi_init(self):
        """Switch on MPI parallelisation and initialise a Mpi object"""
        from mbtrack2.tracking.parallel import Mpi
        self.mpi = Mpi(self.filling_pattern)
        self.mpi_switch = True

    def mpi_gather(self):
        """Gather beam, all bunches of the different processors are sent to 
        all processors. Rather slow"""

        if (self.mpi_switch == False):

        bunch = self[self.mpi.bunch_num]
        bunches = self.mpi.comm.allgather(bunch)
        for rank in range(self.mpi.size):
            self[self.mpi.rank_to_bunch(rank)] = bunches[rank]

    def mpi_close(self):
        """Call mpi_gather and switch off MPI parallelisation"""
        self.mpi_gather()
        self.mpi = None

    def plot(self, var, option=None):
        """

        var : str {"bunch_current", "bunch_charge", "bunch_particle", 
                   "bunch_mean", "bunch_std", "bunch_emit"}
            Variable to be plotted.
        option : str, optional
            If var is "bunch_mean", "bunch_std", or "bunch_emit, option needs 
            to be specified.
            For "bunch_mean" and "bunch_std", 
                option = {"x","xp","y","yp","tau","delta"}.
            For "bunch_emit", option = {"x","y","s"}.
            
        Return
        ------
        """

        var_dict = {
            "bunch_current": self.bunch_current,
            "bunch_charge": self.bunch_charge,
            "bunch_particle": self.bunch_particle,
            "bunch_mean": self.bunch_mean,
            "bunch_std": self.bunch_std,
            "bunch_emit": self.bunch_emit
        }

        fig, ax = plt.subplots()

        if var == "bunch_mean" or var == "bunch_std":
            value_dict = {
                "x": 0,
                "xp": 1,
                "y": 2,
                "yp": 3,
                "tau": 4,
                "delta": 5
            }
            scale = [1e6, 1e6, 1e6, 1e6, 1e12, 1]
            label_mean = [
                "x (um)", "x' ($\\mu$rad)", "y (um)", "y' ($\\mu$rad)",
                "$\\tau$ (ps)", "$\\delta$"
            ]
            label_std = [
                "std x (um)", "std x' ($\\mu$rad)", "std y (um)",
                "std y' ($\\mu$rad)", "std $\\tau$ (ps)", "std $\\delta$"
            ]

            y_axis = var_dict[var][value_dict[option]]

            # Convert NaN in y_axis array into zero
            where_is_nan = np.isnan(y_axis)
            y_axis[where_is_nan] = 0

            ax.plot(np.arange(len(self.filling_pattern)),
                    y_axis * scale[value_dict[option]])
            ax.set_xlabel('bunch number')
            if var == "bunch_mean":
                ax.set_ylabel(label_mean[value_dict[option]])
            else:
                ax.set_ylabel(label_std[value_dict[option]])

        elif var == "bunch_emit":
            value_dict = {"x": 0, "y": 1, "s": 2}
            scale = [1e9, 1e9, 1e15]

            if option == "x": label_y = "hor. emittance (nm.rad)"
            elif option == "y": label_y = "ver. emittance (nm.rad)"
            elif option == "s": label_y = "long. emittance (fm.rad)"

            ax.set_ylabel(label_y)

        elif var == "bunch_current" or var == "bunch_charge" or var == "bunch_particle":
            scale = {
                "bunch_current": 1e3,
                "bunch_charge": 1e9,
                "bunch_particle": 1
            }

                    var_dict[var] * scale[var])

            if var == "bunch_current": label_y = "bunch current (mA)"
            elif var == "bunch_charge": label_y = "bunch chagre (nC)"
            else: label_y = "number of particles"

        elif var == "current" or var == "charge" or var == "particle_number":
                "'{0}'is a total value and cannot be plotted.".format(var))

        return fig

    def save(self, file_name):
        """
        
        """
        if self.mpi_switch is True:
            for i, bunch in enumerate(self):
                if i in self.mpi.table[:, 1]:
                    if i == self.mpi.bunch_num:
                        mp_number = self[self.mpi.bunch_num].mp_number
                        self.mpi.comm.bcast(mp_number, root=self.mpi.rank)
                        self[self.mpi.bunch_num].save(file_name,
                                                      self.mpi.bunch_num,
                                                      self.mpi.comm)
                    else:
                        mp_number = None
                        mp_number = self.mpi.comm.bcast(
                            mp_number, root=self.mpi.bunch_to_rank(i))
                        f = hp.File(file_name + ".hdf5",
                                    "a",
                                    libver='earliest',
                                    driver='mpio',
                                    comm=self.mpi.comm)
                        group_name = "Bunch_" + str(i)
                        g = f.create_group(group_name)
                        g.create_dataset("alive", (mp_number, ), dtype=bool)
                        g.create_dataset("phasespace", (mp_number, 6),
                                         dtype=float)
                        g.create_dataset("current", (1, ), dtype=float)
                        f.close()
                else:
                    bunch.save(file_name, i, self.mpi.comm)
        else:
                bunch.save(file_name, i)

    def load(self, file_name, mpi, track_alive=True):
        """

            Path to the HDF5 file where the Beam data is stored.
        mpi : bool
            The default is True.
        """
            self.__init__(self.ring)
            f = hp.File(file_name, "r")
            filling_pattern = []
            for i in range(self.ring.h):
                current = f["Bunch_" + str(i)]['current'][:][0]
                filling_pattern.append(current)

            self.init_beam(filling_pattern,
                           mp_per_bunch=1,
                           track_alive=track_alive,
                           mpi=True)
            self[self.mpi.bunch_num].load(file_name, self.mpi.bunch_num)
        else:
                bunch.load(file_name, i, track_alive)
            self.update_filling_pattern()
            self.update_distance_between_bunches()

"""
This module handles radio-frequency (RF) cavitiy elements. 
"""

import matplotlib.patches as mpatches
from matplotlib.legend_handler import HandlerPatch

from mbtrack2.instability import lcbi_growth_rate

class RFCavity(Element):
    """
    Perfect RF cavity class for main and harmonic RF cavities.
    Use cosine definition.
    
    m : int
        Harmonic number of the cavity
    Vc : float
        Amplitude of cavity voltage [V]
    theta : float
        Phase of Cavity voltage
    """

    def __init__(self, ring, m, Vc, theta):
        self.m = m
        self.Vc = Vc
        self.theta = theta

        """
        
        """
        bunch["delta"] += self.Vc / self.ring.E0 * np.cos(
            self.m * self.ring.omega1 * bunch["tau"] + self.theta)

    def value(self, val):
        return self.Vc / self.ring.E0 * np.cos(self.m * self.ring.omega1 *
                                               val + self.theta)

class CavityResonator():
    """Cavity resonator class for active or passive RF cavity with beam
    loading or HOM, based on [1,2].
    
    If used with mpi, beam.mpi.share_distributions must be called before the 
    track method call.
    
    Different kind of RF feeback and loops can be added using:
        cavity_resonator.feedback.append(loop)
    
    m : int or float
        Harmonic number of the cavity.
    Rs : float
        Shunt impedance of the cavities in [Ohm], defined as 0.5*Vc*Vc/Pc.
        If Ncav = 1, used for the total shunt impedance.
        If Ncav > 1, used for the shunt impedance per cavity.
    Q : float
        Quality factor of the cavity.
    QL : float
        Loaded quality factor of the cavity.
    detune : float
        Detuing of the cavity in [Hz], defined as (fr - m*ring.f1).
        Number of cavities.
    Vc : float, optinal
        Total cavity voltage objective value in [V].
    theta : float, optional
        Total cavity phase objective value in [rad].
    n_bin : int, optional
        Number of bins used for the beam loading computation. 
        Only used if MPI is not used, otherwise n_bin must be specified in the 
        beam.mpi.share_distributions method.
        The default is 75.
        
    beam_phasor : complex
        Beam phasor in [V].
    beam_phasor_record : array of complex
        Last beam phasor value of each bunch in [V].
    generator_phasor : complex
        Generator phasor in [V].
    generator_phasor_record : array of complex
        Last generator phasor value of each bunch in [V].
    cavity_phasor : complex
        Cavity phasor in [V].
    cavity_phasor_record : array of complex
        Last cavity phasor value of each bunch in [V].
    ig_phasor_record : array of complex
        Last current generator phasor of each bunch in [A].
        Only used for some feedback types.
    cavity_voltage : float
        Cavity total voltage in [V].
    cavity_phase : float
        Cavity total phase in [rad].
    loss_factor : float
        Cavity loss factor in [V/C].
    Rs_per_cavity : float
        Shunt impedance of a single cavity in [Ohm], defined as 0.5*Vc*Vc/Pc.
        Coupling coefficient of the cavity.
    fr : float
        Resonance frequency of the cavity in [Hz].
    wr : float
        Angular resonance frequency in [Hz.rad].
    psi : float
        Tuning angle in [rad].
    filling_time : float
        Cavity filling time in [s].
    Pc : float
        Power dissipated in the cavity walls in [W].
    Pg : float
        Generator power in [W].
    Vgr : float
        Generator voltage at resonance in [V].
    theta_gr : float
        Generator phase at resonance in [rad].
    Vg : float
        Generator voltage in [V].
    theta_g : float
        Generator phase in [rad].
    tracking : bool
        True if the tracking has been initialized.
    bunch_index : int
        Number of the tracked bunch in the current core.
    distance : array
        Distance between bunches.
    valid_bunch_index : array
    
    Methods
    -------
    Vbr(I0)
        Return beam voltage at resonance in [V].
    Vb(I0)
        Return beam voltage in [V].
    Pb(I0)
        Return power transmitted to the beam in [W].
    Pr(I0)
        Return power reflected back to the generator in [W].
    Z(f)
        Cavity impedance in [Ohm] for a given frequency f in [Hz].
    set_optimal_coupling(I0)
        Set coupling to optimal value.
    set_optimal_detune(I0)
        Set detuning to optimal conditions.
    set_generator(I0)
        Set generator parameters.
    plot_phasor(I0)
        Plot phasor diagram.
    is_DC_Robinson_stable(I0)
        Check DC Robinson stability.
    plot_DC_Robinson_stability()
        Plot DC Robinson stability limit.
    init_tracking(beam)
        Initialization of the tracking.
    track(beam)
    phasor_decay(time)
        Compute the beam phasor decay during a given time span.
    phasor_evol(profile, bin_length, charge_per_mp)
        Compute the beam phasor evolution during the crossing of a bunch.
    VRF(z, I0)
        Return the total RF voltage.
    dVRF(z, I0)
        Return derivative of total RF voltage.
    ddVRF(z, I0)
        Return the second derivative of total RF voltage.
    deltaVRF(z, I0)
        Return the generator voltage minus beam loading voltage.
    update_feedback()
        Force feedback update from current CavityResonator parameters.
    sample_voltage()
        Sample the voltage seen by a zero charge particle during an RF period.
    
    [1] Wilson, P. B. (1994). Fundamental-mode rf design in e+ e− storage ring 
    factories. In Frontiers of Particle Beams: Factories with e+ e-Rings 
    (pp. 293-311). Springer, Berlin, Heidelberg.
    
    [2] Naoto Yamamoto, Alexis Gamelin, and Ryutaro Nagaoka. "Investigation 
    of Longitudinal Beam Dynamics With Harmonic Cavities by Using the Code 
    Mbtrack." IPAC’19, Melbourne, Australia, 2019.
    """

                 ring,
                 m,
                 Rs,
                 Q,
                 QL,
                 detune,
                 Ncav=1,
                 Vc=0,
                 theta=0,
                 n_bin=75):
        self.feedback = []
        self.m = m
        self.Ncav = Ncav
        if Ncav != 1:
            self.Rs_per_cavity = Rs
        else:
            self.Rs = Rs
        self.Q = Q
        self.QL = QL
        self.detune = detune
        self.beam_phasor = np.zeros(1, dtype=complex)
        self.beam_phasor_record = np.zeros((self.ring.h), dtype=complex)
        self.generator_phasor_record = np.zeros((self.ring.h), dtype=complex)
        self.tracking = False
        self.Vg = 0
        self.theta_g = 0
        self.Vgr = 0
        self.theta_gr = 0
        self.Pg = 0
        self.n_bin = int(n_bin)

    def init_tracking(self, beam):
        """

        """
        if beam.mpi_switch:
            self.bunch_index = beam.mpi.bunch_num  # Number of the tracked bunch in this processor

        self.distance = beam.distance_between_bunches
        self.valid_bunch_index = beam.bunch_index
        self.tracking = True
        self.nturn = 0

        """
        Track a Beam object through the CavityResonator object.
        
        Can be used with or without mpi.
        If used with mpi, beam.mpi.share_distributions must be called before.
        
        The beam phasor is given at t=0 (synchronous particle) of the first 
        non empty bunch.

        """

        if self.tracking is False:
            self.init_tracking(beam)

        for index, bunch in enumerate(beam):

            if beam.filling_pattern[index]:

                if beam.mpi_switch:
                    # get rank of bunch n° index
                    rank = beam.mpi.bunch_to_rank(index)
                    # mpi -> get shared bunch profile for current bunch
                    center = beam.mpi.tau_center[rank]
                    profile = beam.mpi.tau_profile[rank]
                    bin_length = center[1] - center[0]
                    charge_per_mp = beam.mpi.charge_per_mp_all[rank]
                    if index == self.bunch_index:
                        sorted_index = beam.mpi.tau_sorted_index
                else:
                    # no mpi -> get bunch profile for current bunch
                    if len(bunch) != 0:
                        (bins, sorted_index, profile,
                         center) = bunch.binning(n_bin=self.n_bin)
                        bin_length = center[1] - center[0]
                        charge_per_mp = bunch.charge_per_mp
                        self.bunch_index = index
                    else:
                        # Update filling pattern
                        beam.update_filling_pattern()
                        beam.update_distance_between_bunches()
                        # save beam phasor value
                        self.beam_phasor_record[index] = self.beam_phasor
                        # phasor decay to be at t=0 of the next bunch
                        self.phasor_decay(self.ring.T1, ref_frame="beam")
                        continue

                energy_change = bunch["tau"] * 0

                # remove part of beam phasor decay to be at the start of the binning (=bins[0])
                self.phasor_decay(center[0] - bin_length/2, ref_frame="beam")

                if index != self.bunch_index:
                    self.phasor_evol(profile,
                                     bin_length,
                                     charge_per_mp,
                                     ref_frame="beam")
                else:
                    # modify beam phasor
                    for i, center0 in enumerate(center):
                        mp_per_bin = profile[i]

                        if mp_per_bin == 0:
                            self.phasor_decay(bin_length, ref_frame="beam")
                            continue

                        ind = (sorted_index == i)
                        phase = self.m * self.ring.omega1 * (
                            center0 + self.ring.T1 *
                            (index + self.ring.h * self.nturn))
                        Vgene = np.real(self.generator_phasor_record[index] *
                                        np.exp(1j * phase))
                        Vbeam = np.real(self.beam_phasor)
                        Vtot = Vgene + Vbeam - charge_per_mp * self.loss_factor * mp_per_bin
                        energy_change[ind] = Vtot / self.ring.E0

                        self.beam_phasor -= 2 * charge_per_mp * self.loss_factor * mp_per_bin
                        self.phasor_decay(bin_length, ref_frame="beam")

                # phasor decay to be at t=0 of the current bunch (=-1*bins[-1])
                self.phasor_decay(-1 * (center[-1] + bin_length/2),
                                  ref_frame="beam")

                if index == self.bunch_index:
                    # apply kick
                    bunch["delta"] += energy_change

            # save beam phasor value
            self.beam_phasor_record[index] = self.beam_phasor

            # phasor decay to be at t=0 of the next bunch
            self.phasor_decay(self.ring.T1, ref_frame="beam")

        # apply different kind of RF feedback
        for fb in self.feedback:
            fb.track()

        self.nturn += 1

    def init_phasor_track(self, beam):
        """
        Initialize the beam phasor for a given beam distribution using a
        tracking like method.
        
        Follow the same steps as the track method but in the "rf" reference 
        frame and without any modifications on the beam.

        """

        n_turn = int(self.filling_time / self.ring.T0 * 10)

        for i in range(n_turn):
            for j, bunch in enumerate(beam.not_empty):

                index = self.valid_bunch_index[j]

                    # get shared bunch profile for current bunch
                    center = beam.mpi.tau_center[j]
                    profile = beam.mpi.tau_profile[j]
                    charge_per_mp = beam.mpi.charge_per_mp_all[j]
                else:
                    if i == 0:
                        # get bunch profile for current bunch
                        if j == 0:
                            self.profile_save = np.zeros((
                                len(beam),
                                len(profile),
                            ))
                            self.center_save = np.zeros((
                                len(center),
                            ))
                        self.profile_save[j, :] = profile
                        self.center_save[j, :] = center
                    else:
                        profile = self.profile_save[j, :]
                        center = self.center_save[j, :]

                    charge_per_mp = bunch.charge_per_mp

                self.phasor_decay(center[0] - bin_length/2, ref_frame="rf")
                self.phasor_evol(profile,
                                 bin_length,
                                 charge_per_mp,
                                 ref_frame="rf")
                                  ref_frame="rf")
                self.phasor_decay((self.distance[index] * self.ring.T1),

    def phasor_decay(self, time, ref_frame="beam"):
        """
        Compute the beam phasor decay during a given time span, assuming that 
        no particles are crossing the cavity during the time span.

        time : float
            Time span in [s], can be positive or negative.
        ref_frame : string, optional
            Reference frame to be used, can be "beam" or "rf".

        """
        if ref_frame == "beam":
            delta = self.wr
        elif ref_frame == "rf":
            delta = (self.wr - self.m * self.ring.omega1)
        self.beam_phasor = self.beam_phasor * np.exp(
            (-1 / self.filling_time + 1j*delta) * time)

    def phasor_evol(self,
                    profile,
                    bin_length,
                    charge_per_mp,
                    ref_frame="beam"):
        """
        Compute the beam phasor evolution during the crossing of a bunch using 
        an analytic formula [1].
        
        Assume that the phasor decay happens before the beam loading.

        profile : array
            Longitudinal profile of the bunch in [number of macro-particle].
        bin_length : float
            Length of a bin in [s].
        charge_per_mp : float
            Charge per macro-particle in [C].
            
        [1] mbtrack2 manual.
            
        """

        n_bin = len(profile)

        # Phasor decay during crossing time
        deltaT = n_bin * bin_length
        self.phasor_decay(deltaT, ref_frame)

        # Phasor evolution due to induced voltage by marco-particles
        k = np.arange(0, n_bin)
        var = np.exp(
            (-1 / self.filling_time + 1j*delta) * (n_bin-k) * bin_length)
        sum_tot = np.sum(profile * var)
        sum_val = -2 * sum_tot * charge_per_mp * self.loss_factor
        self.beam_phasor += sum_val

    def init_phasor(self, beam):
        """
        Initialize the beam phasor for a given beam distribution using an
        analytic formula [1].
        
        No modifications on the Beam object.

        """

        # Initialization

        N = self.n_bin - 1
        delta = (self.wr - self.m * self.ring.omega1)

        T = np.ones(self.ring.h) * self.ring.T1
        bin_length = np.zeros(self.ring.h)
        charge_per_mp = np.zeros(self.ring.h)
        profile = np.zeros((N, self.ring.h))
        center = np.zeros((N, self.ring.h))

        # Gather beam distribution data
        for j, bunch in enumerate(beam.not_empty):
            index = self.valid_bunch_index[j]
            if beam.mpi_switch:
                beam.mpi.share_distributions(beam, n_bin=self.n_bin)
                center[:, index] = beam.mpi.tau_center[j]
                profile[:, index] = beam.mpi.tau_profile[j]
                bin_length[index] = center[1, index] - center[0, index]
                charge_per_mp[index] = beam.mpi.charge_per_mp_all[j]
            else:
                (bins, sorted_index, profile[:, index],
                 center[:, index]) = bunch.binning(n_bin=self.n_bin)
                charge_per_mp[index] = bunch.charge_per_mp
            T[index] -= (center[-1, index] + bin_length[index] / 2)
            if index != 0:
                T[index - 1] += (center[0, index] - bin_length[index] / 2)
        T[self.ring.h - 1] += (center[0, 0] - bin_length[0] / 2)

        # Compute matrix coefficients
        k = np.arange(0, N)
        Tkj = np.zeros((N, self.ring.h))
        for j in range(self.ring.h):
            sum_t = np.array(
                [T[n] + N * bin_length[n] for n in range(j + 1, self.ring.h)])
            Tkj[:, j] = (N-k) * bin_length[j] + T[j] + np.sum(sum_t)

        var = np.exp((-1 / self.filling_time + 1j*delta) * Tkj)
        sum_tot = np.sum((profile*charge_per_mp) * var)

        # Use the formula n_turn times
            # Phasor decay during one turn
            self.phasor_decay(self.ring.T0, ref_frame="rf")
            # Phasor evolution due to induced voltage by marco-particles during one turn
            sum_val = -2 * sum_tot * self.loss_factor
            self.beam_phasor += sum_val

        # Replace phasor at t=0 (synchronous particle) of the first non empty bunch.
        idx0 = self.valid_bunch_index[0]
        self.phasor_decay(center[-1, idx0] + bin_length[idx0] / 2,
                          ref_frame="rf")

    @property
    def generator_phasor(self):
        """Generator phasor in [V]"""
        return self.Vg * np.exp(1j * self.theta_g)

    @property
    def cavity_phasor(self):
        """Cavity total phasor in [V]"""
        return self.generator_phasor + self.beam_phasor

    @property
    def cavity_phasor_record(self):
        """Last cavity phasor value of each bunch in [V]"""
        return self.generator_phasor_record + self.beam_phasor_record

    @property
    def ig_phasor_record(self):
        """Last current generator phasor of each bunch in [A]"""
        for FB in self.feedback:
            if isinstance(FB, (ProportionalIntegralLoop, DirectFeedback)):
                return FB.ig_phasor_record
        return np.zeros(self.ring.h)

    @property
    def DFB_ig_phasor(self):
        """Last direct feedback current generator phasor of each bunch in [A]"""
            if isinstance(FB, DirectFeedback):
                return FB.DFB_ig_phasor

    @property
    def cavity_voltage(self):
        """Cavity total voltage in [V]"""
        return np.abs(self.cavity_phasor)

    @property
    def cavity_phase(self):
        """Cavity total phase in [rad]"""
        return np.angle(self.cavity_phasor)

    @property
    def beam_voltage(self):
        """Beam loading voltage in [V]"""
        return np.abs(self.beam_phasor)

    @property
    def beam_phase(self):
        """Beam loading phase in [rad]"""
        return np.angle(self.beam_phasor)

    @property
    def loss_factor(self):
        """Cavity loss factor in [V/C]"""
        return self.wr * self.Rs / (2 * self.Q)

    @property
    def m(self):
        """Harmonic number of the cavity"""
        return self._m

    @m.setter
    def m(self, value):
        self._m = value

    @property
    def Ncav(self):
        """Number of cavities"""
        return self._Ncav

    @Ncav.setter
    def Ncav(self, value):
        self._Ncav = value

    @property
    def Rs_per_cavity(self):
        """Shunt impedance of a single cavity in [Ohm], defined as 
        0.5*Vc*Vc/Pc."""
        return self._Rs_per_cavity

    @Rs_per_cavity.setter
    def Rs_per_cavity(self, value):
        self._Rs_per_cavity = value

    @property
    def Rs(self):
        """Shunt impedance [ohm]"""
        return self.Rs_per_cavity * self.Ncav

    @Rs.setter
    def Rs(self, value):
        self.Rs_per_cavity = value / self.Ncav

    @property
    def RL(self):
        """Loaded shunt impedance [ohm]"""
        return self.Rs / (1 + self.beta)

    @property
    def Q(self):
        """Quality factor"""
        return self._Q

    @Q.setter
    def Q(self, value):
        self._Q = value

    @property
    def QL(self):
        """Loaded quality factor"""
        return self._QL

    @QL.setter
    def QL(self, value):
        self._QL = value
        self._beta = self.Q / self.QL - 1
        self.update_feedback()

    @property
    def beta(self):
        """Coupling coefficient"""
        return self._beta

    @beta.setter
    def beta(self, value):
        self.QL = self.Q / (1+value)

    @property
    def detune(self):
        """Cavity detuning [Hz] - defined as (fr - m*f1)"""
        return self._detune

    @detune.setter
    def detune(self, value):
        self._detune = value
        self._fr = self.detune + self.m * self.ring.f1
        self._wr = self.fr * 2 * np.pi
        self._psi = np.arctan(self.QL * (self.fr / (self.m * self.ring.f1) -
                                         (self.m * self.ring.f1) / self.fr))

    @property
    def fr(self):
        """Resonance frequency of the cavity in [Hz]"""
        return self._fr

    @fr.setter
    def fr(self, value):
        self.detune = value - self.m * self.ring.f1

    @property
    def wr(self):
        """Angular resonance frequency in [Hz.rad]"""
        return self._wr

    @wr.setter
    def wr(self, value):
        self.detune = (value - self.m * self.ring.f1) * 2 * np.pi

    @property
    def psi(self):
        """Tuning angle in [rad]"""
        return self._psi

    @psi.setter
    def psi(self, value):
        delta = (self.ring.f1 * self.m * np.tan(value) /
                 self.QL)**2 + 4 * (self.ring.f1 * self.m)**2
        fr = (self.ring.f1 * self.m * np.tan(value) / self.QL +
              np.sqrt(delta)) / 2
        self.detune = fr - self.m * self.ring.f1

    @property
    def filling_time(self):
        """Cavity filling time in [s]"""
        return 2 * self.QL / self.wr

    @property
    def Pc(self):
        """Power dissipated in the cavity walls in [W]"""
        return self.Vc**2 / (2 * self.Rs)

    def Pb(self, I0):
        """
        Return power transmitted to the beam in [W] - near Eq. (4.2.3) in [1].

        I0 : float
            Beam current in [A].

        Returns
        -------
        float
            Power transmitted to the beam in [W].

        """
        return I0 * self.Vc * np.cos(self.theta)

    def Pr(self, I0):
        """
        Power reflected back to the generator in [W].

        I0 : float

        Returns
        -------
        float
            Power reflected back to the generator in [W].

        """
        return self.Pg - self.Pb(I0) - self.Pc

    def Vbr(self, I0):
        """

        I0 : float

        Returns
        -------
        float
            Beam voltage at resonance in [V].

        """
        return 2 * I0 * self.Rs / (1 + self.beta)

    def Vb(self, I0):
        """

        I0 : float

        Returns
        -------
        float
            Beam voltage in [V].

        """
        return self.Vbr(I0) * np.cos(self.psi)

    def Z(self, f):
        """Cavity impedance in [Ohm] for a given frequency f in [Hz]"""
        return self.RL / (1 + 1j * self.QL * (self.fr / f - f / self.fr))

    def set_optimal_detune(self, I0, F=1):
        """
        Set detuning to optimal conditions - second Eq. (4.2.1) in [1].

        I0 : float

        """
        self.psi = np.arctan(-self.Vbr(I0) * F / self.Vc * np.sin(self.theta))

    def set_optimal_coupling(self, I0):
        """
        Set coupling to optimal value - Eq. (4.2.3) in [1].

        I0 : float

        """
        self.beta = 1 + (2 * I0 * self.Rs * np.cos(self.theta) / self.Vc)

    def set_generator(self, I0):
        """
        Set generator parameters (Pg, Vgr, theta_gr, Vg and theta_g) for a 
        given current and set of parameters.

        I0 : float

        """

        # Generator power [W] - Eq. (4.1.2) [1] corrected with factor (1+beta)**2 instead of (1+beta**2)
        self.Pg = self.Vc**2 * (1 + self.beta)**2 / (
            2 * self.Rs * 4 * self.beta * np.cos(self.psi)**2) * (
                (np.cos(self.theta) + 2 * I0 * self.Rs /
                 (self.Vc * (1 + self.beta)) * np.cos(self.psi)**2)**2 +
                (np.sin(self.theta) + 2 * I0 * self.Rs /
                 (self.Vc *
                  (1 + self.beta)) * np.cos(self.psi) * np.sin(self.psi))**2)
        # Generator voltage at resonance [V] - Eq. (3.2.2) [1]
        self.Vgr = 2 * self.beta**(1 / 2) / (1 + self.beta) * (
            2 * self.Rs * self.Pg)**(1 / 2)
        # Generator phase at resonance [rad] - from Eq. (4.1.1)
        self.theta_gr = np.arctan(
            (self.Vc * np.sin(self.theta) +
             self.Vbr(I0) * np.cos(self.psi) * np.sin(self.psi)) /
            (self.Vc * np.cos(self.theta) +
             self.Vbr(I0) * np.cos(self.psi)**2)) - self.psi
        # Generator voltage [V]
        self.Vg = self.Vgr * np.cos(self.psi)
        # Generator phase [rad]
        self.theta_g = self.theta_gr + self.psi
        # Set generator_phasor_record
        self.generator_phasor_record = np.ones(
            self.ring.h) * self.generator_phasor

    def plot_phasor(self, I0):
        """
        Plot phasor diagram showing the vector addition of generator and beam 
        loading voltage.

        I0 : float
            
        Returns
        -------
        Figure.

        """

        def make_legend_arrow(legend, orig_handle, xdescent, ydescent, width,
                              height, fontsize):
            p = mpatches.FancyArrow(0,
                                    0.5 * height,
                                    width,
                                    0,
                                    length_includes_head=True,
                                    head_width=0.75 * height)
            return p

        ax = fig.add_subplot(111, polar=True)
        ax.set_rmax(
            max([1.2,
                 self.Vb(I0) / self.Vc * 1.2, self.Vg / self.Vc * 1.2]))
        arr1 = ax.arrow(self.theta,
                        0,
                        0,
                        1,
                        alpha=0.5,
                        width=0.015,
                        edgecolor='black',
                        lw=2)

        arr2 = ax.arrow(self.psi + np.pi,
                        0,
                        0,
                        self.Vb(I0) / self.Vc,
                        edgecolor='red',
                        lw=2)

        arr3 = ax.arrow(self.theta_g,
                        0,
                        0,
                        self.Vg / self.Vc,
                        edgecolor='blue',
                        lw=2)

        ax.set_rticks([])  # less radial ticks
        plt.legend([arr1, arr2, arr3], ['Vc', 'Vb', 'Vg'],
                   handler_map={
                       mpatches.FancyArrow:
                       HandlerPatch(patch_func=make_legend_arrow),
                   })

        return fig

    def is_CBI_stable(self,
                      I0,
                      synchrotron_tune=None,
                      modes=None,
                      bool_return=False):
        """
        Check Coupled-Bunch-Instability stability, 
        Effect of Direct RF feedback is not included.

        This method is a wraper around lcbi_growth_rate to caluclate the CBI 
        growth rate from the CavityResonator own impedance.
        
        See lcbi_growth_rate for details.

        I0 : float
        synchrotron_tune : float, optinal
            Fractional number of longitudinal tune.
            If None, synchrotron_tune is computed using self.Vc as total 
            voltage.
        modes : float or list of float, optional
            Coupled Bunch Instability mode numbers to consider.
            If not None, return growth_rates of the input coupled bunch modes.
        bool_return : bool, optional
            If True:
                - return True if the most unstable mode is stronger than 
                longitudinal damping time.
                - return False if it is not the case.
            Default is False.
            
        Returns
        -------
        growth_rate : float
            Maximum coupled bunch instability growth rate in [s-1].
        mu : int
            Coupled bunch mode number corresponding to the maximum coupled bunch 
            instability growth rate.

        """
        growth_rate, mu, growth_rates = lcbi_growth_rate(
            self.ring,
            I0,
            M=self.ring.h,
            Vrf=self.Vc,
            fr=self.fr,
            RL=self.RL,
            QL=self.QL)

        if modes is not None:
            growth_rates = growth_rates[modes]
            return growth_rates

        if bool_return:
            if growth_rate > 1 / self.ring.tau[2]:
                return True
            else:
                return False
        else:
            return growth_rate, mu

    def is_DC_Robinson_stable(self, I0):
        """
        Check DC Robinson stability - Eq. (6.1.1) [1]

        I0 : float

        Returns
        -------
        bool

        """
        return 2 * self.Vc * np.sin(self.theta) + self.Vbr(I0) * np.sin(
            2 * self.psi) > 0

    def plot_DC_Robinson_stability(self, detune_range=[-1e5, 1e5]):
        """

        detune_range : list or array, optional
            Range of tuning to plot in [Hz].

        Returns
        -------
        Figure.

        """
        old_detune = self.psi

        x = np.linspace(detune_range[0], detune_range[1], 1000)
        y = []
        for i in range(0, x.size):
            self.detune = x[i]
            y.append(-self.Vc * (1 + self.beta) /
                     (self.Rs * np.sin(2 * self.psi)) *
                     np.sin(self.theta))  # droite de stabilité

        ax = plt.gca()
        ax.plot(x, y)
        ax.set_xlabel("Detune [Hz]")
        ax.set_ylabel("Threshold current [A]")
        ax.set_title("DC Robinson stability limit")

        self.psi = old_detune

        return fig

    def VRF(self, z, I0, F=1, PHI=0):
        """Total RF voltage taking into account form factor amplitude F and form factor phase PHI"""
        return self.Vg * np.cos(self.ring.k1 * self.m * z +
                                self.theta_g) - self.Vb(I0) * F * np.cos(
                                    self.ring.k1 * self.m * z + self.psi - PHI)

    def dVRF(self, z, I0, F=1, PHI=0):
        """Return derivative of total RF voltage taking into account form factor amplitude F and form factor phase PHI"""
        return -1 * self.Vg * self.ring.k1 * self.m * np.sin(
            self.ring.k1 * self.m * z +
            self.theta_g) + self.Vb(I0) * F * self.ring.k1 * self.m * np.sin(
                self.ring.k1 * self.m * z + self.psi - PHI)

    def ddVRF(self, z, I0, F=1, PHI=0):
        """Return the second derivative of total RF voltage taking into account form factor amplitude F and form factor phase PHI"""
        return -1 * self.Vg * (self.ring.k1 * self.m)**2 * np.cos(
            self.ring.k1 * self.m * z + self.theta_g) + self.Vb(I0) * F * (
                self.ring.k1 * self.m)**2 * np.cos(self.ring.k1 * self.m * z +
                                                   self.psi - PHI)

    def deltaVRF(self, z, I0, F=1, PHI=0):
        """Return the generator voltage minus beam loading voltage taking into account form factor amplitude F and form factor phase PHI"""
            self.ring.k1 * self.m * z + self.theta_g) - self.Vb(I0) * F * (

    def update_feedback(self):
        """Force feedback update from current CavityResonator parameters."""
                FB.init_Ig2Vg_matrix()

    def sample_voltage(self, n_points=1e4, index=0):
        """

        n_points : int or float, optional
            Number of sample points. The default is 1e4.
        index : int, optional
            RF bucket number to sample. Be carful if index > 0 as no new beam 
            loading is not taken into account here.

        Returns
        -------
        pos : array of float
            Array of position from -T1/2 to T1/2.
        voltage_rec : array of float
            Recoring of the voltage.

        """
        # Init
        n_points = int(n_points)
        index = 0
        voltage_rec = np.zeros(n_points)
        pos = np.linspace(-self.ring.T1 / 2, self.ring.T1 / 2, n_points)
        DeltaT = self.ring.T1 / (n_points-1)

        # From t=0 of first non empty bunch to -T1/2
        self.phasor_decay(-self.ring.T1 / 2 + index * self.ring.T1,
                          ref_frame="beam")

        # Goes from (-T1/2) to (T1/2 + DeltaT) in n_points steps
        for i in range(n_points):
            phase = self.m * self.ring.omega1 * (
                pos[i] + self.ring.T1 * (index + self.ring.h * self.nturn))
            Vgene = np.real(self.generator_phasor_record[index] *
                            np.exp(1j * phase))
            Vbeam = np.real(self.beam_phasor)
            Vtot = Vgene + Vbeam
            voltage_rec[i] = Vtot
            self.phasor_decay(DeltaT, ref_frame="beam")

        # Get back to t=0
        self.phasor_decay(-DeltaT * n_points + self.ring.T1 / 2 -
                          index * self.ring.T1,

        return pos, voltage_rec

class ProportionalLoop():
    """
    Proportional feedback loop to control a CavityResonator amplitude and phase.
    
    Feedback setpoints are cav_res.Vc and cav_res.theta.
    
    The loop must be added to the CavityResonator object using:
        cav_res.feedback.append(loop)

    cav_res : CavityResonator
        The CavityResonator which is to be controlled.
    gain_A : float
        Amplitude (voltage) gain of the feedback.
    gain_P : float
        Phase gain of the feedback.
    delay : int
        Feedback delay in unit of turns.
        Must be supperior or equal to 1.

    """

    def __init__(self, ring, cav_res, gain_A, gain_P, delay):
        self.cav_res = cav_res
        self.gain_A = gain_A
        self.gain_P = gain_P
        self.delay = int(delay)
        if delay < 1:
            raise ValueError("delay must be >= 1.")
        self.volt_delay = np.ones(self.delay) * self.cav_res.Vc
        self.phase_delay = np.ones(self.delay) * self.cav_res.theta

    def track(self):
        """
        Tracking method for the amplitude and phase loop.

        Returns
        -------
        None.

        """
        diff_A = self.volt_delay[-1] - self.cav_res.Vc
        diff_P = self.phase_delay[-1] - self.cav_res.theta
        self.cav_res.Vg -= self.gain_A * diff_A
        self.cav_res.theta_g -= self.gain_P * diff_P
        self.cav_res.generator_phasor_record = np.ones(
            self.ring.h) * self.cav_res.generator_phasor
        self.volt_delay = np.roll(self.volt_delay, 1)
        self.phase_delay = np.roll(self.phase_delay, 1)
        self.volt_delay[0] = self.cav_res.cavity_voltage
        self.phase_delay[0] = self.cav_res.cavity_phase

class TunerLoop():
    """
    Cavity tuner loop used to control a CavityResonator tuning (psi or detune)
    in order to keep the phase between cavity and generator current constant.
    
    Only a proportional controller is assumed.
    
        Proportional gain of the tuner loop.
        If not specified, 0.01 is used.
    avering_period:
        Period during which the phase difference is monitored and averaged. 
        Then the feedback correction is applied every avering_period turn.
        Unit is turn number.
        A value longer than one synchrotron period (1/fs) is recommended.
        If not specified, 2-synchrotron period (2/fs) is used, although it is
        too fast compared to the actual situation.
    offset : float
        Tuning offset in [rad].
        
    """

    def __init__(self, ring, cav_res, gain=0.01, avering_period=0, offset=0):
        if avering_period == 0:
            fs = self.ring.synchrotron_tune(
                self.cav_res.Vc) * self.ring.f1 / self.ring.h
            avering_period = 2 / fs / self.ring.T0

        self.Pgain = gain
        self.offset = offset
        self.avering_period = int(avering_period)
        self.diff = 0
        self.count = 0

        """
        Tracking method for the tuner loop.
    
        Returns
        -------
        None.
    
        """
        if self.count == self.avering_period:
            diff = self.diff / self.avering_period - self.offset
            self.cav_res.psi -= diff * self.Pgain
            self.count = 0
            self.diff = 0
        else:
            self.diff += self.cav_res.cavity_phase - self.cav_res.theta_g + self.cav_res.psi
            self.count += 1

class ProportionalIntegralLoop():
    """
    Proportional Integral (PI) loop to control a CavityResonator amplitude and 
    phase via generator current (ig) [1,2].
    
    Feedback reference targets (setpoints) are cav_res.Vc and cav_res.theta.
    
    The ProportionalIntegralLoop should be initialized only after generator 
    parameters are set.
    
    When the reference target is changed, it might be needed to re-initialize 
    the feedforward constant by calling init_FFconst().

    cav_res : CavityResonator object
        CavityResonator in which the loop will be added. 
    gain : float list
        Proportional gain (Pgain) and integral gain (Igain) of the feedback 
        system in the form of a list [Pgain, Igain].
        In case of normal conducting cavity (QL~1e4), the Pgain of ~1.0 and 
        Igain of ~1e4(5) are usually used.
        In case of super conducting cavity (QL > 1e6), the Pgain of ~100 
        can be used.
        In a "bad" parameter set, unstable oscillation of the cavity voltage 
        can be caused. So, a parameter scan of the gain should be made.
        Igain * ring.T1 / dtau is Ki defined as a the coefficients for 
        integral part in of [1], where dtau is a clock period of the PI controller.
    sample_num : int
        Number of bunch over which the mean cavity voltage is computed.
        Units are in bucket numbers.
    every : int
        Sampling and clock period of the feedback controller
        Time interval between two cavity voltage monitoring and feedback.
    delay : int
        Loop delay of the PI feedback system. 
    IIR_cutoff : float, optinal
        Cutoff frequency of the IIR filter in [Hz].
        If 0, cutoff frequency is infinity.
        Default is 0.
    FF : bool, optional
        Boolean switch to use feedforward constant.
        True is recommended to prevent a cavity voltage drop in the beginning 
        of the tracking.
        In case of small Pgain (QL ~ 1e4), cavity voltage drop may cause 
        beam loss due to static Robinson.
        Default is True.
        
    Methods
    -------
    track()
        Tracking method for the Cavity PI control feedback.
    init_Ig2Vg_matrix()
        Initialize matrix for Ig2Vg_matrix.
    init_FFconst()
        Initialize feedforward constant.
    Ig2Vg_matrix()
        Return Vg from Ig using matrix formalism.
    Ig2Vg()
        Go from Ig to Vg and apply values.
    Vg2Ig(Vg)
        Return Ig from Vg (assuming constant Vg).
    IIR_init(cutoff)
        Initialization for the IIR filter.
    IIR(input)
        Return IIR filter output.
    IIRcutoff()
        Return IIR cutoff frequency in [Hz].

    Notes
    -----
    Assumption : delay >> every ~ sample_num
    
    Adjusting ig(Vg) parameter to keep the cavity voltage constant (to target values).
    The "track" method is
       1) Monitoring the cavity voltage phasor
       Mean cavity voltage value between specified bunch number (sample_num) is monitored
       with specified interval period (every).
       2) Changing the ig phasor
       The ig phasor is changed according the difference of the specified
       reference values with specified gain (gain).
       By using ig instead of Vg, the cavity response can be taken account.
       3) ig changes are reflected to Vg after the specifed delay (delay) of the system

    Vc-->(rot)-->IIR-->(-)-->(V->I,fac)-->PI-->Ig --> Vg
                        |
                       Ref

    Examples
    --------
    PF; E0=2.5GeV, C=187m, frf=500MHz
        QL=11800, fs0=23kHz
        ==> gain=[0.5,1e4], sample_num=8, every=7(13ns), delay=500(1us)
                     is one reasonable parameter set.
        The practical clock period is 13ns.
            ==> Igain_PF = Igain_mbtrack * Trf / 13ns = Igain_mbtrack * 0.153
     
    [1] : https://en.wikipedia.org/wiki/PID_controller
    [2] : Yamamoto, N., Takahashi, T., & Sakanaka, S. (2018). Reduction and 
    compensation of the transient beam loading effect in a double rf system 
    of synchrotron light sources. PRAB, 21(1), 012001.
    
    """

                 ring,
                 cav_res,
                 gain,
                 sample_num,
                 every,
                 delay,
                 IIR_cutoff=0,
                 FF=True):
        self.Ig2Vg_mat = np.zeros((self.ring.h, self.ring.h), dtype=complex)
        self.ig_modulation_signal = np.zeros(self.ring.h, dtype=complex)
        self.FF = FF

        if delay > 0:
            self.delay = int(delay)
        else:
            self.delay = 1
        if every > 0:
            self.every = int(every)
        else:
            self.every = 1
        record_size = int(np.ceil(self.delay / self.every))
        if record_size < 1:
            raise ValueError("Bad parameter set : delay or every")
        self.sample_num = int(sample_num)

        # init lists for FB process
        self.ig_phasor = np.ones(self.ring.h, dtype=complex) * self.Vg2Ig(
            self.cav_res.generator_phasor)
        self.ig_phasor_record = self.ig_phasor
        self.vc_previous = np.ones(
            self.sample_num) * self.cav_res.cavity_phasor
        self.diff_record = np.zeros(record_size, dtype=complex)
        self.I_record = 0 + 0j

        self.sample_list = range(0, self.ring.h, self.every)

        self.IIR_init(IIR_cutoff)
        self.init_FFconst()

        # Pre caclulation for Ig2Vg
        self.init_Ig2Vg_matrix()

    def track(self, apply_changes=True):
        """

        Returns
        -------
        None.

        """
        vc_list = np.concatenate([
            self.vc_previous, self.cav_res.cavity_phasor_record
        ])  #This line is slowing down the process.
        self.ig_phasor.fill(self.ig_phasor[-1])

        for index in self.sample_list:
            # 2) updating Ig using last item of the list
            diff = self.diff_record[-1] - self.FFconst
            self.I_record += diff / self.ring.f1
            fb_value = self.gain[0] * diff + self.gain[1] * self.I_record
            self.ig_phasor[index:] = self.Vg2Ig(fb_value) + self.FFconst
            # Shift the record
            self.diff_record = np.roll(self.diff_record, 1)
            # 1) recording diff as a first item of the list
            mean_vc = np.mean(vc_list[index:self.sample_num + index]) * np.exp(
                -1j * self.cav_res.theta)
            self.diff_record[0] = self.cav_res.Vc - self.IIR(mean_vc)
        # update sample_list for next turn
        self.sample_list = range(index + self.every - self.ring.h, self.ring.h,
                                 self.every)
        # update vc_previous for next turn
        self.vc_previous = self.cav_res.cavity_phasor_record[-self.sample_num:]

        self.ig_phasor = self.ig_phasor + self.ig_modulation_signal

        if apply_changes:
            self.Ig2Vg()

    def init_Ig2Vg_matrix(self):
        """
        
        Shoud be called before first use of Ig2Vg_matrix and after each cavity 
        parameter change.
        """
        k = np.arange(0, self.ring.h)
        self.Ig2Vg_vec = np.exp(-1 / self.cav_res.filling_time *
                                (1 - 1j * np.tan(self.cav_res.psi)) *
                                self.ring.T1 * (k+1))
        tempV = np.exp(-1 / self.cav_res.filling_time * self.ring.T1 * k *
                       (1 - 1j * np.tan(self.cav_res.psi)))
        for idx in np.arange(self.ring.h):
            self.Ig2Vg_mat[idx:, idx] = tempV[:self.ring.h - idx]

    def init_FFconst(self):
        """Initialize feedforward constant."""
        if self.FF:
            self.FFconst = np.mean(self.ig_phasor)
        else:
            self.FFconst = 0

    def Ig2Vg_matrix(self):
        """
        Warning: self.init_Ig2Vg should be called after each CavityResonator 
        """
        generator_phasor_record = (
            self.Ig2Vg_vec * self.cav_res.generator_phasor_record[-1] +
            self.Ig2Vg_mat.dot(self.ig_phasor_record) *
            self.cav_res.loss_factor * self.ring.T1)
        return generator_phasor_record

    def Ig2Vg(self):
        """
        Go from Ig to Vg.
        
        Apply new values to cav_res.generator_phasor_record, cav_res.Vg and
        cav_res.theta_g from ig_phasor_record.
        """
        self.cav_res.generator_phasor_record = self.Ig2Vg_matrix()
        self.cav_res.Vg = np.mean(np.abs(self.cav_res.generator_phasor_record))
        self.cav_res.theta_g = np.mean(
            np.angle(self.cav_res.generator_phasor_record))

    def Vg2Ig(self, Vg):
        """

        Eq.25 of ref [2] assuming the dVg/dt = 0.
        """
        return Vg * (1 - 1j * np.tan(self.cav_res.psi)) / self.cav_res.RL

    def IIR_init(self, cutoff):
        """

        cutoff : float
            Cutoff frequency of the IIR filter in [Hz].
            If 0, cutoff frequency is infinity.
            
        """
        if cutoff == 0:
            self.IIRcoef = 1.0
        else:
            omega = 2.0 * np.pi * cutoff
            T = self.ring.T1 * self.every
            alpha = np.cos(omega * T) - 1
            tmp = alpha*alpha - 2*alpha
            if tmp > 0:
                self.IIRcoef = alpha + np.sqrt(tmp)
            else:
                self.IIRcoef = T * cutoff * 2 * np.pi
        self.IIRout = self.cav_res.Vc

    def IIR(self, input):
        """Return IIR filter output."""
        self.IIRout = (1 - self.IIRcoef) * self.IIRout + self.IIRcoef * input
        return self.IIRout

    @property
    def IIRcutoff(self):
        """Return IIR cutoff frequency in [Hz]."""
        T = self.ring.T1 * self.every
        return 1.0 / 2.0 / np.pi / T * np.arccos(
            (2 - 2 * self.IIRcoef - self.IIRcoef * self.IIRcoef) / 2 /
            (1 - self.IIRcoef))

class DirectFeedback(ProportionalIntegralLoop):
    """
    Direct RF FB (DFB) via generator current (ig) using PI controller [1,2].
    
    The DirectFeedback inherits from ProportionalIntegralLoop so all 
    mandatory parameters from ProportionalIntegralLoop should be passed as
    kwargs when initializing a DirectFeedback object.
        
    To avoid cavity-beam unmatching (large synchrotron oscilation of beam),
    the CavityResonator generator parameters should be set before 
    initialization.
    
    DFB_gain : float
        Gain of the DFB.
    DFB_phase_shift : float
        Phase shift of the DFB.
    DFB_sample_num : int, optional
        Sample number to monitor Vc for the DFB.
        The averaged Vc value in DFB_sample_num is monitored.
        If None, value from the ProportionalIntegralLoop is used.
    DFB_every : int, optional
        Interval to monitor and change Vc for the DFB.
    DFB_delay : int, optional
        Loop delay of the DFB.
        
    phase_shift : float
        Phase shift applied to DFB, defined as psi - DFB_phase_shift.
    DFB_psi: float
        Return detuning angle value with Direct RF feedback in [rad].
    DFB_alpha : float
        Return the amplitude ratio alpha of the DFB.
    DFB_gamma : float
        Return the gain gamma of the DFB.
    DFB_Rs : float
        Return the shunt impedance of the DFB in [ohm].
    
    Methods
    -------
    DFB_parameter_set(DFB_gain, DFB_phase_shift)
        Set DFB gain and phase shift parameters.
    track()
        Tracking method for the Direct RF feedback.
    DFB_Vg()
        Return the generator voltage main and DFB components in [V].
    DFB_fs()
        Return the modified synchrotron frequency in [Hz].
        
    [1] : Akai, K. (2022). Stability analysis of rf accelerating mode with 
    feedback loops under heavy beam loading in SuperKEKB. PRAB, 25(10), 
    102002.
    [2] : N. Yamamoto et al. (2023). Stability survey of a double RF system 
    with RF feedback loops for bunch lengthening in a low-emittance synchrotron 
    ring. In Proc. IPAC'23. doi:10.18429/JACoW-IPAC2023-WEPL161

    """

                 DFB_gain,
                 DFB_phase_shift,
                 DFB_sample_num=None,
                 DFB_every=None,
                 DFB_delay=None,
                 **kwargs):
        super(DirectFeedback, self).__init__(**kwargs)

        if DFB_delay is not None:
            self.DFB_delay = int(DFB_delay)
        else:
            self.DFB_delay = self.delay

        if DFB_sample_num is not None:
            self.DFB_sample_num = int(DFB_sample_num)
        else:
            self.DFB_sample_num = self.sample_num

        if DFB_every is not None:
            self.DFB_every = int(DFB_every)
        else:
            self.DFB_every = self.every

        record_size = int(np.ceil(self.DFB_delay / self.DFB_every))
            raise ValueError("Bad parameter set : DFB_delay or DFB_every")

        self.DFB_parameter_set(DFB_gain, DFB_phase_shift)
        if np.sum(np.abs(self.cav_res.beam_phasor)) == 0:
            cavity_phasor = self.cav_res.Vc * np.exp(1j * self.cav_res.theta)
        else:
            cavity_phasor = np.mean(self.cav_res.cavity_phasor_record)
        self.DFB_VcRecord = np.ones(record_size, dtype=complex) * cavity_phasor
        self.DFB_vc_previous = np.ones(self.DFB_sample_num,
                                       dtype=complex) * cavity_phasor

        self.DFB_sample_list = range(0, self.ring.h, self.DFB_every)

    @property
    def DFB_phase_shift(self):
        """Return DFB phase shift."""
        return self._DFB_phase_shift

    @DFB_phase_shift.setter
    def DFB_phase_shift(self, value):
        """Set DFB_phase_shift and phase_shift"""
        self._DFB_phase_shift = value
        self._phase_shift = self.cav_res.psi - value

    @property
    def phase_shift(self):
        """
        Return phase shift applied to DFB.
        Defined as self.cav_res.psi - self.DFB_phase_shift.
        """
        return self._phase_shift

    @property
    def DFB_psi(self):
        """
    
        Fig.4 of ref [1].
        """
        return (np.angle(np.mean(self.cav_res.cavity_phasor_record)) -
                np.angle(np.mean(self.ig_phasor_record)))

    @property
    def DFB_alpha(self):
        """
        
        Near Eq. (13) of [1].
        """
        fac = np.abs(
            np.mean(self.DFB_ig_phasor) / np.mean(self.ig_phasor_record))
        return 20 * np.log10(fac)

    @property
    def DFB_gamma(self):
        """
        
        """
        return fac / (1-fac)

    @property
    def DFB_Rs(self):
        """
        
        Eq. (15) of [1].
        """
        return self.cav_res.Rs / (1 + self.DFB_gamma * np.cos(self.DFB_psi))

    def DFB_parameter_set(self, DFB_gain, DFB_phase_shift):
        """

        DFB_gain : float
            Gain of the DFB.
        DFB_phase_shift : float
            Phase shift of the DFB.

        """
        self.DFB_gain = DFB_gain
        self.DFB_phase_shift = DFB_phase_shift

            vc = np.ones(self.ring.h) * self.cav_res.Vc * np.exp(
                1j * self.cav_res.theta)
        else:
            vc = self.cav_res.cavity_phasor_record
        vg_drf = self.DFB_gain * vc * np.exp(1j * self.phase_shift)
        self.DFB_ig_phasor = self.Vg2Ig(vg_drf)

        self.ig_phasor = self.ig_phasor_record - self.DFB_ig_phasor

        """

        Returns
        -------
        None.

        """
        super(DirectFeedback, self).track(False)

        vc_list = np.concatenate(
            [self.DFB_vc_previous, self.cav_res.cavity_phasor_record])
        self.DFB_ig_phasor = np.roll(self.DFB_ig_phasor, 1)
        for index in self.DFB_sample_list:
            vg_drf = self.DFB_gain * self.DFB_VcRecord[-1] * np.exp(
                1j * self.phase_shift)
            self.DFB_ig_phasor[index:] = self.Vg2Ig(vg_drf)
            self.DFB_VcRecord = np.roll(self.DFB_VcRecord, 1)
            # 1) recording Vc
            mean_vc = np.mean(vc_list[index:self.DFB_sample_num + index])
            self.DFB_VcRecord[0] = mean_vc
        self.DFB_sample_list = range(index + self.DFB_every - self.ring.h,
                                     self.ring.h, self.DFB_every)
        self.DFB_vc_previous = self.cav_res.cavity_phasor_record[
            -self.DFB_sample_num:]

        self.ig_phasor_record = self.ig_phasor + self.DFB_ig_phasor

        self.Ig2Vg()

    def DFB_Vg(self, vc=-1):
        """Return the generator voltage main and DFB components in [V]."""
        if vc == -1:
            vc = np.mean(self.cav_res.cavity_phasor_record)
        vg_main = np.mean(self.cav_res.generator_phasor_record) - vg_drf
        return vg_main, vg_drf

    def DFB_fs(self, vg_main=-1, vg_drf=-1):
        """Return the modified synchrotron frequency in [Hz]."""
        vc = np.mean(self.cav_res.cavity_phasor_record)
        if vg_drf == -1:
            vg_drf = self.DFB_gain * vc * np.exp(1j * self.phase_shift)
        if vg_main == -1:
            vg_main = np.mean(self.cav_res.generator_phasor_record) - vg_drf
        vg_sum = np.abs(vg_main) * np.sin(
            np.angle(vg_main)) + np.abs(vg_drf) * np.sin(np.angle(vg_drf))
        omega_s = 0
        if (vg_sum) > 0.0:
            omega_s = np.sqrt(self.ring.ac * self.ring.omega1 * (vg_sum) /
                              self.ring.E0 / self.ring.T0)
        return omega_s / 2 / np.pi

"""
Module where the Synchrotron class is defined.
"""

from scipy.constants import c, e

class Synchrotron:
    """
    Synchrotron class to store main properties.
    
    Optional parameters are optional only if the Optics object passed to the 
    class uses a loaded lattice.
    
    h : int
        Harmonic number of the accelerator.
    optics : Optics object
        Object where the optic functions are stored.
    particle : Particle object
        Particle which is accelerated.
    tau : array of shape (3,)
        Horizontal, vertical and longitudinal damping times in [s].
    sigma_delta : float
        Equilibrium energy spread.
    sigma_0 : float
        Natural bunch length in [s].
    emit : array of shape (2,)
        Horizontal and vertical equilibrium emittance in [m.rad].
    L : float, optional
        Ring circumference in [m].
    E0 : float, optional
        Nominal (total) energy of the ring in [eV].
    ac : float, optional
        Momentum compaction factor.
    tune : array of shape (2,), optional
        Horizontal and vertical tunes.
    chro : array of shape (2,), optional
        Horizontal and vertical (non-normalized) chromaticities.
    U0 : float, optional
        Energy loss per turn in [eV].
    adts : list of arrays or None, optional
        List that contains arrays of polynomial's coefficients, in decreasing 
        powers, used to determine Amplitude-Dependent Tune Shifts (ADTS). 
        The order of the elements strictly needs to be
        [coef_xx, coef_yx, coef_xy, coef_yy], where x and y denote the horizontal
        and the vertical plane, respectively, and coef_PQ means the polynomial's
        coefficients of the ADTS in plane P due to the amplitude in plane Q.

        For example, if the tune shift in y due to the Courant-Snyder invariant 
        Jx is characterized by the equation dQy(x) = 3*Jx**2 + 2*Jx + 1, then 
        coef_yx takes the form np.array([3, 2, 1]).
        
        Use None, to exclude the ADTS from calculation.
    mcf_order : array, optional
        Higher-orders momentum compaction factor in decreasing powers.
        Input here the coefficients considering only the derivatives and which
        do not include any factorial factor.
        The 1st order should be included in the array and be at the last 
        position.
        
    T0 : float
        Revolution time in [s].
    f0 : float
        Revolution frequency in [Hz].
    omega0 : float
        Angular revolution frequency in [Hz.rad]
    T1 : flaot
        Fundamental RF period in [s].
    f1 : float
        Fundamental RF frequency in [Hz].
    omega1 : float
        Fundamental RF angular frequency in [Hz.rad].
    k1 : float
        Fundamental RF wave number in [m**-1].
    gamma : float
        Relativistic Lorentz gamma.
        Relativistic Lorentz beta.
    long_alpha : float
        Longitudinal alpha Twiss parameter at the tracking location.
        Initialized at zero.
    long_beta : float
        Longitudinal beta Twiss parameter at the tracking location in [s].
    long_gamma : float
        Longitudinal gamma Twiss parameter at the tracking location in [s-1].
        
    Methods
    -------
    synchrotron_tune(V)
        Compute the (unperturbed) synchrotron tune from main RF voltage.
    sigma(position)
        Return the RMS beam size at equilibrium in [m].
    eta(delta)
        Momentum compaction taking into account higher orders if provided in
        mcf_order.
    mcf(delta)
        Momentum compaction factor taking into account higher orders if 
        provided in mcf_order.
    get_adts()
        Compute and add Amplitude-Dependent Tune Shifts (ADTS) sextupolar 
        componenet from AT lattice.
    get_mcf_order()
        Compute momentum compaction factor up to 3rd order from AT lattice.
    get_longitudinal_twiss(V)
        Compute the longitudinal Twiss parameters and the synchrotron tune for 
        single or multi-harmonic RF systems.
    to_pyat(Vrf)
        Return a pyAT simple_ring element from the Synchrotron element data.
    """

    def __init__(self, h, optics, particle, **kwargs):
        self._h = h
        self.particle = particle
        self.optics = optics
        self.long_alpha = 0
        self.long_beta = 0
        self.long_gamma = 0

        if self.optics.use_local_values == False:
            self.L = kwargs.get('L', self.optics.lattice.circumference)
            self.E0 = kwargs.get('E0', self.optics.lattice.energy)
            self.ac = kwargs.get('ac', self.optics.ac)
            self.tune = kwargs.get('tune', self.optics.tune)
            self.chro = kwargs.get('chro', self.optics.chro)
            self.U0 = kwargs.get('U0', self.optics.lattice.energy_loss)
        else:
            self.L = kwargs.get('L')  # Ring circumference [m]
            self.E0 = kwargs.get(
                'E0')  # Nominal (total) energy of the ring [eV]
            self.ac = kwargs.get('ac')  # Momentum compaction factor
            self.tune = kwargs.get('tune')  # X/Y/S tunes
            self.chro = kwargs.get(
                'chro')  # X/Y (non-normalized) chromaticities
            self.U0 = kwargs.get('U0')  # Energy loss per turn [eV]

        self.tau = kwargs.get('tau')  # X/Y/S damping times [s]
        self.sigma_delta = kwargs.get(
            'sigma_delta')  # Equilibrium energy spread
        self.sigma_0 = kwargs.get('sigma_0')  # Natural bunch length [s]
        self.emit = kwargs.get('emit')  # X/Y emittances in [m.rad]
        self.adts = kwargs.get('adts')  # Amplitude-Dependent Tune Shift (ADTS)
        self.mcf_order = kwargs.get(
            'mcf_order', self.ac)  # Higher-orders momentum compaction factor

    @property
    def h(self):
        """Harmonic number"""
        return self._h

    @h.setter
    def h(self, value):
        self._h = value
        self.L = self.L  # call setter

    @property
    def L(self):
        """Ring circumference [m]"""
        return self._L

    @L.setter
    def L(self, value):
        self._L = value
        self._T0 = self.L / c
        self._T1 = self.T0 / self.h
        self._f0 = 1 / self.T0
        self._omega0 = 2 * np.pi * self.f0
        self._f1 = self.h * self.f0
        self._omega1 = 2 * np.pi * self.f1
        self._k1 = self.omega1 / c

    @property
    def T0(self):
        """Revolution time [s]"""
        return self._T0

    @T0.setter
    def T0(self, value):
        self.L = c * value

    @property
    def T1(self):
        """"Fundamental RF period [s]"""
        return self._T1

    @T1.setter
    def T1(self, value):
        self.L = c * value * self.h

    @property
    def f0(self):
        """Revolution frequency [Hz]"""
        return self._f0

    @f0.setter
    def f0(self, value):
        self.L = c / value

    @property
    def omega0(self):
        """Angular revolution frequency [Hz rad]"""
        return self._omega0

    @omega0.setter
    def omega0(self, value):
        self.L = 2 * np.pi * c / value

    @property
    def f1(self):
        """Fundamental RF frequency [Hz]"""
        return self._f1

    @f1.setter
    def f1(self, value):
        self.L = self.h * c / value

    @property
    def omega1(self):
        """Fundamental RF angular frequency[Hz rad]"""
        return self._omega1

    @omega1.setter
    def omega1(self, value):
        self.L = 2 * np.pi * self.h * c / value

    @property
    def k1(self):
        """Fundamental RF wave number [m**-1]"""
        return self._k1

    @k1.setter
    def k1(self, value):
        self.L = 2 * np.pi * self.h / value

    @property
    def gamma(self):
        """Relativistic gamma"""
        return self._gamma

    @gamma.setter
    def gamma(self, value):
        self._gamma = value
        self._beta = np.sqrt(1 - self.gamma**-2)
        self._E0 = self.gamma * self.particle.mass * c**2 / e

    @property
        """Relativistic beta"""

        self.gamma = 1 / np.sqrt(1 - value**2)

    @property
    def E0(self):
        """Nominal (total) energy of the ring [eV]"""
        return self._E0

    @E0.setter
    def E0(self, value):
        self.gamma = value / (self.particle.mass * c**2 / e)

    @property
    def mcf_order(self):
        """Higher-orders momentum compaction factor"""
        return self._mcf_order

    @mcf_order.setter
    def mcf_order(self, value):
        self._mcf_order = value
        self.mcf = np.poly1d(self.mcf_order)

    def eta(self, delta=0):
        """

        delta : float or array like bunch["delta"], optional
            Relative energy deviation. The default is 0.

        Returns
        -------
        float or array
            Momentum compaction.

        """
        return self.mcf(delta) - 1 / (self.gamma**2)

    def sigma(self, position=None):
        """

        position : float or array, optional
            Longitudinal position in [m] where the beam size is computed. 
            If None, the local values are used.

        Returns
        -------
        sigma : array
            RMS beam size in [m] at position location or at local positon if 
            position is None.

        """
        if position is None:
            sigma = np.zeros((4, ))
            sigma[0] = (
                self.emit[0] * self.optics.local_beta[0] +
                self.optics.local_dispersion[0]**2 * self.sigma_delta**2)**0.5
            sigma[1] = (
                self.emit[0] * self.optics.local_gamma[0] +
                self.optics.local_dispersion[1]**2 * self.sigma_delta**2)**0.5
            sigma[2] = (
                self.emit[1] * self.optics.local_beta[1] +
                self.optics.local_dispersion[2]**2 * self.sigma_delta**2)**0.5
            sigma[3] = (
                self.emit[1] * self.optics.local_gamma[1] +
                self.optics.local_dispersion[3]**2 * self.sigma_delta**2)**0.5
        else:
            if isinstance(position, (float, int)):
                n = 1
            else:
                n = len(position)
            sigma = np.zeros((4, n))
            sigma[0, :] = (self.emit[0] * self.optics.beta(position)[0] +
                           self.optics.dispersion(position)[0]**2 *
                           self.sigma_delta**2)**0.5
            sigma[1, :] = (self.emit[0] * self.optics.gamma(position)[0] +
                           self.optics.dispersion(position)[1]**2 *
            sigma[2, :] = (self.emit[1] * self.optics.beta(position)[1] +
                           self.optics.dispersion(position)[2]**2 *
            sigma[3, :] = (self.emit[1] * self.optics.gamma(position)[1] +
                           self.optics.dispersion(position)[3]**2 *
        return sigma

    def synchrotron_tune(self, V):
        """
        
        V : float
            Main RF voltage in [V].

        Returns
        -------
            
        """
        Vsum = V * np.sin(np.arccos(self.U0 / V))
        phi = np.arccos(1 - self.eta(0) * np.pi * self.h / self.E0 * Vsum)
        tuneS = phi / (2 * np.pi)
        return tuneS

    def get_adts(self):
        """
        """
        import at
        if self.optics.use_local_values:
            raise ValueError("ADTS needs to be provided manualy as no AT" +
                             " lattice file is loaded.")

        det = at.physics.nonlinear.gen_detuning_elem(self.optics.lattice)
        coef_xx = np.array([det.A1 / 2, 0])
        coef_yx = np.array([det.A2 / 2, 0])
        coef_xy = np.array([det.A2 / 2, 0])
        coef_yy = np.array([det.A3 / 2, 0])
        self.adts = [coef_xx, coef_yx, coef_xy, coef_yy]

    def get_mcf_order(self, add=True, show_fit=False):
        """
        
        add : bool, optional
            If True, add the computed coefficient to the Synchrotron object.
            If False, the result is returned.
        show_fit : bool, optional
            If True, show a plot with the data and fit.
            The default is False.
            
        Return
        ------
        pvalue : array, optional
            Momentum compaction factor up to 3rd order
        """
        import at
        deltamin = -1e-4
        deltamax = 1e-4

        delta = np.linspace(deltamin, deltamax, 7)
        delta4eval = np.linspace(deltamin, deltamax, 21)
        alpha = np.zeros_like(delta)

        for i in range(len(delta)):
            alpha[i] = at.physics.revolution.get_mcf(self.optics.lattice,
                                                     delta[i])
        pvalue = np.polyfit(delta, alpha, 2)

        if show_fit:
            pvalue = np.polyfit(delta, alpha, 2)
            palpha = np.polyval(pvalue, delta4eval)

            plt.plot(delta * 100, alpha, 'k.')
            plt.grid()
            plt.plot(delta4eval * 100, palpha, 'r')
            plt.xlabel('Energy (%)')
            plt.ylabel('Momemtum compaction factor')
            plt.legend(['Data', 'polyfit'])

        if add:
            self.mcf_order = pvalue
        else:
            return pvalue

    def get_longitudinal_twiss(self, V, phase=None, harmonics=None, add=True):
        """
        
        Hypthesis:
        - Use the linear approximation (tau ~ 0).
        - For higher haromics only, cos(phi) ~ 0 must be true.
        
        V : float or array-like of float
            RF voltage in [V].
        phase : float or array-like of float, optional
            RF phase using cosine convention in [rad].
        harmonics : int or array-like of int, optional
            Harmonics of the cavities to consider.
            If True, 
            
        Usage
        -----
        - If a single voltage value is given, assume a single RF system set at 
        the synchronous phase.
        - If a single pair of voltage/phase value is given, assume a single RF 
        system at this setting.
        - Otherwise, compute the synchrotron tune for a multi-harmonic RF 
        system.
            
        Returns
        -------
            Synchrotron tune in the linear approximation.
        long_alpha : float
            Longitudinal alpha Twiss parameter at the tracking location.
        long_beta : float
            Longitudinal beta Twiss parameter at the tracking location in [s].   
        long_gamma : float
            Longitudinal gamma Twiss parameter at the tracking location in [s-1].
        """

        if isinstance(V, float) or isinstance(V, int):
            V = [V]
            if phase is None:
                phase = [np.arccos(self.U0 / V[0])]
            elif isinstance(phase, float) or isinstance(phase, int):
                phase = [phase]
            if harmonics is None:
                harmonics = [1]

        if not (len(V) == len(phase) == len(harmonics)):
            raise ValueError("You must provide array of the same length for"
                             " V, phase and harmonics")

        Vsum = 0
        for i in range(len(V)):
            Vsum += harmonics[i] * V[i] * np.sin(phase[i])
        long_alpha = -self.eta(0) * np.pi * self.h / (self.E0 *
                                                      np.sin(phi)) * Vsum
        long_beta = self.eta(0) * self.T0 / np.sin(phi)
        long_gamma = self.omega1 * Vsum / (self.E0 * np.sin(phi))
        if add:
            self.tuneS = tuneS
            self.long_alpha = long_alpha
            self.long_beta = long_beta
            self.long_gamma = long_gamma
        else:
            return tuneS, long_alpha, long_beta, long_gamma

    def to_pyat(self, Vrf, harmonic_number=None, TimeLag=False):
        """
        
        See pyAT documentation for informations about simple_ring.

        Vrf : float or array-like of float
            RF Voltage in [V]. If sevral cavities are provided, harmonic_number
            and TimeLag should be provided.
        harmonic_number : float or array-like of float, optional
            Harmonic numbers of the RF cavities. The default is None.
        TimeLag : float or array-like of float, optional
            Set the timelag of the cavities in pyAT definition. 

        Returns
        -------
        at_simple_ring : at.physics.fastring.simple_ring
            A pyAT simple_ring element.

        """
        from at import simple_ring
        optics = self.optics

        if (harmonic_number is None) and isinstance(Vrf, (float, int)):
            harmonic_number = self.h

        if isinstance(Vrf, (list, np.ndarray)):
            if (harmonic_number is None) or (TimeLag is None):
                raise ValueError(
                    "If sevral cavities are provided, "
                    "harmonic_number and TimeLag should be provided.")

        if self.adts is not None:
            try:
                A1 = self.adts[0][-2] * 2
                A2 = self.adts[1][-2] * 2
                A3 = self.adts[3][-2] * 2
            except IndexError:
                A1 = None
                A2 = None
                A3 = None
        else:
            A1 = None
            A2 = None
            A3 = None

        at_simple_ring = simple_ring(energy=self.E0,
                                     circumference=self.L,
                                     harmonic_number=harmonic_number,
                                     Qx=self.tune[0],
                                     Qy=self.tune[1],
                                     Vrf=Vrf,
                                     alpha=self.ac,
                                     betax=optics.local_beta[0],
                                     betay=optics.local_beta[1],
                                     alphax=optics.local_alpha[0],
                                     alphay=optics.local_alpha[1],
                                     Qpx=self.chro[0],
                                     Qpy=self.chro[1],
                                     A1=A1,
                                     A2=A2,
                                     A3=A3,
                                     emitx=self.emit[0],
                                     emity=self.emit[1],
                                     espread=self.sigma_delta,
                                     taux=self.tau[0] / self.T0,
                                     tauy=self.tau[1] / self.T0,
                                     tauz=self.tau[2] / self.T0,
                                     U0=self.U0,
                                     TimeLag=TimeLag)
        return at_simple_ring

"""
This module defines the WakePotential and LongRangeResistiveWall classes which 
deal with the single bunch and multi-bunch wakes.
"""

from scipy import signal

from mbtrack2.utilities.spectrum import gaussian_bunch

class WakePotential(Element):
    """
    Compute a wake potential from uniformly sampled wake functions by 
    performing a convolution with a bunch charge profile.
    
    Two different time bases are used. The first one is controled by the n_bin
    parameter and is used to compute the bunch profile. Then the bunch profile
    is interpolated on the wake function time base which is used to perform the
    convolution to get the wake potential.
    
    wakefield : Wakefield object
        Wakefield object which contains the wake functions to be used. The wake
        functions must be uniformly sampled!
        Number of bins for constructing the longitudinal bunch profile.
    interp_on_postion : bool, optional
        If True, the computed wake potential is interpolated on the exact 
        particle location. If False, the wake potential is interpolated on the 
        bin center and each particle of the bin get the same value.
        
    rho : array of shape (n_bin, )
        Bunch charge density profile in the unit [1/s].
    Wp : array 
        Wake potential profile.
    Wp_interp : array of shape (mp_number, )
        Wake potential, obtained from interpolating Wp, exerted on each macro-particle.
    
    Methods
    -------
    charge_density(bunch)
        Compute bunch charge density profile in [1/s].
    dipole_moment(bunch, plane, tau0)
        Return the dipole moment of the bunch computed on the same time array 
        as the wake function.
    prepare_wakefunction(wake_type)
        Prepare the wake function of a given wake_type to be used for the wake
        potential computation. 
    get_wakepotential(bunch, wake_type)
        Return the wake potential computed on the wake function time array 
        limited to the bunch profile.
    track(bunch)
    plot_last_wake(wake_type)
        Plot the last wake potential of a given type computed during the last
        call of the track method.
    reference_loss(bunch)
        Calculate the loss factor and kick factor from the wake potential and 
        compare it to a reference value assuming a Gaussian bunch computed in 
        the frequency domain.
    check_sampling()
        Check if the wake function sampling is uniform.
    reduce_sampling(factor)
        Reduce wake function samping by an integer factor.
        
    """

    def __init__(self, ring, wakefield, n_bin=80, interp_on_postion=True):
        self.wakefield = wakefield
        self.types = self.wakefield.wake_components
        self.n_types = len(self.wakefield.wake_components)
        self.n_bin = n_bin
        self.check_sampling()
        self.interp_on_postion = interp_on_postion

        # Suppress numpy warning for floating-point operations.
        np.seterr(invalid='ignore')

    def charge_density(self, bunch):
        """

        bunch : Bunch object

        """

        # Get binning data
        a, b, c, d = bunch.binning(n_bin=self.n_bin)
        self.bins = a
        self.sorted_index = b
        self.profile = c
        self.center = d
        self.bin_size = self.bins[1] - self.bins[0]

        # Compute charge density
        self.rho = bunch.charge_per_mp * self.profile / (self.bin_size *
                                                         bunch.charge)
        self.rho = np.array(self.rho)

        # Compute time array
        self.tau = np.array(self.center)
        self.dtau = self.tau[1] - self.tau[0]

        # Add N values before and after rho and tau
        N = int(self.n_bin / 2)
        self.tau = np.arange(self.tau[0] - self.dtau * N,
                             self.tau[-1] + self.dtau * N, self.dtau)
        self.rho = np.pad(self.rho, (N, N + self.n_bin % 2),
                          mode="constant",
                          constant_values=(0, 0))

        if len(self.tau) != len(self.rho):
            self.tau = np.append(self.tau, self.tau[-1] + self.dtau)

        self.tau_mean = np.mean(self.tau)
        self.tau -= self.tau_mean

    def dipole_moment(self, bunch, plane, tau0):
        """

        plane : str
            Plane on which the dipole moment is computed, "x" or "y".
        tau0 : array
            Time array on which the dipole moment will be interpolated, in [s].

        Returns
        -------
        dipole : array
            Dipole moment of the bunch.

        """
        dipole = np.empty((self.n_bin - 1, ))
        for i in range(self.n_bin - 1):
            dipole[i] = bunch[plane][self.sorted_index == i].sum()
        dipole = dipole / self.profile
        dipole[np.isnan(dipole)] = 0

        # Add N values to get same size as tau/profile
        if self.n_bin % 2 == 0:
            N = int(self.n_bin / 2)
            dipole = np.append(dipole, np.zeros(N))
            dipole = np.insert(dipole, 0, np.zeros(N))
        else:
            N = int(np.floor(self.n_bin / 2))
            dipole = np.insert(dipole, 0, np.zeros(N + 1))

        # Interpole on tau0 to get the same size as W0
        dipole0 = np.interp(tau0, self.tau, dipole, 0, 0)

        setattr(self, "dipole_" + plane, dipole0)
        return dipole0

    def prepare_wakefunction(self, wake_type, tau, save_data=True):
        """
        
        The new time array keeps the same sampling time as given in the 
        WakeFunction definition but is restricted to the bunch profile time 
        array.

        wake_type : str
            Type of the wake function to prepare: "Wlong", "Wxdip", ...
        tau : array
            Time domain array of the bunch profile in [s].
        save_data : bool, optional
            If True, the results are saved as atributes.

        Returns
        -------
            Time base of the wake function in [s].
        dtau0 : float
            Difference between two points of the wake function time base in 
            [s].
        W0 : array
            Wake function array in [V/C] or [V/C/m].

        """

        tau0 = np.array(getattr(self.wakefield, wake_type).data.index)
        dtau0 = tau0[1] - tau0[0]
        W0 = np.array(getattr(self.wakefield, wake_type).data["real"])

        # Keep only the wake function on the rho window
        ind = np.all([min(tau[0], 0) < tau0, max(tau[-1], 0) > tau0], axis=0)
        tau0 = tau0[ind]
        W0 = W0[ind]

        # Check the wake function window for assymetry
        assym = (np.abs(tau0[-1]) - np.abs(tau0[0])) / dtau0
        n_assym = int(np.floor(assym))
        if np.floor(assym) > 1:

            # add at head
            if np.abs(tau0[-1]) > np.abs(tau0[0]):
                tau0 = np.arange(tau0[0] - dtau0*n_assym, tau0[-1] + dtau0,
                                 dtau0)
                n_to_add = len(tau0) - len(W0)
                W0 = np.concatenate((np.zeros(n_to_add), W0))

            # add at tail
            elif np.abs(tau0[0]) > np.abs(tau0[-1]):
                tau0 = np.arange(tau0[0], tau0[-1] + dtau0 * (n_assym+1),
                                 dtau0)

        # Check is the wf is shorter than rho then add zeros
        if (tau0[0] > tau[0]) or (tau0[-1] < tau[-1]):
            n = max(int(np.ceil((tau0[0] - tau[0]) / dtau0)),
                    int(np.ceil((tau[-1] - tau0[-1]) / dtau0)))

            tau0 = np.arange(tau0[0] - dtau0*n, tau0[-1] + dtau0 * (n+1),
                             dtau0)
            W0 = np.insert(W0, 0, np.zeros(n))
            n_to_add = len(tau0) - len(W0)
            W0 = np.pad(W0, (0, n_to_add),
                        mode="constant",
                        constant_values=(0, 0))

        if save_data:
            setattr(self, "tau0_" + wake_type, tau0)
            setattr(self, "dtau0_" + wake_type, dtau0)
            setattr(self, "W0_" + wake_type, W0)

        return (tau0, dtau0, W0)

    def get_wakepotential(self, bunch, wake_type):
        """

            Wake function type: "Wlong", "Wxdip", ...

        Returns
        -------
        Wp : array
            Wake potential.

        """

        (tau0, dtau0, W0) = self.prepare_wakefunction(wake_type, self.tau)

        profile0 = np.interp(tau0, self.tau, self.rho, 0, 0)

        if wake_type == "Wlong" or wake_type == "Wxquad" or wake_type == "Wyquad" or wake_type == "Wxcst" or wake_type == "Wycst":
            Wp = signal.convolve(profile0, W0 * -1, mode='same') * dtau0
        elif wake_type == "Wxdip":
            dipole0 = self.dipole_moment(bunch, "x", tau0)
            Wp = signal.convolve(profile0 * dipole0, W0, mode='same') * dtau0
        elif wake_type == "Wydip":
            dipole0 = self.dipole_moment(bunch, "y", tau0)
        else:
            raise ValueError("This type of wake is not taken into account.")

        setattr(self, "profile0_" + wake_type, profile0)
        setattr(self, wake_type, Wp)
        return tau0, Wp

        """

        bunch : Bunch or Beam object.
        
        """

            self.charge_density(bunch)
            for wake_type in self.types:
                tau0, Wp = self.get_wakepotential(bunch, wake_type)
                if self.interp_on_postion:
                    Wp_interp = np.interp(bunch["tau"], tau0 + self.tau_mean,
                                          Wp, 0, 0)
                else:
                    Wp_interp = np.interp(self.center, tau0 + self.tau_mean,
                                          Wp, 0, 0)
                    Wp_interp = Wp_interp[self.sorted_index]
                if wake_type == "Wlong":
                    bunch["delta"] += Wp_interp * bunch.charge / self.ring.E0
                elif wake_type == "Wxdip":
                    bunch["xp"] += Wp_interp * bunch.charge / self.ring.E0
                elif wake_type == "Wydip":
                    bunch["yp"] += Wp_interp * bunch.charge / self.ring.E0
                elif wake_type == "Wxquad":
                    bunch["xp"] += (bunch["x"] * Wp_interp * bunch.charge /
                                    self.ring.E0)
                elif wake_type == "Wyquad":
                    bunch["yp"] += (bunch["y"] * Wp_interp * bunch.charge /
                elif wake_type == "Wxcst":
                elif wake_type == "Wycst":

    def plot_last_wake(self,
                       wake_type,
                       plot_rho=True,
                       plot_dipole=False,
                       plot_wake_function=True):
        """

            Type of the wake to plot: "Wlong", "Wxdip", ...
        plot_rho : bool, optional
            Plot the normalised bunch profile. The default is True.
        plot_dipole : bool, optional
            Plot the normalised dipole moment. The default is False.
        plot_wake_function : bool, optional
            Plot the normalised wake function. The default is True.

        Returns
        -------
        fig : figure

        """

        labels = {
            "Wlong": r"$W_{p,long}$ (V/pC)",
            "Wxdip": r"$W_{p,x}^{D} (V/pC)$",
            "Wydip": r"$W_{p,y}^{D} (V/pC)$",
            "Wxquad": r"$W_{p,x}^{Q} (V/pC/m)$",
            "Wyquad": r"$W_{p,y}^{Q} (V/pC/m)$",
            "Wxcst": r"$W_{p,x}^{M}$ (V/pC)",
            "Wycst": r"$W_{p,y}^{M}$ (V/pC)",
        }

        Wp = getattr(self, wake_type)
        tau0 = getattr(self, "tau0_" + wake_type)

        ax.plot(tau0 * 1e12, Wp * 1e-12, label=labels[wake_type])
        ax.set_xlabel("$\\tau$ (ps)")
        ax.set_ylabel(labels[wake_type])

        if plot_rho is True:
            profile0 = getattr(self, "profile0_" + wake_type)
            profile_rescaled = profile0 / max(profile0) * max(np.abs(Wp))
            rho_rescaled = self.rho / max(self.rho) * max(np.abs(Wp))
            ax.plot(tau0 * 1e12,
                    profile_rescaled * 1e-12,
                    label=r"$\rho$ interpolated (a.u.)")
            ax.plot((self.tau + self.tau_mean) * 1e12,
                    rho_rescaled * 1e-12,
                    label=r"$\rho$ (a.u.)",
                    linestyle='dashed')
            plt.legend()

        if plot_wake_function is True:
            W0 = getattr(self, "W0_" + wake_type)
            W0_rescaled = W0 / max(W0) * max(np.abs(Wp))
                    W0_rescaled * 1e-12,
                    label=r"$W_{function}$ (a.u.)")

        if plot_dipole is True:
            dipole = getattr(self, "dipole_" + wake_type[1])
            dipole_rescaled = dipole / max(dipole) * max(np.abs(Wp))
                    dipole_rescaled * 1e-12,
                    label=r"Dipole moment (a.u.)")

        return fig

    def get_gaussian_wakepotential(self, sigma, wake_type, dipole=1e-3):
        """
        Return the wake potential computed using a perfect gaussian profile.

        dipole : float, optional
            Dipole moment to consider in [m], (uniform dipole moment).
        Returns
        -------
            Time base in [s].
        W0 : array
            Wake function.
        Wp : array
        profile0 : array
            Gaussian bunch profile.
        dipole0 : array
            Dipole moment.

        """

        tau = np.linspace(-10 * sigma, 10 * sigma, int(1e3))
        (tau0, dtau0, W0) = self.prepare_wakefunction(wake_type, tau, False)

        profile0 = gaussian_bunch(tau0, sigma)
        dipole0 = np.ones_like(profile0) * dipole

        if wake_type == "Wlong" or wake_type == "Wxquad" or wake_type == "Wyquad":
        else:

        return tau0, W0, Wp, profile0, dipole0

    def plot_gaussian_wake(self,
                           sigma,
                           wake_type,
                           dipole=1e-3,
                           plot_rho=True,
                           plot_dipole=False,
                           plot_wake_function=True):
        """
        Plot the wake potential of a given type for a perfect gaussian bunch.

        Returns
        -------

        """

        labels = {
            "Wyquad": r"$W_{p,y}^{Q} (V/pC/m)$"
        }

        tau0, W0, Wp, profile0, dipole0 = self.get_gaussian_wakepotential(
            sigma, wake_type, dipole)

                    label=r"$\rho$ (a.u.)")

            dipole_rescaled = dipole0 / max(dipole0) * max(np.abs(Wp))

        return fig

    def reference_loss(self, bunch):
        """

        Returns
        -------
        loss_data : DataFrame
            An output showing the loss/kick factors compared to the reference 
            values.

        """
        loss = []
        loss_0 = []
        delta_loss = []
        index = []
        for wake_type in self.types:
            tau0, Wp = self.get_wakepotential(bunch, wake_type)
            factorTD = np.trapz(Wp * profile0, tau0)

            if wake_type == "Wlong":
                factorTD *= -1
            if wake_type == "Wxdip":
                factorTD /= bunch["x"].mean()
            if wake_type == "Wydip":
                factorTD /= bunch["y"].mean()

            Z = getattr(self.wakefield, "Z" + wake_type[1:])
            sigma = bunch['tau'].std()
            factorFD = Z.loss_factor(sigma)

            loss.append(factorTD)
            loss_0.append(factorFD)
            delta_loss.append((factorTD-factorFD) / factorFD * 100)
                index.append("Wlong [V/C]")
            else:
                index.append(wake_type + " [V/C/m]")

            column = ['TD factor', 'FD factor', 'Relative error [%]']

        loss_data = pd.DataFrame(np.array([loss, loss_0, delta_loss]).T,
                                 columns=column,
                                 index=index)
        return loss_data

    def check_sampling(self):
        """

        Raises
        ------
        ValueError

        """
            idx = getattr(self.wakefield, wake_type).data.index
            diff = idx[1:] - idx[:-1]
            result = np.all(np.isclose(diff, diff[0], atol=1e-15))
            if result is False:
                    "The wake function must be uniformly sampled.")

    def reduce_sampling(self, factor):
        """
        
        Used to reduce computation time for long bunches.

        factor : int

        """
            idx = getattr(self.wakefield, wake_type).data.index[::factor]
            getattr(self.wakefield,
                    wake_type).data = getattr(self.wakefield,
                                              wake_type).data.loc[idx]

class LongRangeResistiveWall(Element):
    """
    Element to deal with multi-bunch and multi-turn wakes from resistive wall 
    using the algorithm defined in [1].
    
    Main approximations:
        - Bunches are treated as point charge.
        - Assymptotic expression for the resistive wall wake functions are 
        used.
        - Multi-turn wakes are limited to nt turns.
    
    Self-bunch interaction is not included in this element and should be dealed
    with the WakePotential class.
    
    beam : Beam object
    length : float
        Length of the resistive pipe to consider in [m].
    rho : float
        Effective resistivity to consider in [ohm.m] as in [1].
        Beam pipe radius to consider in [m].
    types : str or list, optional
        Wake types to consider. 
        The default is ["Wlong","Wxdip","Wydip"].
    nt : int or float, optional
        Number of turns to consider for the long range wakes. 
        The default is 50.
    x3 : float, optional
        Horizontal effective radius of the 3rd power in [m], as Eq.27 in [1].
        The default is radius.
    y3 : float, optional
        Vertical effective radius of the 3rd power in [m], as Eq.27 in [1].
    
    [1] : Skripka, Galina, et al. "Simultaneous computation of intrabunch and 
    interbunch collective beam motions in storage rings." NIM.A (2016).
    """

                 ring,
                 beam,
                 length,
                 rho,
                 radius,
                 types=["Wlong", "Wxdip", "Wydip"],
                 nt=50,
                 x3=None,
                 y3=None):
        # parameters
        self.nt = int(nt)
        self.nb = len(beam)
        if isinstance(types, str):
            self.types = [types]
        elif isinstance(types, list):
            self.types = types

        # effective radius for RW
        if x3 is not None:
            self.x3 = x3
        else:
            self.x3 = radius
        if y3 is not None:
            self.y3 = y3
        else:
            self.y3 = radius

        # constants

        # check approximation
        if 20 * self.t0 > ring.T1:
            raise ValueError("The approximated wake functions are not valid.")

        # init tables
        self.tau = np.ones((self.nb, self.nt)) * 1e100
        self.x = np.zeros((self.nb, self.nt))
        self.y = np.zeros((self.nb, self.nt))
        self.charge = np.zeros((self.nb, self.nt))

    def Wlong(self, t):
        """
        Approxmiate expression for the longitudinal resistive wall wake 
        function - Eq.24 of [1].

        t : float
            Time in [s].

        Returns
        -------
        wl : float
            Wake function in [V/C].

        """
        wl = (1 / (4 * pi * self.radius) * np.sqrt(self.Z0 * self.rho /
                                                   (c*pi)) /
              t**(3 / 2)) * self.length * -1
        return wl

    def Wdip(self, t, plane):
        """
        Approxmiate expression for the transverse resistive wall wake 
        function - Eq.26 of [1].

        t : float
        plane : str
            "x" or "y".

        Returns
        -------
        wdip : float
            Wake function in [V/C/m].

        """
        if plane == "x":
            r3 = self.x3
        elif plane == "y":
            r3 = self.y3
        else:
            raise ValueError()

        wdip = (1 / (pi * r3**3) * np.sqrt(self.Z0 * c * self.rho / pi) /
                t**(1 / 2) * self.length)
        return wdip

    def update_tables(self, beam):
        """
        Update tables.
        
        Table tau[i,j] is defined as the time difference of the bunch i center 
        of mass with respect to center of the RF bucket number 0 at turn j.
        Turn 0 corresponds to the tracked turn.
        
        Positive time corresponds to past events and negative time to future 
        events.

        Returns
        -------
        None.

        """
        # shift tables to next turn
        self.tau += self.ring.T0
        self.tau = np.roll(self.tau, shift=1, axis=1)
        self.x = np.roll(self.x, shift=1, axis=1)
        self.y = np.roll(self.y, shift=1, axis=1)
        self.charge = np.roll(self.charge, shift=1, axis=1)

        # update tables
            beam.mpi.share_means(beam)
            # negative sign as when bunch 0 is tracked, the others are not yet passed
            self.tau[:,
                     0] = beam.mpi.mean_all[:,
                                            4] - beam.bunch_index * self.ring.T1
            self.x[:, 0] = beam.mpi.mean_all[:, 0]
            self.y[:, 0] = beam.mpi.mean_all[:, 2]
            self.charge[:, 0] = beam.mpi.charge_all
        else:
            mean_all = beam.bunch_mean
            charge_all = beam.bunch_charge
            self.tau[:, 0] = mean_all[
                4, beam.filling_pattern] - beam.bunch_index * self.ring.T1
            self.x[:, 0] = mean_all[0, beam.filling_pattern]
            self.y[:, 0] = mean_all[2, beam.filling_pattern]
            self.charge[:, 0] = charge_all[beam.filling_pattern]

    def get_kick(self, rank, wake_type):
        """
        Compute the wake kick to apply.

        rank : int
            Rank of the bunch, as defined in Mpi class.
        wake_type : float
            Type of the wake to compute.

        Returns
        -------
        sum_kick : float
            Sum of the kicks from the different bunches at different turns.

        """
        sum_kick = 0
        for j in range(self.nt):
            for i in range(self.nb):
                if (j == 0) and (rank <= i):
                    continue
                deltaT = self.tau[i, j] - self.tau[rank, 0]
                    sum_kick += self.Wlong(deltaT) * self.charge[i, j]
                    sum_kick += self.Wdip(
                        deltaT, "x") * self.charge[i, j] * self.x[i, j]
                        deltaT, "y") * self.charge[i, j] * self.y[i, j]
                    raise NotImplementedError()

        return sum_kick

    def track_bunch(self, bunch, rank):
        """
        Track a bunch.
        
        Should only be used within the track method and not standalone.

        rank : int

        Returns
        -------
        None.

        """
            kick = self.get_kick(rank, wake_type)
                bunch["delta"] -= kick / self.ring.E0
            elif wake_type == "Wxdip":
                bunch["xp"] += kick / self.ring.E0
            elif wake_type == "Wydip":
                bunch["yp"] += kick / self.ring.E0
            elif wake_type == "Wxquad":
                bunch["xp"] += (bunch["x"] * kick / self.ring.E0)
            elif wake_type == "Wyquad":
                bunch["yp"] += (bunch["y"] * kick / self.ring.E0)

        """
        Track a beam.

        Returns
        -------
        None.

        """
        self.update_tables(beam)

            rank = beam.mpi.rank
            bunch_index = beam.mpi.bunch_num  # Number of the tracked bunch in this processor
            bunch = beam[bunch_index]
            self.track_bunch(bunch, rank)
        else:
            for rank, bunch in enumerate(beam.not_empty):
                self.track_bunch(bunch, rank)

"""
This module defines the different monitor class which are used to save data
during tracking.
"""

import random
from abc import ABCMeta

from scipy.fft import rfft, rfftfreq

from mbtrack2.tracking.rf import CavityResonator

class Monitor(Element, metaclass=ABCMeta):
    """
    Abstract Monitor class used for subclass inheritance to define all the
    different kind of monitors objects. 
    
    The Monitor class is based on h5py module to be able to write data on 
    structured binary files. The class provides a common file where the 
    different Monitor subclass can write.
    
    file : HDF5 file
        Common file where all monitors, Monitor subclass elements, write the
        saved data. Based on class attribute _file_storage.
    file_name : string
        Name of the HDF5 file where the data is stored. Based on class 
        attribute _file_name_storage.
        
    Methods
    -------
    monitor_init(group_name, save_every, buffer_size, total_size,
                     dict_buffer, dict_file, file_name=None, mpi_mode=True)
        Method called to initialize Monitor subclass.
    write()
        Write data from buffer to the HDF5 file.
    to_buffer(object_to_save)
        Save data to buffer.
    close()
        Close the HDF5 file shared by all Monitor subclass, must be called 
        by at least an instance of a Montior subclass at the end of the 
        tracking.
    track_bunch_data(object_to_save)
        Track method to use when saving bunch data.
    """

    _file_name_storage = []
    _file_storage = []

    @property
    def file_name(self):
        """Common file where all monitors, Monitor subclass elements, write the
        saved data."""
        try:
            return self._file_name_storage[0]
            print("The HDF5 file name for monitors is not set.")
            raise ValueError

    @property
    def file(self):
        """Name of the HDF5 file where the data is stored."""
        try:
            return self._file_storage[0]
            print("The HDF5 file to store data is not set.")

    def monitor_init(self,
                     group_name,
                     save_every,
                     buffer_size,
                     total_size,
                     dict_buffer,
                     dict_file,
                     file_name=None,
                     mpi_mode=False,
                     dict_dtype=None):
        """
        Method called to initialize Monitor subclass. 
        
        group_name : string
            Name of the HDF5 group in which the data for the current monitor 
            will be saved.
        save_every : int or float
            Set the frequency of the save. The data is saved every save_every 
            call of the montior.
        buffer_size : int or float
            Size of the save buffer.
        total_size : int or float
            Total size of the save. The following relationships between the 
            parameters must exist: 
                total_size % buffer_size == 0
                number of call to track / save_every == total_size
        dict_buffer : dict
            Dictionary with keys as the attribute name to save and values as
            the shape of the buffer to create to hold the attribute, like 
            (key.shape, buffer_size)
        dict_file : dict
            the shape of the dataset to create to hold the attribute, like 
            (key.shape, total_size)
        file_name : string, optional
            Name of the HDF5 where the data will be stored. Must be specified
            the first time a subclass of Monitor is instancied and must be None
            the following times.
        mpi_mode : bool, optional
            If True, open the HDF5 file in parallel mode, which is needed to
            allow several cores to write in the same file at the same time.
            If False, open the HDF5 file in standard mode.
        dict_dtype : dict, optional
            the dtype to use to save the values.
            If None, float is used for all attributes.
        """

        # setup and open common file for all monitors
        if file_name is not None:
            if len(self._file_name_storage) == 0:
                self._file_name_storage.append(file_name + ".hdf5")
                if len(self._file_storage) == 0:
                    if mpi_mode == True:
                        from mpi4py import MPI
                        f = hp.File(self.file_name,
                                    "a",
                                    comm=MPI.COMM_WORLD)
                    else:
                        f = hp.File(self.file_name, "a", libver='earliest')
                    self._file_storage.append(f)
                else:
                    raise ValueError("File is already open.")
            else:
                    "File name for monitors is already attributed.")

        self.group_name = group_name
        self.save_every = int(save_every)
        self.total_size = int(total_size)
        self.buffer_size = int(buffer_size)
        if total_size % buffer_size != 0:
            raise ValueError("total_size must be divisible by buffer_size.")
        self.buffer_count = 0
        self.write_count = 0
        self.track_count = 0

        # setup attribute buffers from values given in dict_buffer
        for key, value in dict_buffer.items():
            if dict_dtype == None:
                self.__setattr__(key, np.zeros(value))
            else:
                self.__setattr__(key, np.zeros(value, dtype=dict_dtype[key]))
        self.time = np.zeros((self.buffer_size, ), dtype=int)

        # create HDF5 groups and datasets to save data from group_name and
        # dict_file
        self.g = self.file.require_group(self.group_name)
        self.g.require_dataset("time", (self.total_size, ), dtype=int)
        for key, value in dict_file.items():
                self.g.require_dataset(key, value, dtype=float)
            else:
                self.g.require_dataset(key, value, dtype=dict_dtype[key])

        # create a dictionary which handle slices
        slice_dict = {}
            slice_dict[key] = []
            for i in range(len(value) - 1):
                slice_dict[key].append(slice(None))
        self.slice_dict = slice_dict

    def write(self):
        """Write data from buffer to the HDF5 file."""

        self.file[self.group_name]["time"][self.write_count *
                                           self.buffer_size:(self.write_count +
                                                             1) *
                                           self.buffer_size] = self.time
        for key, value in self.dict_buffer.items():
            slice_list = list(self.slice_dict[key])
            slice_list.append(
                slice(self.write_count * self.buffer_size,
                      (self.write_count + 1) * self.buffer_size))
            slice_tuple = tuple(slice_list)
            self.file[
                self.group_name][key][slice_tuple] = self.__getattribute__(key)

        self.file.flush()
        self.write_count += 1

    def to_buffer(self, object_to_save):
        """
        
        object_to_save : python object
            Depends on the Monitor subclass, typically a Beam or Bunch object.
        """
        self.time[self.buffer_count] = self.track_count
            slice_list.append(self.buffer_count)
            self.__getattribute__(
                key)[slice_tuple] = object_to_save.__getattribute__(key)
        self.buffer_count += 1

        if self.buffer_count == self.buffer_size:
            self.write()
            self.buffer_count = 0

    def close(self):
        """
        tracking.
        """
        try:
            self.file.close()
            Monitor._file_name_storage = []
            Monitor._file_storage = []
        except ValueError:
            pass

    def track_bunch_data(self, object_to_save, check_empty=False):
        """
        
        object_to_save : Beam or Bunch
        check_emptiness: bool
            If True, check if the bunch is empty. If it is, then do nothing.
        """
        save = True
        if self.track_count % self.save_every == 0:

            if isinstance(object_to_save, Beam):
                if (object_to_save.mpi_switch == True):
                    if object_to_save.mpi.bunch_num == self.bunch_number:
                        bunch = object_to_save[object_to_save.mpi.bunch_num]
                    else:
                        save = False
                else:
                    bunch = object_to_save[self.bunch_number]
            elif isinstance(object_to_save, Bunch):
                bunch = object_to_save
            else:
                raise TypeError(
                    "object_to_save should be a Beam or Bunch object.")

            if save:
                if (check_empty == False) or (bunch.is_empty == False):
                    self.to_buffer(bunch)

        self.track_count += 1

class BunchMonitor(Monitor):
    """
    Monitor a single bunch and save attributes 
    (mean, std, emit, current, and cs_invariant).
    
    bunch_number : int
        Bunch to monitor
    save_every : int or float
        Set the frequency of the save. The data is saved every save_every 
        call of the montior.
    buffer_size : int or float
        Size of the save buffer.
    total_size : int or float
        Total size of the save. The following relationships between the 
        parameters must exist: 
            total_size % buffer_size == 0
            number of call to track / save_every == total_size
    file_name : string, optional
        Name of the HDF5 where the data will be stored. Must be specified
        the first time a subclass of Monitor is instancied and must be None
        the following times.
    mpi_mode : bool, optional
        If True, open the HDF5 file in parallel mode, which is needed to
        allow several cores to write in the same file at the same time.
        If False, open the HDF5 file in standard mode.

    Methods
    -------
    track(object_to_save)
        Save data
    """

                 bunch_number,
                 save_every,
                 buffer_size,
                 total_size,
                 file_name=None,
                 mpi_mode=False):

        self.bunch_number = bunch_number
        group_name = "BunchData_" + str(self.bunch_number)
        dict_buffer = {
            "mean": (6, buffer_size),
            "std": (6, buffer_size),
            "emit": (3, buffer_size),
            "current": (buffer_size, ),
            "cs_invariant": (3, buffer_size)
        }
        dict_file = {
            "mean": (6, total_size),
            "std": (6, total_size),
            "emit": (3, total_size),
            "current": (total_size, ),
            "cs_invariant": (3, total_size)
        }
        self.monitor_init(group_name, save_every, buffer_size, total_size,
                          dict_buffer, dict_file, file_name, mpi_mode)

        self.dict_buffer = dict_buffer
        self.dict_file = dict_file

    def track(self, object_to_save):
        """
        Save data
        
        object_to_save : Bunch or Beam object
        """
        self.track_bunch_data(object_to_save)

class PhaseSpaceMonitor(Monitor):
    """
    Monitor a single bunch and save the full phase space.
    
    mp_number : int or float
        Number of macroparticle in the phase space to save. If less than the 
        total number of macroparticles, a random fraction of the bunch is saved.

    Methods
    -------
        Save data
    """

                 mp_number,

        self.mp_number = int(mp_number)
        group_name = "PhaseSpaceData_" + str(self.bunch_number)
            "particles": (self.mp_number, 6, buffer_size),
            "alive": (self.mp_number, buffer_size)
        }
            "particles": (self.mp_number, 6, total_size),
            "alive": (self.mp_number, total_size)
        }

        """
        Save data
        
        """

    def to_buffer(self, bunch):
        """
        
        """

        if len(bunch.alive) != self.mp_number:
            index = np.arange(len(bunch.alive))
            samples_meta = random.sample(list(index), self.mp_number)
            samples = sorted(samples_meta)
        else:
            samples = slice(None)

        self.alive[:, self.buffer_count] = bunch.alive[samples]
        for i, dim in enumerate(bunch):
            self.particles[:, i,
                           self.buffer_count] = bunch.particles[dim][samples]

class BeamMonitor(Monitor):
    """
    Monitor the full beam and save each bunch attributes (mean, std, emit and 
    current).
    
    h : int
        Harmonic number of the ring.

    Methods
    -------
        Save data    
    """

                 h,

        group_name = "Beam"
            "mean": (6, h, buffer_size),
            "std": (6, h, buffer_size),
            "emit": (3, h, buffer_size),
            "current": (h, buffer_size),
            "cs_invariant": (3, h, buffer_size)
        }
            "mean": (6, h, total_size),
            "std": (6, h, total_size),
            "emit": (3, h, total_size),
            "current": (h, total_size),
            "cs_invariant": (3, h, total_size)
        }

        """
        Save data
        
        """
                self.to_buffer(beam[beam.mpi.bunch_num], beam.mpi.bunch_num)
            else:
                self.to_buffer_no_mpi(beam)

    def to_buffer(self, bunch, bunch_num):
        """
        Save data to buffer, if mpi is being used.
        
        """

        self.mean[:, bunch_num, self.buffer_count] = bunch.mean
        self.std[:, bunch_num, self.buffer_count] = bunch.std
        self.emit[:, bunch_num, self.buffer_count] = bunch.emit
        self.current[bunch_num, self.buffer_count] = bunch.current
        self.cs_invariant[:, bunch_num, self.buffer_count] = bunch.cs_invariant

            self.write(bunch_num)

    def to_buffer_no_mpi(self, beam):
        """
        Save data to buffer, if mpi is not being used.
        
        """

        self.mean[:, :, self.buffer_count] = beam.bunch_mean
        self.std[:, :, self.buffer_count] = beam.bunch_std
        self.emit[:, :, self.buffer_count] = beam.bunch_emit
        self.current[:, self.buffer_count] = beam.bunch_current
        self.cs_invariant[:, :, self.buffer_count] = beam.bunch_cs

            self.write_no_mpi()

    def write(self, bunch_num):
        """
        Write data from buffer to the HDF5 file, if mpi is being used.
        
        """
                                                             1) *

        self.file[self.group_name][
            "mean"][:, bunch_num, self.write_count *
                    self.buffer_size:(self.write_count + 1) *
                    self.buffer_size] = self.mean[:, bunch_num, :]

        self.file[
            self.group_name]["std"][:, bunch_num, self.write_count *
                                    self.buffer_size:(self.write_count + 1) *
                                    self.buffer_size] = self.std[:,
                                                                 bunch_num, :]

            "emit"][:, bunch_num, self.write_count *
                    self.buffer_size] = self.emit[:, bunch_num, :]

        self.file[self.group_name]["current"][
            bunch_num,
            self.write_count * self.buffer_size:(self.write_count + 1) *
            self.buffer_size] = self.current[bunch_num, :]

            "cs_invariant"][:, bunch_num, self.write_count *
                            self.buffer_size:(self.write_count + 1) *
                            self.buffer_size] = self.cs_invariant[:,
                                                                  bunch_num, :]

    def write_no_mpi(self):
        """
        Write data from buffer to the HDF5 file, if mpi is not being used.
        """

                                                             1) *

        self.file[self.group_name]["mean"][:, :, self.write_count *
                                                             1) *
                                           self.buffer_size] = self.mean

        self.file[self.group_name]["std"][:, :, self.write_count *
                                          self.buffer_size:(self.write_count +
                                                            1) *
                                          self.buffer_size] = self.std

        self.file[self.group_name]["emit"][:, :, self.write_count *
                                                             1) *
                                           self.buffer_size] = self.emit

        self.file[self.group_name]["current"][:, self.write_count *
                                              self.buffer_size:
                                              (self.write_count + 1) *
                                              self.buffer_size] = self.current

            "cs_invariant"][:, :, self.write_count *
                            self.buffer_size] = self.cs_invariant

class ProfileMonitor(Monitor):
    """
    Monitor a single bunch and save bunch profiles.
    
        Bunch to monitor.
    dimensions : str or list of str, optional
        Dimensions to save.
    n_bin : int or list of int, optional
        Number of bin to use in each dimension.

    Methods
    -------
        Save data.
    """

                 dimensions="tau",
                 n_bin=75,

        group_name = "ProfileData_" + str(self.bunch_number)

            self.dimensions = [dimensions]
        else:
            self.dimensions = dimensions

            self.n_bin = np.ones((len(self.dimensions), ), dtype=int) * n_bin
        else:
            self.n_bin = n_bin

        dict_buffer = {}
        dict_file = {}
        for index, dim in enumerate(self.dimensions):
            dict_buffer.update({dim: (self.n_bin[index] - 1, buffer_size)})
            dict_buffer.update(
                {dim + "_bin": (self.n_bin[index] - 1, buffer_size)})
            dict_file.update({dim: (self.n_bin[index] - 1, total_size)})
            dict_file.update(
                {dim + "_bin": (self.n_bin[index] - 1, total_size)})

        """
        
        """

            bins, sorted_index, profile, center = bunch.binning(
                dim, self.n_bin[index])
            self.__getattribute__(dim + "_bin")[:, self.buffer_count] = center
            self.__getattribute__(dim)[:, self.buffer_count] = profile

                                                             1) *

        for dim in self.dimensions:
            self.file[self.group_name][
                dim][:, self.write_count *
                     self.buffer_size:(self.write_count + 1) *
                     self.buffer_size] = self.__getattribute__(dim)
                dim + "_bin"][:, self.write_count *
                              self.buffer_size:(self.write_count + 1) *
                              self.buffer_size] = self.__getattribute__(dim +
                                                                        "_bin")

        """
        Save data.
        
        """
        self.track_bunch_data(object_to_save, check_empty=True)

class WakePotentialMonitor(Monitor):
    """
    Monitor the wake potential from a single bunch and save attributes (tau, 
    ...).
    
    wake_types : str or list of str
        Wake types to save: "Wlong, "Wxdip", ...
    n_bin : int
        Number of bin to be used to interpolate the wake potential on a fixed
        grid.

    Methods
    -------
    track(object_to_save, wake_potential_to_save)
        Save data.
    """

                 wake_types,
                 n_bin,

        group_name = "WakePotentialData_" + str(self.bunch_number)

        if isinstance(wake_types, str):
            self.wake_types = [wake_types]
        else:
            self.wake_types = wake_types

        self.n_bin = n_bin * 2

        for index, dim in enumerate(self.wake_types):
            dict_buffer.update({"tau_" + dim: (self.n_bin, buffer_size)})
            dict_file.update({"tau_" + dim: (self.n_bin, total_size)})
            dict_buffer.update({"profile_" + dim: (self.n_bin, buffer_size)})
            dict_file.update({"profile_" + dim: (self.n_bin, total_size)})
            dict_buffer.update({dim: (self.n_bin, buffer_size)})
            dict_file.update({dim: (self.n_bin, total_size)})
            if dim == "Wxdip" or dim == "Wydip":
                dict_buffer.update(
                    {"dipole_" + dim: (self.n_bin, buffer_size)})
                dict_file.update({"dipole_" + dim: (self.n_bin, total_size)})

    def to_buffer(self, wp):
        """
        
        wp : WakePotential object
        """

            tau0 = wp.__getattribute__("tau0_" + dim)
            profile0 = wp.__getattribute__("profile0_" + dim)
            WP0 = wp.__getattribute__(dim)
            if dim == "Wxdip":
                dipole0 = wp.__getattribute__("dipole_x")
            if dim == "Wydip":
                dipole0 = wp.__getattribute__("dipole_y")

            tau = np.linspace(tau0[0], tau0[-1], self.n_bin)
            f = interp1d(tau0, WP0, fill_value=0, bounds_error=False)
            WP = f(tau)
            g = interp1d(tau0, profile0, fill_value=0, bounds_error=False)
            profile = g(tau)
                h = interp1d(tau0, dipole0, fill_value=0, bounds_error=False)
                dipole = h(tau)

            self.__getattribute__("tau_" +
                                  dim)[:,
                                       self.buffer_count] = tau + wp.tau_mean
            self.__getattribute__("profile_" +
                                  dim)[:, self.buffer_count] = profile
            self.__getattribute__(dim)[:, self.buffer_count] = WP
                self.__getattribute__("dipole_" +
                                      dim)[:, self.buffer_count] = dipole

                                                             1) *

        for dim in self.wake_types:
                "tau_" +
                     self.buffer_size] = self.__getattribute__("tau_" + dim)
                "profile_" +
                     self.buffer_size] = self.__getattribute__("profile_" +
                                                               dim)
                self.file[self.group_name][
                    "dipole_" +
                    dim][:, self.write_count *
                         self.buffer_size:(self.write_count + 1) *
                         self.buffer_size] = self.__getattribute__("dipole_" +
                                                                   dim)

    def track(self, object_to_save, wake_potential_to_save):
        """
        Save data.
        
        wake_potential_to_save : WakePotential object
        """
        if isinstance(object_to_save, Beam):
            if (object_to_save.mpi_switch == True):
                if object_to_save.mpi.bunch_num == self.bunch_number:
                    save = True
                else:
                    save = False
            else:
                raise NotImplementedError("WakePotentialMonitor for Beam " +
                                          "objects is only available " +
                                          "with MPI mode.")
        elif isinstance(object_to_save, Bunch):
            save = True
        else:
            raise TypeError("object_to_save should be a Beam or Bunch object.")

        if save and (self.track_count % self.save_every == 0):
            self.to_buffer(wake_potential_to_save)

class BunchSpectrumMonitor(Monitor):
    """
    Monitor the coherent and incoherent bunch spectrums. 
    
        Total number of macro-particles in the bunch.
    sample_size : int or float
        Number of macro-particles to be used for tune and FFT computation.
        This number cannot exceed mp_number.  
        Set the frequency of the save. The spectrums are computed every 
        save_every call of the montior.
            number of call to track / save_every == total_size - 1 
    dim : str, optional
        Dimensions in which the spectrums have to be computed.
        Can be:
                - "all"
                - "tau"
                - "x"
                - "y"
                - "xy" or "yx"
                - "xtau" or "taux"
                - "ytau" or "tauy"
    n_fft : int or float, optional
        The number of points used for FFT computation, if n_fft is bigger than
        save_every zero-padding is applied.
        If None, save_every is used.
        
    fft_resolution : float
        Return the fft resolution in [Hz].
    signal_resolution : float
        Return the signal resolution in [Hz].
    frequency_samples : array of float
        Return the fft frequency samples in [Hz].        
        
    Methods
    -------
    track(bunch):
        Save spectrum data.
    
    """

                 ring,
                 sample_size,
                 dim="all",
                 n_fft=None,

        if n_fft is None:
            self.n_fft = int(save_every)
        else:
            self.n_fft = int(n_fft)

        self.sample_size = int(sample_size)
        self.store_dict = {"x": 0, "y": 1, "tau": 2}

        if dim == "all":
            self.track_dict = {"x": 0, "y": 1, "tau": 2}
            self.mean_index = [True, False, True, False, True, False]
        elif dim == "tau":
            self.track_dict = {"tau": 0}
            self.mean_index = [False, False, False, False, True, False]
        elif dim == "x":
            self.track_dict = {"x": 0}
            self.mean_index = [True, False, False, False, False, False]
        elif dim == "y":
            self.track_dict = {"y": 0}
            self.mean_index = [False, False, True, False, False, False]
        elif dim == "xy" or dim == "yx":
            self.track_dict = {"x": 0, "y": 1}
            self.mean_index = [True, False, True, False, False, False]
        elif dim == "xtau" or dim == "taux":
            self.track_dict = {"x": 0, "tau": 1}
            self.mean_index = [True, False, False, False, True, False]
        elif dim == "ytau" or dim == "tauy":
            self.track_dict = {"y": 0, "tau": 1}
            self.mean_index = [False, False, True, False, True, False]
        else:
            raise ValueError("dim is not correct.")

        self.size_list = len(self.track_dict)

        group_name = "BunchSpectrum_" + str(self.bunch_number)

            "incoherent": (3, self.n_fft // 2 + 1, buffer_size),
            "coherent": (3, self.n_fft // 2 + 1, buffer_size),
            "mean_incoherent": (3, buffer_size),
            "std_incoherent": (3, buffer_size)
        }
            "incoherent": (3, self.n_fft // 2 + 1, total_size),
            "coherent": (3, self.n_fft // 2 + 1, total_size),
            "mean_incoherent": (3, total_size),
            "std_incoherent": (3, total_size)
        }

        self.save_count = 0

        self.positions = np.zeros(
            (self.size_list, self.sample_size, self.save_every + 1))
        self.mean = np.zeros((self.size_list, self.save_every + 1))

        index = np.arange(0, int(mp_number))
        self.index_sample = sorted(random.sample(list(index),
                                                 self.sample_size))

        self.incoherent = np.zeros((3, self.n_fft // 2 + 1, self.buffer_size))
        self.coherent = np.zeros((3, self.n_fft // 2 + 1, self.buffer_size))

        self.file[self.group_name].create_dataset("freq",
                                                  data=self.frequency_samples)

    @property
    def fft_resolution(self):
        """
        
        It is defined as the sampling frequency over the number of samples.
        """
        return self.ring.f0 / self.n_fft

    @property
    def signal_resolution(self):
        """
        
        It is defined as the inverse of the signal length.
        """
        return 1 / (self.ring.T0 * self.save_every)

    @property
    def frequency_samples(self):
        """
        Return the fft frequency samples in [Hz].
        """
        return rfftfreq(self.n_fft, self.ring.T0)

        """

        object_to_save : Beam or Bunch object

        """
        save = True
                    bunch = object_to_save[object_to_save.mpi.bunch_num]
                else:
            else:
                bunch = object_to_save[self.bunch_number]
            bunch = object_to_save
        else:

        if save:
            try:
                for key, value in self.track_dict.items():
                    self.positions[value, :, self.save_count] = bunch[key][
                        self.index_sample]
                self.positions[value, :, self.save_count] = np.nan

            self.mean[:, self.save_count] = bunch.mean[self.mean_index]

            self.save_count += 1

            if self.track_count > 0 and self.track_count % self.save_every == 0:
                self.to_buffer(bunch)
                self.save_count = 0

            self.track_count += 1

        """
        A method to hold saved data before writing it to the output file.

        """

        for key, value in self.track_dict.items():
            incoherent, mean_incoherent, std_incoherent = self.get_incoherent_spectrum(
                self.positions[value, :, :])
            self.incoherent[self.store_dict[key], :,
                            self.buffer_count] = incoherent
            self.mean_incoherent[self.store_dict[key],
                                 self.buffer_count] = mean_incoherent
            self.std_incoherent[self.store_dict[key],
                                self.buffer_count] = std_incoherent
            self.coherent[self.store_dict[key], :,
                          self.buffer_count] = self.get_coherent_spectrum(
                              self.mean[value])

        """
        Write data from buffer to output file.

        """
                                                             1) *

            self.group_name]["incoherent"][:, :, self.write_count *
                                                             1) *
                                           self.buffer_size] = self.incoherent
            "mean_incoherent"][:, self.write_count *
                               self.buffer_size:(self.write_count + 1) *
                               self.buffer_size] = self.mean_incoherent
            "std_incoherent"][:, self.write_count *
                              self.buffer_size] = self.std_incoherent
            self.group_name]["coherent"][:, :, self.write_count *
                                         self.buffer_size:(self.write_count +
                                                           1) *
                                         self.buffer_size] = self.coherent

    def get_incoherent_spectrum(self, positions):
        """
        Compute the incoherent spectrum i.e. the average of the absolute value 
        of the FT of the position of every particule of the bunch. 

        Returns
        -------
        incoherent : array
            Bunch incoherent spectrum.
        mean_incoherent : float
            Mean frequency of the maximum of each individual particle spectrum
            in [Hz].
        std_incoherent : float
            Standard deviation of the frequency of the maximum of each 
            individual particle spectrum in [Hz].

        """
        fourier = rfft(positions, n=self.n_fft)
        fourier_abs = np.abs(fourier)
        max_array = np.argmax(fourier_abs, axis=1)
        freq_array = self.frequency_samples[max_array]
        mean_incoherent = np.mean(freq_array)
        std_incoherent = np.std(freq_array)
        incoherent = np.mean(fourier_abs, axis=0)

        return incoherent, mean_incoherent, std_incoherent

    def get_coherent_spectrum(self, mean):
        """
        Compute the coherent spectrum i.e. the absolute value of the FT of the
        mean position of the bunch.

        Returns
        -------
        coherent : array
            Bunch coherent spectrum.

        """
        coherent = np.abs(rfft(mean, n=self.n_fft))

        return coherent

class BeamSpectrumMonitor(Monitor):
    """
    Monitor coherent beam spectrum. 
    
        Can be:
                - "all"
                - "tau"
                - "x"
                - "y"
        
    Methods
    -------
    
    """

                 ring,

        else:

        else:

        group_name = "BeamSpectrum"

        dict_buffer = {"coherent": (3, self.n_fft // 2 + 1, buffer_size)}
        dict_file = {"coherent": (3, self.n_fft // 2 + 1, total_size)}

        self.mean = np.zeros((self.size_list, ring.h, self.save_every))

    @property
        """
        
        """
        return self.ring.f1 / self.n_fft

    @property
        """
        
        """

    @property
        """
        """
        return rfftfreq(self.n_fft, self.ring.T1)

        """
        Save mean data.

        """
        if (beam.mpi_switch == True):
            bunch_num = beam.mpi.bunch_num
            bunch = beam[bunch_num]
            self.mean[:, bunch_num,
                      self.save_count] = bunch.mean[self.mean_index]
        else:
            self.mean[:, :,
                      self.save_count] = beam.bunch_mean[self.mean_index, :]

        self.save_count += 1

        if self.save_count == self.save_every:
            self.to_buffer(beam)
            self.save_count = 0

    def to_buffer(self, beam):
        """

        """

                data_core = self.mean[value, beam.mpi.bunch_num, :]
                full_data = beam.mpi.comm.allgather(data_core)
                data = np.reshape(full_data, (-1), 'F')
            else:
                data = np.reshape(self.mean[value, :, :], (-1), 'F')
                          self.buffer_count] = self.get_beam_spectrum(data)

        """

        """
                                                             1) *

                                                           1) *

    def get_beam_spectrum(self, mean):
        """
        Compute the beam coherent spectrum i.e. the absolute value of the FT 
        of the mean position of every bunch.

        Returns
        -------
            The beam coherent spectrum.

        """

class CavityMonitor(Monitor):
    """
    Monitor a CavityResonator object and save attributes.
    
    cavity_name : str
        Name of the CavityResonator object to monitor.

    Methods
    -------
    track(beam, cavity)
        Save data
    """

                 cavity_name,
                 ring,

        self.cavity_name = cavity_name

        group_name = cavity_name
            "cavity_phasor_record": (
                ring.h,
                buffer_size,
            ),
            "beam_phasor_record": (
                ring.h,
            ),
            "generator_phasor_record": (
                ring.h,
            ),
            "ig_phasor_record": (
                ring.h,
            ),
            "DFB_ig_phasor": (
                ring.h,
            ),
            "detune": (buffer_size, ),
            "psi": (buffer_size, ),
            "Vg": (buffer_size, ),
            "theta_g": (buffer_size, ),
            "Pg": (buffer_size, ),
            "Rs": (buffer_size, ),
            "Q": (buffer_size, ),
            "QL": (buffer_size, ),
            "Vc": (buffer_size, ),
            "theta": (buffer_size, ),
        }
                ring.h,
                total_size,
            ),
                ring.h,
            ),
                ring.h,
            ),
                ring.h,
            ),
                ring.h,
            ),
            "detune": (total_size, ),
            "psi": (total_size, ),
            "Vg": (total_size, ),
            "theta_g": (total_size, ),
            "Pg": (total_size, ),
            "Rs": (total_size, ),
            "Q": (total_size, ),
            "QL": (total_size, ),
            "Vc": (total_size, ),
            "theta": (total_size, )
        }
        dict_dtype = {
            "cavity_phasor_record": complex,
            "beam_phasor_record": complex,
            "generator_phasor_record": complex,
            "ig_phasor_record": complex,
            "DFB_ig_phasor": complex,
            "detune": float,
            "psi": float,
            "Vg": float,
            "theta_g": float,
            "Pg": float,
            "Rs": float,
            "Q": float,
            "QL": float,
            "Vc": float,
            "theta": float
        }

                          dict_buffer, dict_file, file_name, mpi_mode,
                          dict_dtype)

    def track(self, beam, cavity):
        """
        Save data
        
        cavity : CavityResonator object
        """
            if isinstance(cavity, CavityResonator):
                if beam.mpi_switch == False:
                    self.to_buffer(cavity)
                elif beam.mpi.rank == 0:
                else:
                    pass
            else:
                raise TypeError("cavity should be a CavityResonator object.")

"""
Module for plotting the data recorded by the monitor module during the 
tracking.
"""

import matplotlib as mpl
from scipy.stats import gmean

def plot_beamdata(filenames,
                  dataset="mean",
                  dimension="tau",
                  stat_var="mean",
                  x_var="time",
                  turn=None,
                  legend=None):
    """
    Plot 2D data recorded by BeamMonitor.

    filenames : str or list of str
        Names of the HDF5 files to be plotted.
    dataset : {"current","emit","mean","std","cs_invariant"}
        HDF5 file's dataset to be plotted. The default is "mean".
    dimension : str
         The dimension of the dataset to plot:
            for "emit", dimension = {"x","y","s"},
            for "cs_invariant", dimension = {"x","y","s"},
            for "mean" and "std", dimension = {"x","xp","y","yp","tau","delta"}.
            not used if "current".
        The default is "tau".
    stat_var : {"mean", "std"}
        Statistical value of the dimension.
    x_var : {"time", "index"}
        Choice of the horizontal axis:
            "time" corresponds to turn number.
            "index" corresponds to bunch index.
    turn : int or float, optional
        Choice of the turn to plot when x_var = "index".
        If None, the last turn is plotted.
    legend : list of str, optional
        Legend to add for each file.

    Return
    ------
        Figure object with the plot on it.

    """

    if isinstance(filenames, str):
        filenames = [filenames]

    for filename in filenames:
        file = hp.File(filename, "r")
        time = np.array(file["Beam"]["time"])
        data = np.array(file["Beam"][dataset])

        if x_var == "time":
            x = time
            x_label = "Number of turns"
            bunch_index = (file["Beam"]["current"][:, 0] != 0).nonzero()[0]
            if dataset == "current":
                y = np.nansum(data[bunch_index, :], 0) * 1e3
                y_label = "Total current (mA)"
            elif dataset == "emit":
                dimension_dict = {"x": 0, "y": 1, "s": 2}
                axis = dimension_dict[dimension]
                label = [
                    "$\\epsilon_{x}$ (m.rad)", "$\\epsilon_{y}$ (m.rad)",
                    "$\\epsilon_{s}$ (s)"
                ]
                if stat_var == "mean":
                    y = np.nanmean(data[axis, bunch_index, :], 0)
                elif stat_var == "std":
                    y = np.nanstd(data[axis, bunch_index, :], 0)
                y_label = stat_var + " " + label[axis]
            elif dataset == "cs_invariant":
                label = ['$J_x$ (m)', '$J_y$ (m)', '$J_s$ (s)']
            elif dataset == "mean" or dataset == "std":
                dimension_dict = {
                    "x": 0,
                    "xp": 1,
                    "y": 2,
                    "yp": 3,
                    "tau": 4,
                    "delta": 5
                }
                scale = [1e6, 1e6, 1e6, 1e6, 1e12, 1]
                label = [
                    "x (um)", "x' ($\\mu$rad)", "y (um)", "y' ($\\mu$rad)",
                    "$\\tau$ (ps)", "$\\delta$"
                ]
                    y = np.nanmean(data[axis, bunch_index, :], 0) * scale[axis]
                    y = np.nanstd(data[axis, bunch_index, :], 0) * scale[axis]
                label_sup = {"mean": "mean of ", "std": "std of "}
                y_label = label_sup[stat_var] + dataset + " " + label[axis]

        elif x_var == "index":
            h = len(file["Beam"]["mean"][0, :, 0])
            x = np.arange(h)
            x_label = "Bunch index"
            if turn is None:
                idx = -1
            else:
                idx = np.where(time == int(turn))[0]
                if (idx.size == 0):
                    raise ValueError("Turn is not valid.")

                y = data[:, idx] * 1e3
                y_label = "Bunch current (mA)"
                label = [
                ]
                y = data[axis, :, idx]
                y_label = label[axis]
                    "x": 0,
                    "xp": 1,
                    "y": 2,
                    "yp": 3,
                    "tau": 4,
                    "delta": 5
                }
                label = [
                ]
                y = data[axis, :, idx] * scale[axis]
                y_label = dataset + " " + label[axis]
        else:
            raise ValueError("x_var should be time or index")

        y = np.squeeze(y)

        ax.set_xlabel(x_label)
        ax.set_ylabel(y_label)
        if legend is not None:
            plt.legend(legend)

        file.close()

    return fig

def streak_beamdata(filename, dataset="mean", dimension="tau", cm_lim=None):
    """
    Plot 3D data recorded by BeamMonitor.

    filename : str
        Name of the HDF5 file that contains the data.
    cm_lim : list [vmin, vmax], optional
        Colormap scale for the "streak" plot.

    Return
    ------

    """

    file = hp.File(filename, "r")
    data = file["Beam"]
    time = np.array(data["time"])

    h = len(data["mean"][0, :, 0])
    x = np.arange(h)
    x_label = "Bunch index"
    y = time
    y_label = "Number of turns"
    if dataset == "current":
        z = (np.array(data["current"]) * 1e3).T
        z_label = "Bunch current (mA)"
        title = z_label
    elif dataset == "emit":
        dimension_dict = {"x": 0, "y": 1, "s": 2}
        axis = dimension_dict[dimension]
        label = [
            "$\\epsilon_{x}$ (m.rad)", "$\\epsilon_{y}$ (m.rad)",
            "$\\epsilon_{s}$ (s)"
        ]
        z = np.array(data["emit"][axis, :, :]).T
        z_label = label[axis]
    elif dataset == "cs_invariant":
        label = ['$J_x$ (m)', '$J_y$ (m)', '$J_s$ (s)']
        z = np.array(data["cs_invariant"][axis, :, :]).T
    else:
        dimension_dict = {
            "x": 0,
            "xp": 1,
            "y": 2,
            "yp": 3,
            "tau": 4,
            "delta": 5
        }
        scale = [1e6, 1e6, 1e6, 1e6, 1e12, 1]
        label = [
            "x (um)", "x' ($\\mu$rad)", "y (um)", "y' ($\\mu$rad)",
            "$\\tau$ (ps)", "$\\delta$"
        ]
        z = np.array(data[dataset][axis, :, :]).T * scale[axis]
        if dataset == "mean":
            title = label[axis] + " CM"
        elif dataset == "std":
            title = label[axis] + " RMS"

    ax.set_xlabel(x_label)
    ax.set_ylabel(y_label)
    ax.set_title(title)

    if dataset == "mean":
        cmap = mpl.cm.coolwarm  # diverging
    else:
        cmap = mpl.cm.inferno  # sequential

    c = ax.imshow(z,
                  cmap=cmap,
                  origin='lower',
                  aspect='auto',
                  extent=[x.min(), x.max(), y.min(),
                          y.max()])
    if cm_lim is not None:
        c.set_clim(vmin=cm_lim[0], vmax=cm_lim[1])
    cbar = fig.colorbar(c, ax=ax)
    cbar.set_label(z_label)

    file.close()

    return fig

def plot_bunchdata(filenames,
                   bunch_number,
                   dataset,
                   dimension="x",
                   legend=None):
    """
    Plot data recorded by BunchMonitor.
    
    bunch_number : int or list of int
        Bunch to plot. This has to be identical to 'bunch_number' parameter in 
        'BunchMonitor' object.
    dataset : {"current", "emit", "mean", "std", "cs_invariant"}
        HDF5 file's dataset to be plotted.
    dimension : str, optional
        The dimension of the dataset to plot. Use "None" for "current",
        otherwise use the following : 
            for "mean" and "std", dimension = {"x","xp","y","yp","tau","delta"},
            for "action", dimension = {"x","y"}.
        
    Return
    ------

    """

    if isinstance(bunch_number, int):
        ll = []
        for i in range(len(filenames)):
            ll.append(bunch_number)
        bunch_number = ll

    for i, filename in enumerate(filenames):
        group = "BunchData_{0}".format(
            bunch_number[i])  # Data group of the HDF5 file

        if dataset == "current":
            y_var = file[group][dataset][:] * 1e3
            label = "current (mA)"

        elif dataset == "emit":
            dimension_dict = {"x": 0, "y": 1, "s": 2}

            y_var = file[group][dataset][dimension_dict[dimension]]

            if dimension == "x": label = "hor. emittance (m.rad)"
            elif dimension == "y": label = "ver. emittance (m.rad)"
            elif dimension == "s": label = "long. emittance (s)"

        elif dataset == "mean" or dataset == "std":
            dimension_dict = {
                "x": 0,
                "xp": 1,
                "y": 2,
                "yp": 3,
                "tau": 4,
                "delta": 5
            }
            axis_index = dimension_dict[dimension]

            y_var = file[group][dataset][axis_index] * scale[axis_index]
            if dataset == "mean":
                label_list = [
                    "x ($\\mu$m)", "x' ($\\mu$rad)", "y ($\\mu$m)",
                    "y' ($\\mu$rad)", "$\\tau$ (ps)", "$\\delta$"
                ]
            else:
                    "$\\sigma_x$ ($\\mu$m)", "$\\sigma_{x'}$ ($\\mu$rad)",
                    "$\\sigma_y$ ($\\mu$m)", "$\\sigma_{y'}$ ($\\mu$rad)",
                    "$\\sigma_{\\tau}$ (ps)", "$\\sigma_{\\delta}$"
                ]

            label = label_list[axis_index]

        elif dataset == "cs_invariant":
            y_var = file[group][dataset][axis_index]
            label_list = ['$J_x$ (m)', '$J_y$ (m)', '$J_s$ (s)']

        x_axis = file[group]["time"][:]
        xlabel = "Number of turns"

        ax.plot(x_axis, y_var)
        ax.set_xlabel(xlabel)

    return fig

def plot_phasespacedata(filename,
                        bunch_number,
                        x_var,
                        y_var,
                        turn,
                        only_alive=True,
                        plot_size=1,
                        plot_kind='kde'):
    """
    Plot data recorded by PhaseSpaceMonitor.

        'PhaseSpaceMonitor' object.
    x_var, y_var : str {"x", "xp", "y", "yp", "tau", "delta"}
        If dataset is "particles", the variables to be plotted on the 
        horizontal and the vertical axes need to be specified.
    turn : int
        Turn at which the data will be extracted.
    only_alive : bool, optional
        When only_alive is True, only alive particles are plotted and dead 
        particles will be discarded.
    plot_size : [0,1], optional
        Number of macro-particles to plot relative to the total number 
        of macro-particles recorded. This option helps reduce processing time
        when the data is big.
    plot_kind : {'scatter', 'kde', 'hex', 'reg', 'resid'}, optional
        The plot style. The default is 'kde'. 
        
    Return
    ------
    """

    group = "PhaseSpaceData_{0}".format(bunch_number)
    dataset = "particles"

    option_dict = {"x": 0, "xp": 1, "y": 2, "yp": 3, "tau": 4, "delta": 5}
    scale = [1e3, 1e3, 1e3, 1e3, 1e12, 1]
    label = [
        "x (mm)", "x' (mrad)", "y (mm)", "y' (mrad)", "$\\tau$ (ps)",
        "$\\delta$"
    ]

    # find the index of "turn" in time array
    turn_index = np.where(file[group]["time"][:] == turn)

    if len(turn_index[0]) == 0:
        raise ValueError("Turn {0} is not found. Enter turn from {1}.".format(
            turn, file[group]["time"][:]))

    path = file[group][dataset]
    mp_number = path[:, 0, 0].size

    if only_alive is True:
        data = np.array(file[group]["alive"])
        index = np.where(data[:, turn_index])[0]
    else:
        index = np.arange(mp_number)

    if plot_size == 1:
        samples = index
    elif plot_size < 1:
        samples_meta = random.sample(list(index), int(plot_size * mp_number))
        samples = sorted(samples_meta)
    else:
        raise ValueError("plot_size must be in range [0,1].")

    # format : sns.jointplot(x_axis, yaxis, kind)
    x_axis = path[samples, option_dict[x_var], turn_index[0][0]]
    y_axis = path[samples, option_dict[y_var], turn_index[0][0]]

    fig = sns.jointplot(x_axis * scale[option_dict[x_var]],
                        y_axis * scale[option_dict[y_var]],
                        kind=plot_kind)

    plt.xlabel(label[option_dict[x_var]])
    plt.ylabel(label[option_dict[y_var]])

    return fig

def plot_profiledata(filename,
                     bunch_number,
                     dimension="tau",
                     start=0,
                     stop=None,
                     step=None,
                     profile_plot=True,
                     streak_plot=True):
    """
    Plot data recorded by ProfileMonitor

        'ProfileMonitor' object.
        Dimension to plot. The default is "tau"
    start : int, optional
        First turn to plot. The default is 0.
    stop : int, optional
        Last turn to plot. If None, the last turn of the record is selected.
    step : int, optional
        Plotting step. This has to be divisible by 'save_every' parameter in
        'ProfileMonitor' object, i.e. step % save_every == 0. If None, step is
        equivalent to save_every.
    profile_plot : bool, optional
        If Ture, bunch profile plot is plotted.
    streak_plot : bool, optional
        If True, strek plot is plotted.

    Returns
    -------

    """

    path = file['ProfileData_{0}'.format(bunch_number)]
    l_bound = np.array(path["{0}_bin".format(dimension)])
    data = np.array(path[dimension])
    time = np.array(path["time"])

    if stop is None:
        stop = time[-1]
    elif stop not in time:
        raise ValueError("stop not found. Choose from {0}".format(time[:]))

    if start not in time:
        raise ValueError("start not found. Choose from {0}".format(time[:]))

    save_every = time[1] - time[0]

    if step is None:
        step = save_every

    if step % save_every != 0:
        raise ValueError("step must be divisible by the recording step "
                         "which is {0}.".format(save_every))

    dimension_dict = {"x": 0, "xp": 1, "y": 2, "yp": 3, "tau": 4, "delta": 5}
    scale = [1e6, 1e6, 1e6, 1e6, 1e12, 1]
    label = [
        "x (um)", "x' ($\\mu$rad)", "y (um)", "y' ($\\mu$rad)", "$\\tau$ (ps)",
    ]

    num = int((stop-start) / step)
    n_bin = len(data[:, 0])

    start_index = np.where(time[:] == start)[0][0]

    x_var = np.zeros((num + 1, n_bin))
    turn_index_array = np.zeros((num + 1, ), dtype=int)
    for i in range(num + 1):
        turn_index = int(start_index + i*step/save_every)
        turn_index_array[i] = turn_index
        # construct an array of bin mids
        x_var[i, :] = l_bound[:, turn_index]

    if profile_plot is True:
        for i in range(num + 1):
            ax.plot(x_var[i] * scale[dimension_dict[dimension]],
                    data[:, turn_index_array[i]],
                    label="turn {0}".format(time[turn_index_array[i]]))
        ax.set_xlabel(label[dimension_dict[dimension]])
        ax.set_ylabel("number of macro-particles")
        ax.legend()

    if streak_plot is True:
        turn = np.reshape(time[turn_index_array], (num + 1, 1))
        y_var = np.ones((num + 1, n_bin)) * turn
        z_var = np.transpose(data[:, turn_index_array])
        fig2, ax2 = plt.subplots()
        c = ax2.imshow(z_var,
                       cmap=cmap,
                       origin='lower',
                       aspect='auto',
                       extent=[
                           x_var.min() * scale[dimension_dict[dimension]],
                           x_var.max() * scale[dimension_dict[dimension]],
                           y_var.min(),
                           y_var.max()
                       ])
        ax2.set_xlabel(label[dimension_dict[dimension]])
        ax2.set_ylabel("Number of turns")
        cbar = fig2.colorbar(c, ax=ax2)
        cbar.set_label("Number of macro-particles")

    if profile_plot is True and streak_plot is True:
        return fig, fig2
    elif profile_plot is True:
        return fig
    elif streak_plot is True:
        return fig2

def plot_wakedata(filename,
                  bunch_number,
                  wake_type="Wlong",
                  start=0,
                  stop=None,
                  step=None,
                  profile_plot=False,
                  streak_plot=True,
                  bunch_profile=False,
                  dipole=False):
    """
    Plot data recorded by WakePotentialMonitor

        'WakePotentialMonitor' object.
    wake_type : str, optional
        Wake type to plot: "Wlong", "Wxdip", ... 
        'WakePotentialMonitor' object, i.e. step % save_every == 0. If None, 
        step is equivalent to save_every.
        If Ture, wake potential profile plot is plotted.
    bunch_profile : bool, optional.
        If True, the bunch profile is plotted.
    dipole : bool, optional
        If True, the dipole moment is plotted.

    Returns
    -------

    """

    path = file['WakePotentialData_{0}'.format(bunch_number)]

    dimension_dict = {
        "Wlong": 0,
        "Wxdip": 1,
        "Wydip": 2,
        "Wxquad": 3,
        "Wyquad": 4
    }
    scale = [1e-12, 1e-12, 1e-12, 1e-15, 1e-15]
    label = [
        "$W_p$ (V/pC)", "$W_{p,x}^D (V/pC)$", "$W_{p,y}^D (V/pC)$",
        "$W_{p,x}^Q (V/pC/mm)$", "$W_{p,y}^Q (V/pC/mm)$"
    ]

    if bunch_profile == True:
        tau_name = "tau_" + wake_type
        wake_type = "profile_" + wake_type
        dimension_dict = {wake_type: 0}
        scale = [1]
        label = ["$\\rho$ (a.u.)"]
    elif dipole == True:
        wake_type = "dipole_" + wake_type
        scale = [1]
        label = ["Dipole moment (m)"]
    else:

    data = np.array(path[wake_type])

        x_var[i, :] = np.array(path[tau_name])[:, turn_index]

            ax.plot(x_var[i] * 1e12,
                    data[:, turn_index_array[i]] *
                    scale[dimension_dict[wake_type]],
        ax.set_ylabel(label[dimension_dict[wake_type]])

        z_var = np.transpose(data[:, turn_index_array] *
                             scale[dimension_dict[wake_type]])
                       extent=[
                           x_var.min() * 1e12,
                           x_var.max() * 1e12,
                       ])
        ax2.set_xlabel("$\\tau$ (ps)")
        cbar.set_label(label[dimension_dict[wake_type]])

        return fig

def plot_bunchspectrum(filenames,
                       bunch_number,
                       dataset="incoherent",
                       dim="tau",
                       turns=None,
                       fs=None,
                       log_scale=True,
                       legend=None,
                       norm=False):
    """
    Plot coherent and incoherent spectrum data.

        'BunchSpectrumMonitor' object.
    dataset : {"mean_incoherent", "coherent", "incoherent"}
        HDF5 file's dataset to be plotted. 
        The default is "incoherent".
    dim :  {"x","y","tau"}, optional
        The dimension of the dataset to plot.
    turns : array or None, optional
        Numbers of the turns to plot.
        If None, all turns are shown. 
    fs : float or None, optional
        If not None, the frequency axis is noramlised by fs. 
    log_scale : bool, optional
        If True, the spectrum plots are shown in y-log scale. 
        The default is True.
    norm : bool, optional
        If True, normalise the data of each spectrum by its geometric mean.
        The default is False.

    Return
    ------

    """

        ll = []

        group = file["BunchSpectrum_{0}".format(bunch_number[i])]

        time = np.array(group["time"])
        freq = np.array(group["freq"])
        dim_dict = {"x": 0, "y": 1, "tau": 2}

        if dataset == "mean_incoherent":
            y_var = group["mean_incoherent"][dim_dict[dim], :]
            y_err = group["std_incoherent"][dim_dict[dim], :]
            ax.errorbar(time, y_var, y_err)
            xlabel = "Turn number"
            ylabel = "Mean incoherent frequency [Hz]"
        elif dataset == "incoherent" or dataset == "coherent":

            if turns is None:
                turn_index = np.where(time == time)[0]
            else:
                turn_index = []
                for turn in turns:
                    idx = np.where(time == turn)[0][0]
                    turn_index.append(idx)
                turn_index = np.array(turn_index)

            if fs is None:
                x_var = freq
                xlabel = "Frequency [Hz]"
            else:
                x_var = freq / fs
                xlabel = r"$f/f_{s}$"

            for idx in turn_index:
                y_var = group[dataset][dim_dict[dim], :, idx]
                if norm is True:
                    y_var = y_var / gmean(y_var)
                ax.plot(x_var, y_var)

            if log_scale is True:
                ax.set_yscale('log')

            ylabel = "FFT amplitude [a.u.]"
            if dataset == "incoherent":
                ax.set_title("Incoherent spectrum")
            elif dataset == "coherent":
                ax.set_title("Coherent spectrum")

        ax.set_ylabel(ylabel)

    return fig

def streak_bunchspectrum(filename,
                         bunch_number,
                         dataset="incoherent",
                         dim="tau",
                         fs=None,
                         log_scale=True,
                         fmin=None,
                         fmax=None,
                         turns=None,
                         norm=False,
                         ylim=None):
    """
    Plot 3D data recorded by the BunchSpectrumMonitor.

    filenames : str
        Name of the HDF5 file to be plotted.
    dataset : {"coherent", "incoherent"}
    fmin : float, optional
        If not None, the plot is limitted to values bigger than fmin.
    fmax : float, optional
        If not None, the plot is limitted to values smaller than fmax.
    turns : array, optional
        If not None, only the turn numbers in the turns array are plotted.
    ylim : array, optional
        If not None, should be array like in the form [ymin, ymax] where ymin 
        and ymax are the minimum and maxmimum values used in the y axis.
        
    Returns
    -------

    """

    group = file["BunchSpectrum_{0}".format(bunch_number)]

    time = np.array(group["time"])
    freq = np.array(group["freq"])
    dim_dict = {"x": 0, "y": 1, "tau": 2}

    if turns is None:
        turn_index = np.where(time == time)[0]
        if ylim is None:
            tmin = time.min()
            tmax = time.max()
        else:
            tmin = ylim[0]
            tmax = ylim[1]
    else:
        tmin = turns.min()
        tmax = turns.max()
        turn_index = []
        for turn in turns:
            idx = np.where(time == turn)[0][0]
            turn_index.append(idx)
        turn_index = np.array(turn_index)

    data = group[dataset][dim_dict[dim], :, turn_index]

    if log_scale is True:
        option = mpl.colors.LogNorm()
    else:
        option = None

    if fs is None:
        x_var = freq
        xlabel = "Frequency [Hz]"
    else:
        x_var = freq / fs
        xlabel = r"$f/f_{s}$"

    if fmin is None:
        fmin = x_var.min()
    if fmax is None:
        fmax = x_var.max()

    ind = (x_var > fmin) & (x_var < fmax)
    x_var = x_var[ind]
    data = data[ind, :]

    if norm is True:
        data = data / gmean(data)

    if ylim is None:
        ylabel = "Turn number"
    else:
        ylabel = ""

    if dataset == "incoherent":
        ax.set_title("Incoherent spectrum")
    elif dataset == "coherent":
        ax.set_title("Coherent spectrum")

    cmap = mpl.cm.inferno  # sequential
    c = ax.imshow(data.T,
                  extent=[x_var.min(), x_var.max(), tmin, tmax],
                  norm=option,
                  interpolation="none")
    cbar.ax.set_ylabel("FFT amplitude [a.u.]", rotation=270)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)

    return fig

def plot_beamspectrum(filenames,
                      dim="tau",
                      turns=None,
                      f0=None,
                      log_scale=True,
                      legend=None,
                      norm=False):
    """
    Plot coherent beam spectrum data.

    f0 : float or None, optional
        If not None, the frequency axis is noramlised by f0. 

    Return
    ------

    """

        group = file["BeamSpectrum"]

        dataset = "coherent"

        if turns is None:
            turn_index = np.where(time == time)[0]
        else:
            turn_index = []
            for turn in turns:
                idx = np.where(time == turn)[0][0]
                turn_index.append(idx)
            turn_index = np.array(turn_index)

        if f0 is None:
            x_var = freq
            xlabel = "Frequency [Hz]"
        else:
            x_var = freq / f0
            xlabel = r"$f/f_{0}$"

        for idx in turn_index:
            y_var = group[dataset][dim_dict[dim], :, idx]
            if norm is True:
                y_var = y_var / gmean(y_var)
            ax.plot(x_var, y_var)

        if log_scale is True:
            ax.set_yscale('log')

        ylabel = "FFT amplitude [a.u.]"
        ax.set_title("Beam coherent spectrum")

    return fig

def streak_beamspectrum(filename,
                        dim="tau",
                        f0=None,
                        log_scale=True,
                        fmin=None,
                        fmax=None,
                        turns=None,
                        norm=False,
                        ylim=None):
    """
    Plot 3D data recorded by the BeamSpectrumMonitor.

    Returns
    -------

    """

    group = file["BeamSpectrum"]
    dataset = "coherent"

        else:
    else:

    else:

    if f0 is None:
    else:
        x_var = freq / f0
        xlabel = r"$f/f_{0}$"

    else:
        ylabel = ""

    ax.set_title("Beam coherent spectrum")

    return fig

def plot_cavitydata(filename,
                    cavity_name,
                    phasor="cavity",
                    plot_type="bunch",
                    bunch_number=0,
                    turn=None,
                    cm_lim=None,
                    show_objective=False):
    """
    Plot data recorded by CavityMonitor.

    filename : str 
        Name of the CavityResonator object.
    phasor : str, optional
        Type of the phasor to plot. Can be "beam", "cavity", "generator" or 
        "ig".
    plot_type : str, optional
        Type of plot:
            - "bunch" plots the phasor voltage and angle versus time for a 
            given bunch.
            - "turn" plots the phasor voltage and ange versus bunch index for
            a given turn.
            - "streak_amplitude" plots the phasor amplitude versus bunch index 
            and time.
            - "streak_angle" plots the phasor angle versus bunch index and 
            time.
            - "detune" or "psi" plots the detuning or tuning angle versus time.
            - "power" plots the generator, cavity, beam and reflected power
            versus time. Needs BeamMonitor data with same save_every as 
            CavityMonitor.
            - "mean" plots the mean phasor voltage and angle versus time for 
            non empty bunches. Needs BeamMonitor data with same save_every as 
            - "tuner_diff" plots the tuner difference from optimal tuning 
            condition (for a given bunch) and the tuning angle versus time.
    bunch_number : int, optional
        Bunch number to select. The default is 0.
    turn : int, optional
        Turn to plot. The default is None.
        Colormap scale for the "streak" plots.
    show_objective : bool, optional
        If True, shows the Cavity voltage and phase objective value (Vc and 
        theta) in plot type "mean" and "bunch" for "cavity" phasor.
        Default is False.

    Returns
    -------

    """
    cavity_data = file[cavity_name]

    time = np.array(cavity_data["time"])

    ph = {"cavity": 0, "beam": 1, "generator": 2, "ig": 3}
    labels = ["Cavity", "Beam", "Generator", "Generator"]
    units = [" voltage [MV]", " voltage [MV]", " voltage [MV]", " current [A]"]
    units_val = [1e-6, 1e-6, 1e-6, 1]

    if (plot_type == "bunch") or (plot_type == "mean"):

        if plot_type == "bunch":
            data = [
                cavity_data["cavity_phasor_record"][bunch_number, :],
                cavity_data["beam_phasor_record"][bunch_number, :],
                cavity_data["generator_phasor_record"][bunch_number, :],
                cavity_data["ig_phasor_record"][bunch_number, :]
            ]
        elif plot_type == "mean":
            try:
                bunch_index = (file["Beam"]["current"][:, 0] != 0).nonzero()[0]
            except:
                ValueError("Beam monitor is needed to show mean voltage.")
            data = [
                np.mean(cavity_data["cavity_phasor_record"][bunch_index, :],
                        0),
                np.mean(cavity_data["beam_phasor_record"][bunch_index, :], 0),
                np.mean(cavity_data["generator_phasor_record"][bunch_index, :],
                        0),
                np.mean(cavity_data["ig_phasor_record"][bunch_index, :], 0)
            ]

        ylabel1 = labels[ph[phasor]] + units[ph[phasor]]
        ylabel2 = labels[ph[phasor]] + " phase [rad]"

        twin = ax.twinx()
        p1, = ax.plot(time,
                      np.abs(data[ph[phasor]]) * units_val[ph[phasor]],
                      color="r",
                      label=ylabel1)
        p2, = twin.plot(time,
                        np.angle(data[ph[phasor]]),
                        color="b",
                        label=ylabel2)
        if show_objective:
            o1, = ax.plot(time,
                          np.array(cavity_data["Vc"]) * 1e-6,
                          "r--",
                          label="Cavity voltage objective value [MV]")
            o2, = twin.plot(time,
                            np.array(cavity_data["theta"]),
                            "b--",
                            label="Cavity phase objective value [rad]")
        ax.set_xlabel("Turn number")
        ax.set_ylabel(ylabel1)
        twin.set_ylabel(ylabel2)

            plots = [p1, o1, p2, o2]
        else:
        ax.legend(handles=plots, loc="best")

        ax.yaxis.label.set_color("r")
        twin.yaxis.label.set_color("b")

    if plot_type == "turn":

        index = np.array(time) == turn
        if (index.size == 0):
            raise ValueError("Turn is not valid.")
        data = [
            cavity_data["cavity_phasor_record"][:, index],
            cavity_data["beam_phasor_record"][:, index],
            cavity_data["generator_phasor_record"][:, index],
            cavity_data["ig_phasor_record"][:, index]
        ]

        h = len(data[0])
        x = np.arange(h)

        p1, = ax.plot(x,
        p2, = twin.plot(x,
        ax.set_xlabel("Bunch index")

        plots = [p1, p2]

    if plot_type == "streak_amplitude" or plot_type == "streak_phase":

        data = [
            cavity_data["cavity_phasor_record"][:, :],
            cavity_data["beam_phasor_record"][:, :],
            cavity_data["generator_phasor_record"][:, :],
            cavity_data["ig_phasor_record"][:, :]
        ]

        if plot_type == "streak_amplitude":
            data = np.transpose(
                np.abs(data[ph[phasor]]) * units_val[ph[phasor]])
            ylabel = labels[ph[phasor]] + units[ph[phasor]]
            cmap = mpl.cm.coolwarm  # diverging
        elif plot_type == "streak_phase":
            data = np.transpose(np.angle(data[ph[phasor]]))
            ylabel = labels[ph[phasor]] + " phase [rad]"

        c = ax.imshow(data, cmap=cmap, origin='lower', aspect='auto')
        if cm_lim is not None:
            c.set_clim(vmin=cm_lim[0], vmax=cm_lim[1])
        ax.set_ylabel("Number of turns")
        cbar = fig.colorbar(c, ax=ax)
        cbar.set_label(ylabel)

    if plot_type == "detune" or plot_type == "psi":

        if plot_type == "detune":
            data = np.array(cavity_data["detune"]) * 1e-3
            ylabel = r"Detuning $\Delta f$ [kHz]"
        elif plot_type == "psi":
            data = np.array(cavity_data["psi"])
            ylabel = r"Tuning angle $\psi$"

        ax.plot(time, data)
        ax.set_xlabel("Number of turns")

    if plot_type == "power":
        Vc = np.mean(np.abs(cavity_data["cavity_phasor_record"]), 0)
        theta = np.mean(np.angle(cavity_data["cavity_phasor_record"]), 0)
        try:
            I0 = np.nansum(file["Beam"]["current"][bunch_index, :], 0)
        except:
            print("Beam monitor is needed to compute power.")

        Rs = np.array(cavity_data["Rs"])
        Pc = Vc**2 / (2*Rs)
        Pb = I0 * Vc * np.cos(theta)
        Pg = np.array(cavity_data["Pg"])
        Pr = Pg - Pb - Pc

        ax.plot(time, Pg * 1e-3, label="Generator power $P_g$ [kW]")
        ax.plot(time, Pb * 1e-3, label="Beam power $P_b$ [kW]")
        ax.plot(time, Pc * 1e-3, label="Dissipated cavity power $P_c$ [kW]")
        ax.plot(time, Pr * 1e-3, label="Reflected power $P_r$ [kW]")
        ax.set_ylabel("Power [kW]")
        plt.legend()

    if plot_type == "tuner_diff":
        data0 = np.angle(cavity_data["cavity_phasor_record"][bunch_number, :])
        data1 = np.array(cavity_data["psi"])
        data2 = np.array(cavity_data["theta_g"])

        ylabel1 = "Tuner diff. from optimal [rad]"
        ylabel2 = r"Tuning angle $\psi$ [rad]"
        p1, = ax.plot(time, data0 - data2 + data1, color="r", label=ylabel1)
        p2, = twin.plot(time, data1, color="b", label=ylabel2)

    return fig

"""
This module defines utilities functions, helping to deals with tracking output 
and hdf5 files.
"""

def merge_files(files_prefix, files_number, start_idx=0, file_name=None):
    """
    Merge several hdf5 files into one.
    
    The function assumes that the files to merge have names in the follwing 
    format:
        - "files_prefix_0.hdf5"
        - "files_prefix_1.hdf5"
        ...
        - "files_prefix_files_number.hdf5"

    files_prefix : str
        Name of the files to merge.
    files_number : int
        Number of files to merge.
    start_idx : int, optional
        Start index of the hdf5 files.        
    file_name : str, optional
        Name of the file with the merged data. If None, files_prefix without
        number is used.

    """
    if file_name == None:
        file_name = files_prefix
    f = hp.File(file_name + ".hdf5", "a")

    ## Create file architecture
    f0 = hp.File(files_prefix + "_" + str(start_idx) + ".hdf5", "r")
    for group in list(f0):
        f.require_group(group)
        for dataset_name in list(f0[group]):
            if dataset_name == "freq":
                f0[group].copy(dataset_name, f[group])
                continue
            shape = f0[group][dataset_name].shape
            dtype = f0[group][dataset_name].dtype
            shape_needed = list(shape)
            shape_needed[-1] = shape_needed[-1] * files_number
            shape_needed = tuple(shape_needed)
            f[group].create_dataset(dataset_name, shape_needed, dtype)

    f0.close()

    ## Copy data
    for i, file_num in enumerate(range(start_idx, start_idx + files_number)):
        fi = hp.File(files_prefix + "_" + str(file_num) + ".hdf5", "r")
        for group in list(fi):
            for dataset_name in list(fi[group]):
                shape = fi[group][dataset_name].shape
                n_slice = int(len(shape) - 1)
                length = shape[-1]
                slice_list = []
                for n in range(n_slice):
                    slice_list.append(slice(None))
                slice_list.append(slice(length * i, length * (i+1)))
                if (dataset_name == "freq"):
                    continue
                if (dataset_name == "time") and (file_num != start_idx):
                    f[group][dataset_name][tuple(slice_list)] = np.max(
                        f[group][dataset_name][:]) + fi[group][dataset_name]
                else:
                    f[group][dataset_name][tuple(
                        slice_list)] = fi[group][dataset_name]
        fi.close()
    f.close()

def copy_files(source, copy, version=None):
    """
    Copy a source hdf5 file into another hdf5 file using a different HDF5 
    version.
    
    The function assumes that the source file has only a single group layer.

    source : str
        Name of the source file.
    copy : str
        Name of the copy file.
    version : str, optional
        Version number of the copy file.

    """
    if version == None:
        version = 'v108'
    f = hp.File(source + ".hdf5", "r")
    h = hp.File(copy + ".hdf5", "a", libver=('earliest', version))

    ## Copy file
    for group in list(f):
        h.require_group(group)
        for dataset_name in list(f[group]):
            h[group][dataset_name] = f[group][dataset_name][()]

    f.close()
    h.close()

"""
Module where the BeamLoadingEquilibrium class is defined.
"""

from scipy.optimize import root

class BeamLoadingEquilibrium():
    """Class used to compute beam equilibrium profile for a given storage ring 
    and a list of RF cavities of any harmonic. The class assumes an uniform 
    filling of the storage ring. Based on an extension of [1].

    [1] Venturini, M. (2018). Passive higher-harmonic rf cavities with general
    settings and multibunch instabilities in electron storage rings.
    Physical Review Accelerators and Beams, 21(11), 114404.

    cavity_list : list of CavityResonator objects
    I0 : beam current in [A].
    auto_set_MC_theta : if True, allow class to change cavity phase for
        CavityResonator objetcs with m = 1 (i.e. main cavities)
    F : list of form factor amplitude
    PHI : list of form factor phase
    B1 : lower intergration boundary
    B2 : upper intergration boundary
    """

                 ring,
                 cavity_list,
                 I0,
                 auto_set_MC_theta=False,
                 F=None,
                 PHI=None,
                 B1=-0.2,
                 B2=0.2):
        self.cavity_list = cavity_list
        self.I0 = I0
        self.n_cavity = len(cavity_list)
        self.auto_set_MC_theta = auto_set_MC_theta
        if F is None:
            self.F = np.ones((self.n_cavity, ))
        else:
            self.F = F
        if PHI is None:
            self.PHI = np.zeros((self.n_cavity, ))
        else:
            self.PHI = PHI
        self.B1 = B1
        self.B2 = B2
        self.mpi = False
        self.__version__ = "1.0"

        # Define constants for scaled potential u(z)
        self.u0 = self.ring.U0 / (self.ring.ac * self.ring.sigma_delta**2 *
                                  self.ring.E0 * self.ring.L)
        self.ug = np.zeros((self.n_cavity, ))
        self.ub = np.zeros((self.n_cavity, ))
        self.update_potentials()

    def update_potentials(self):
        """Update potentials with cavity and ring data."""
        for i in range(self.n_cavity):
            cavity = self.cavity_list[i]
            self.ug[i] = cavity.Vg / (self.ring.ac * self.ring.sigma_delta**2 *
                                      self.ring.E0 * self.ring.L *
                                      self.ring.k1 * cavity.m)
            self.ub[i] = 2 * self.I0 * cavity.Rs / (
                self.ring.ac * self.ring.sigma_delta**2 * self.ring.E0 *
                self.ring.L * self.ring.k1 * cavity.m * (1 + cavity.beta))

    def energy_balance(self):
        """Return energy balance for the synchronous particle
        (z = 0 ,delta = 0)."""
        delta = self.ring.U0
            delta += cavity.Vb(
                self.I0) * self.F[i] * np.cos(cavity.psi - self.PHI[i])
            delta -= cavity.Vg * np.cos(cavity.theta_g)
        return delta

    def center_of_mass(self):
        """Return center of mass position in [s]"""
        z0 = np.linspace(self.B1, self.B2, 1000)
        rho = self.rho(z0)
        CM = np.average(z0, weights=rho)
        return CM / c

    def u(self, z):
        """Scaled potential u(z)"""
        pot = self.u0 * z
            pot += -self.ug[i] * (
                np.sin(self.ring.k1 * cavity.m * z + cavity.theta_g) -
                np.sin(cavity.theta_g))
            pot += self.ub[i] * self.F[i] * np.cos(cavity.psi) * (
                np.sin(self.ring.k1 * cavity.m * z + cavity.psi - self.PHI[i])
                - np.sin(cavity.psi - self.PHI[i]))
        return pot

    def du_dz(self, z):
        """Partial derivative of the scaled potential u(z) by z"""
        pot = self.u0
            pot += -self.ug[i] * self.ring.k1 * cavity.m * np.cos(
                self.ring.k1 * cavity.m * z + cavity.theta_g)
            pot += self.ub[i] * self.F[i] * self.ring.k1 * cavity.m * np.cos(
                cavity.psi) * np.cos(self.ring.k1 * cavity.m * z + cavity.psi -
                                     self.PHI[i])
        return pot

    def uexp(self, z):
        return np.exp(-1 * self.u(z))

    def integrate_func(self, f, g):
        """Return Integral[f*g]/Integral[f] between B1 and B2"""
        A = quad(lambda x: f(x) * g(x), self.B1, self.B2)
        B = quad(f, self.B1, self.B2)
        return A[0] / B[0]

    def to_solve(self, x, CM=True):
        """System of non-linear equation to solve to find the form factors F
        and PHI at equilibrum.
        The system is composed of Eq. (B6) and (B7) of [1] for each cavity.
        If auto_set_MC_theta == True, the system also find the main cavity 
        phase to impose energy balance or cancel center of mass offset.
        If CM is True, the system imposes zero center of mass offset,
        if False, the system imposes energy balance.
        """
        # Update values of F, PHI and theta
        if self.auto_set_MC_theta:
            self.F = x[:-1:2]
            for i in range(self.n_cavity):
                cavity = self.cavity_list[i]
                if cavity.m == 1:
                    cavity.theta = x[-1]
                    cavity.set_generator(self.I0)
                    self.update_potentials()
        else:
            self.F = x[::2]
        self.PHI = x[1::2]

        # Compute system
            res = np.zeros((self.n_cavity * 2 + 1, ))
                res[2 *
                    i] = self.F[i] * np.cos(self.PHI[i]) - self.integrate_func(
                        lambda y: self.uexp(y),
                        lambda y: np.cos(self.ring.k1 * cavity.m * y))
                res[2*i +
                    1] = self.F[i] * np.sin(self.PHI[i]) - self.integrate_func(
                        lambda y: np.sin(self.ring.k1 * cavity.m * y))
            # Factor 1e-8 or 1e12 for better convergence
            if CM is True:
                res[self.n_cavity * 2] = self.center_of_mass() * 1e12
            else:
                res[self.n_cavity * 2] = self.energy_balance() * 1e-8
        else:
            res = np.zeros((self.n_cavity * 2, ))
                res[2 *
                res[2*i +
        return res

    def rho(self, z):
        """Return bunch equilibrium profile at postion z"""
        A = quad(lambda y: self.uexp(y), self.B1, self.B2)
        return self.uexp(z) / A[0]

    def plot_rho(self, z1=None, z2=None):
        """Plot the bunch equilibrium profile between z1 and z2"""
        if z1 is None:
            z1 = self.B1
        if z2 is None:
            z2 = self.B2
        z0 = np.linspace(z1, z2, 1000)
        plt.plot(z0 / c * 1e12, self.rho(z0))
        plt.xlabel(r"$\tau$ [ps]")
        plt.title("Equilibrium bunch profile")

    def voltage(self, z):
        """Return the RF system total voltage at position z"""
        Vtot = 0
            Vtot += cavity.VRF(z, self.I0, self.F[i], self.PHI[i])
        return Vtot

    def dV(self, z):
        """Return derivative of the RF system total voltage at position z"""
        Vtot = 0
            Vtot += cavity.dVRF(z, self.I0, self.F[i], self.PHI[i])

    def ddV(self, z):
        """Return the second derivative of the RF system total voltage at position z"""
        Vtot = 0
            Vtot += cavity.ddVRF(z, self.I0, self.F[i], self.PHI[i])

    def deltaVRF(self, z):
        """Return the generator voltage minus beam loading voltage of the total RF system at position z"""
        Vtot = 0
            Vtot += cavity.deltaVRF(z, self.I0, self.F[i], self.PHI[i])

    def plot_dV(self, z1=None, z2=None):
        """Plot the derivative of RF system total voltage between z1 and z2"""
        plt.plot(z0, self.dV(z0))
        plt.xlabel("z [m]")
        plt.ylabel("Total RF voltage (V)")

    def plot_voltage(self, z1=None, z2=None):
        """Plot the RF system total voltage between z1 and z2"""
        plt.plot(z0, self.voltage(z0))

    def std_rho(self, z1=None, z2=None):
        """Return the rms bunch equilibrium size in [m]"""
        values = self.rho(z0)
        average = np.average(z0, weights=values)
        variance = np.average((z0 - average)**2, weights=values)
        return np.sqrt(variance)

    def beam_equilibrium(self,
                         x0=None,
                         tol=1e-4,
                         method='hybr',
                         options=None,
                         plot=False,
                         CM=True,
                         verbiose=False):
        """Solve system of non-linear equation to find the form factors F
        and PHI at equilibrum. Can be used to compute the equilibrium bunch
        profile.
        
        x0 : initial guess
        tol : tolerance for termination of the algorithm
        method : method used by scipy.optimize.root to solve the system
        options : options given to scipy.optimize.root
        plot : if True, plot the equilibrium bunch profile
        CM : if True, the system imposes zero center of mass offset,
        verbiose : if True, print out informations about convergence.
        
        Returns
        -------
        sol : OptimizeResult object representing the solution
        """
        if x0 is None:
            x0 = [1, 0] * self.n_cavity
            if self.auto_set_MC_theta:
                x0 = x0 + [self.cavity_list[0].theta]

        if verbiose:
            if CM:
                print("The initial center of mass offset is " +
                      str(self.center_of_mass() * 1e12) + " ps")
            else:
                print("The initial energy balance is " +
                      str(self.energy_balance()) + " eV")

        sol = root(lambda x: self.to_solve(x, CM),
                   x0,
                   tol=tol,
                   method=method,
                   options=options)

        # Update values of F, PHI and theta_g
            self.F = sol.x[:-1:2]
                    cavity.theta = sol.x[-1]
        else:
            self.F = sol.x[::2]
        self.PHI = sol.x[1::2]

            if CM:
                print("The final center of mass offset is " +
            else:
                print("The final energy balance is " +
            if not sol.success:
                print("The algorithm has converged: " + str(sol.success))

        if plot:
            self.plot_rho(self.B1 / 4, self.B2 / 4)

        return sol

    def PTBL_threshold(self, I0, m=None, MC_index=0, HC_index=1):
        """
        Compute the periodic transient beam loading (PTBL) instability 
        threshold based on [1].
        
        The beam_equilibrium method should be called before the PTBL_threshold 
        one.
        
        Eq. (22) and (23) have been modified for cosine voltage convention 
        for the main cavity.

        I0 : float
            Beam total current in [A].
        m : int, optional
            Number of bunches in the beam. 
            The default is None. If None, uniform filling is assumed.
        MC_index : int, optional
            Index of the main cavity in cavity_list. The default is 0.
        HC_index : int, optional
            Index of the harmonic cavity in cavity_list. The default is 1.

        Returns
        -------
        eta : float
            Amplification factor, if bigger than 1 then the beam is possibly 
            in the PTBL regime, Eq. (22) of [1].
        RQth : float
            R/Q threshold, Eq. (24) of [1].
        f : float
            f function, Eq. (23) of [1].
            
        Reference
        ---------
        [1] : He, T. (2022). Novel perturbation method for judging the 
        stability of the equilibrium solution in the presence of passive 
        harmonic cavities. Physical Review Accelerators and Beams, 25(9), 
        094402.

        """
        if m is None:
            m = self.ring.h

        MC = self.cavity_list[MC_index]
        HC = self.cavity_list[HC_index]
        a = np.exp(-HC.wr * self.ring.T0 / (2 * HC.Q))
        theta = HC.detune * 2 * np.pi * self.ring.T0
        dtheta = np.arcsin((1-a) * np.cos(theta / 2) /
                           (np.sqrt(1 + a**2 - 2 * a * np.cos(theta))))

        k = np.arange(1, m)
        d_k = np.exp(-1 * HC.wr * self.ring.T0 * (k-1) / (2 * HC.Q * m))
        theta_k = (theta/2 + dtheta - (k-1) / m * theta)
        eps_k = 1 - np.sin(np.pi / 2 - k * 2 * np.pi / m)

        num = np.sum(eps_k * d_k * np.cos(theta_k))
        f = num / (m * np.sqrt(1 + a**2 - 2 * a * np.cos(theta)))

        eta = (2 * np.pi * HC.m**2 * self.F[HC_index] * self.ring.h * I0 *
               HC.Rs / HC.Q * f /
               (MC.Vc * np.sin(MC.theta - self.PHI[HC_index] / HC.m)))

        RQth = (MC.Vc * np.sin(MC.theta - self.PHI[HC_index] / HC.m) /
                (2 * np.pi * HC.m**2 * self.F[1] * self.ring.h * I0 * f))

        return (eta, RQth, f)

"""
This module defines miscellaneous utilities functions.
"""

from pathlib import Path

from mbtrack2.impedance.wakefield import Impedance
from mbtrack2.utilities.spectrum import spectral_density

def effective_impedance(ring,
                        imp,
                        m,
                        mu,
                        sigma,
                        M,
                        tuneS,
                        xi=None,
                        mode="Hermite"):
    """
    Compute the effective (longitudinal or transverse) impedance. 
    Formulas from Eq. (1) and (2) p238 of [1].

    imp : Impedance object
    mu : int
        coupled bunch mode number, goes from 0 to (M-1) where M is the
        number of bunches
    m : int
        head-tail (or azimutal/synchrotron) mode number
    M : int
    tuneS : float
    xi : float, optional
        (non-normalized) chromaticity
    mode: str, optional
        type of the mode taken into account for the computation:
        -"Hermite" modes for Gaussian bunches

    Returns
    -------
    Zeff : float 
        effective impedance in [ohm] or in [ohm/m] depanding on the impedance
        type.

    [1] : Handbook of accelerator physics and engineering, 3rd printing.
    """

    if not isinstance(imp, Impedance):
        raise TypeError("{} should be an Impedance object.".format(imp))

    fmin = imp.data.index.min()
    fmax = imp.data.index.max()
    if fmin > 0:
        double_sided_impedance(imp)

    if mode in ["Hermite", "Legendre", "Sinusoidal", "Sacherer", "Chebyshev"]:

        def h(f):
            return spectral_density(frequency=f, sigma=sigma, m=m, mode=mode)
    else:
        raise NotImplementedError("Not implemanted yet.")

    if imp.component_type == "long":
        pmax = fmax / (ring.f0 * M) - 1
        pmin = fmin / (ring.f0 * M) + 1
        p = np.arange(pmin, pmax + 1)

        fp = ring.f0 * (p*M + mu + m*tuneS)
        fp = fp[np.nonzero(fp)]  # Avoid division by 0
        num = np.sum(imp(fp) * h(fp) / (fp * 2 * np.pi))
        den = np.sum(h(fp))
        Zeff = num / den

    elif imp.component_type == "xdip" or imp.component_type == "ydip":
        if imp.component_type == "xdip":
            tuneXY = ring.tune[0] - np.floor(ring.tune[0])
            if xi is None:
                xi = ring.chro[0]
        elif imp.component_type == "ydip":
            tuneXY = ring.tune[1] - np.floor(ring.tune[1])
                xi = ring.chro[1]
        fp = ring.f0 * (p*M + mu + tuneXY + m*tuneS)
        f_xi = xi / ring.eta() * ring.f0
        num = np.sum(imp(fp) * h(fp - f_xi))
        den = np.sum(h(fp - f_xi))
    else:
        raise TypeError("Effective impedance is only defined for long, xdip"
                        " and ydip impedance type.")

    return Zeff

def head_tail_form_factor(ring,
                          imp,
                          m,
                          sigma,
                          tuneS,
                          xi=None,
                          mode="Hermite",
                          mu=0):
    M = 1

    if fmin > 0:

        def h(f):
    else:

    pmax = np.floor(fmax / (ring.f0 * M))
    pmin = np.ceil(fmin / (ring.f0 * M))

    p = np.arange(pmin, pmax + 1)

            tuneXY = ring.tune[1] - np.floor(ring.tune[0])
    else:

    return den

def tune_shift_from_effective_impedance(Zeff):
    pass

def yokoya_elliptic(x_radius, y_radius):
    """
    Compute Yokoya factors for an elliptic beam pipe.
    Function adapted from N. Mounet IW2D.

    x_radius : float
        Horizontal semi-axis of the ellipse in [m].
    y_radius : float
        Vertical semi-axis of the ellipse in [m].

    Returns
    -------
    yoklong : float
        Yokoya factor for the longitudinal impedance.
    yokxdip : float
        Yokoya factor for the dipolar horizontal impedance.
    yokydip : float
        Yokoya factor for the dipolar vertical impedance.
    yokxquad : float
        Yokoya factor for the quadrupolar horizontal impedance.
    yokyquad : float
        Yokoya factor for the quadrupolar vertical impedance.
    """
    if y_radius < x_radius:
        small_semiaxis = y_radius
        large_semiaxis = x_radius
    else:
        small_semiaxis = x_radius
        large_semiaxis = y_radius

    path_to_file = Path(__file__).parent
    file = path_to_file / "data" / "Yokoya_elliptic_from_Elias_USPAS.csv"

    # read Yokoya factors interpolation file
    # BEWARE: columns are ratio, dipy, dipx, quady, quadx
    yokoya_file = pd.read_csv(file)
    ratio_col = yokoya_file["x"]
    # compute semi-axes ratio (first column of this file)
    ratio = (large_semiaxis-small_semiaxis) / (large_semiaxis+small_semiaxis)

    # interpolate Yokoya file at the correct ratio
    yoklong = 1

        yokydip = np.interp(ratio, ratio_col, yokoya_file["dipy"])
        yokxdip = np.interp(ratio, ratio_col, yokoya_file["dipx"])
        yokyquad = np.interp(ratio, ratio_col, yokoya_file["quady"])
        yokxquad = np.interp(ratio, ratio_col, yokoya_file["quadx"])
    else:
        yokxdip = np.interp(ratio, ratio_col, yokoya_file["dipy"])
        yokydip = np.interp(ratio, ratio_col, yokoya_file["dipx"])
        yokxquad = np.interp(ratio, ratio_col, yokoya_file["quady"])
        yokyquad = np.interp(ratio, ratio_col, yokoya_file["quadx"])

    return (yoklong, yokxdip, yokydip, yokxquad, yokyquad)

def beam_loss_factor(impedance, frequency, spectrum, ring):
    """
    Compute "beam" loss factor using the beam spectrum, uses a sum instead of 
    integral compared to loss_factor [1].

    impedance : Impedance of type "long"
    frequency : array
        Sample points of spectrum.
    spectrum : array
        Beam spectrum to consider.

    Returns
    -------
    kloss_beam : float
        Beam loss factor in [V/C].

    [1] : Handbook of accelerator physics and engineering, 3rd printing. 
        Eq (3) p239.
    """
    pmax = np.floor(impedance.data.index.max() / ring.f0)
    pmin = np.floor(impedance.data.index.min() / ring.f0)

    if pmin >= 0:
        double_sided_impedance(impedance)
        pmin = -1 * pmax

    p = np.arange(pmin + 1, pmax)
    pf0 = p * ring.f0
    ReZ = np.real(impedance(pf0))
    spectral_density = np.abs(spectrum)**2
    # interpolation of the spectrum is needed to avoid problems liked to
    # division by 0
    # computing the spectrum directly to the frequency points gives
    # wrong results
    spect = interp1d(frequency, spectral_density)
    kloss_beam = ring.f0 * np.sum(ReZ * spect(pf0))

    return kloss_beam

def double_sided_impedance(impedance):
    """
    Add negative frequency points to single sided impedance spectrum following
    symetries depending on impedance type.

    impedance : Impedance object
        Single sided impedance.
    """
    fmin = impedance.data.index.min()

    if fmin >= 0:
        negative_index = impedance.data.index * -1
        negative_data = impedance.data.set_index(negative_index)

        imp_type = impedance.component_type

        if imp_type == "long":
            negative_data["imag"] = -1 * negative_data["imag"]

        elif (imp_type == "xdip") or (imp_type == "ydip"):
            negative_data["real"] = -1 * negative_data["real"]

        elif (imp_type == "xquad") or (imp_type == "yquad"):

        else:
            raise ValueError("Wrong impedance type")

        try:
            negative_data = negative_data.drop(0)
            pass

        all_data = pd.concat([impedance.data, negative_data])
        all_data = all_data.sort_index()
        impedance.data = all_data

"""
Module where the class to store the optic functions and the lattice physical 
parameters are defined.
"""

class Optics:
    """
    Class used to handle optic functions.
    
    lattice_file : str, optional if local_beta, local_alpha and 
        local_dispersion are specified.
        AT lattice file path.
    local_beta : array of shape (2,), optional if lattice_file is specified.
        Beta function at the location of the tracking. Default is mean beta if
        lattice has been loaded.
    local_alpha : array of shape (2,), optional if lattice_file is specified.
        Alpha function at the location of the tracking. Default is mean alpha 
        if lattice has been loaded.
    local_dispersion : array of shape (4,), optional if lattice_file is 
        specified.
        Dispersion function and derivative at the location of the tracking. 
        Default is zero if lattice has been loaded.  
        
    use_local_values : bool
        True if no lattice has been loaded.
    local_gamma : array of shape (2,)
        Gamma function at the location of the tracking.
    lattice : AT lattice
        
    Methods
    -------
    load_from_AT(lattice_file, **kwargs)
        Load a lattice from accelerator toolbox (AT).
    setup_interpolation()
        Setup interpolation of the optic functions.
    beta(position)
        Return beta functions at specific locations given by position.
    alpha(position)
        Return alpha functions at specific locations given by position.
    gamma(position)
        Return gamma functions at specific locations given by position.
    dispersion(position)
        Return dispersion functions at specific locations given by position.
    plot(self, var, option, n_points=1000)
        Plot optical variables.
    """

                 lattice_file=None,
                 local_beta=None,
                 local_alpha=None,
                 local_dispersion=None,

        if lattice_file is not None:
            self.use_local_values = False
            self.load_from_AT(lattice_file, **kwargs)
            if local_beta is None:
                self._local_beta = np.mean(self.beta_array, axis=1)
            else:
                self._local_beta = local_beta
            if local_alpha is None:
                self._local_alpha = np.mean(self.alpha_array, axis=1)
            else:
                self._local_alpha = local_alpha
            if local_dispersion is None:
                self.local_dispersion = np.zeros((4, ))
            else:
                self.local_dispersion = local_dispersion
            self._local_gamma = (1 + self._local_alpha**2) / self._local_beta

        else:
            self.use_local_values = True
            self._local_beta = local_beta
            self._local_alpha = local_alpha
            self.local_dispersion = local_dispersion

    def load_from_AT(self, lattice_file, **kwargs):
        """

        lattice_file : str
            AT lattice file path.
            Minimum number of points to use for the optic function arrays.
        periodicity : int, optional
            Lattice periodicity, if not specified the AT lattice periodicity is
            used.
        """
        import at
        self.n_points = int(kwargs.get("n_points", 1e3))
        periodicity = kwargs.get("periodicity")

        self.lattice = at.load_lattice(lattice_file)
        if self.lattice.radiation:
            self.lattice.radiation_off()
        lattice = self.lattice.slice(slices=self.n_points)
        refpts = np.arange(0, len(lattice))
        twiss0, tune, chrom, twiss = at.linopt(lattice,
                                               refpts=refpts,
                                               get_chrom=True)

        if periodicity is None:
            self.periodicity = lattice.periodicity
        else:
            self.periodicity = periodicity

        if self.periodicity > 1:
            for i in range(self.periodicity - 1):
                pos = np.append(twiss.s_pos,
                                twiss.s_pos + twiss.s_pos[-1] * (i+1))
        else:
            pos = twiss.s_pos

        self.position = pos
        self.beta_array = np.tile(twiss.beta.T, self.periodicity)
        self.alpha_array = np.tile(twiss.alpha.T, self.periodicity)
        self.dispersion_array = np.tile(twiss.dispersion.T, self.periodicity)
        self.mu_array = np.tile(twiss.mu.T, self.periodicity)

        self.position = np.append(self.position, self.lattice.circumference)
        self.beta_array = np.append(self.beta_array,
                                    self.beta_array[:, 0:1],
                                    axis=1)
        self.alpha_array = np.append(self.alpha_array,
                                     self.alpha_array[:, 0:1],
                                     axis=1)
        self.dispersion_array = np.append(self.dispersion_array,
                                          self.dispersion_array[:, 0:1],
                                          axis=1)
        self.mu_array = np.append(self.mu_array, self.mu_array[:, 0:1], axis=1)

        self.gamma_array = (1 + self.alpha_array**2) / self.beta_array
        self.tune = tune * self.periodicity
        self.chro = chrom * self.periodicity
        self.ac = at.get_mcf(self.lattice)

        self.mu_array[:, -1] = (np.floor(self.mu_array[:, -2] /
                                         (2 * np.pi)) + self.tune) * 2 * np.pi

        self.setup_interpolation()

    def setup_interpolation(self):
        """Setup interpolation of the optic functions."""
        self.betaX = interp1d(self.position,
                              self.beta_array[0, :],
                              kind='linear')
        self.betaY = interp1d(self.position,
                              self.beta_array[1, :],
        self.alphaX = interp1d(self.position,
                               self.alpha_array[0, :],
                               kind='linear')
        self.alphaY = interp1d(self.position,
                               self.alpha_array[1, :],
        self.gammaX = interp1d(self.position,
                               self.gamma_array[0, :],
        self.gammaY = interp1d(self.position,
                               self.gamma_array[1, :],
        self.dispX = interp1d(self.position,
                              self.dispersion_array[0, :],
        self.disppX = interp1d(self.position,
                               self.dispersion_array[1, :],
        self.dispY = interp1d(self.position,
                              self.dispersion_array[2, :],
        self.disppY = interp1d(self.position,
                               self.dispersion_array[3, :],
        self.muX = interp1d(self.position, self.mu_array[0, :], kind='linear')
        self.muY = interp1d(self.position, self.mu_array[1, :], kind='linear')

    @property
    def local_beta(self):
        """
        Return beta function at the location defined by the lattice file.

        """
        return self._local_beta

    @local_beta.setter
    def local_beta(self, beta_array):
        """
        Set the values of beta function. Gamma function is automatically 
        recalculated after the new value of beta function is set.

        beta_array : array of shape (2,)
            Beta function in the horizontal and vertical plane.

        """
        self._local_beta = beta_array
        self._local_gamma = (1 + self._local_alpha**2) / self._local_beta

    @property
    def local_alpha(self):
        """
        Return alpha function at the location defined by the lattice file.

        """
        return self._local_alpha

    @local_alpha.setter
    def local_alpha(self, alpha_array):
        """
        Set the value of beta functions. Gamma function is automatically 
        recalculated after the new value of alpha function is set.

        alpha_array : array of shape (2,)
            Alpha function in the horizontal and vertical plane.

        """
        self._local_alpha = alpha_array

    @property
    def local_gamma(self):
        """

        """
        return self._local_gamma

    def beta(self, position):
        """
        Return beta functions at specific locations given by position. If no
        lattice has been loaded, local values are returned.

        position : array or float
            Longitudinal position at which the beta functions are returned.

        Returns
        -------
        beta : array
            Beta functions.
        """
        if self.use_local_values:
            return np.outer(self.local_beta, np.ones_like(position))
        else:
            beta = [self.betaX(position), self.betaY(position)]
            return np.array(beta)

    def alpha(self, position):
        """
        Return alpha functions at specific locations given by position. If no

            Longitudinal position at which the alpha functions are returned.

        Returns
        -------
        alpha : array
            Alpha functions.
        """
            return np.outer(self.local_alpha, np.ones_like(position))
        else:
            alpha = [self.alphaX(position), self.alphaY(position)]
            return np.array(alpha)

    def gamma(self, position):
        """
        Return gamma functions at specific locations given by position. If no

            Longitudinal position at which the gamma functions are returned.

        Returns
        -------
        gamma : array
            Gamma functions.
        """
            return np.outer(self.local_gamma, np.ones_like(position))
        else:
            gamma = [self.gammaX(position), self.gammaY(position)]
            return np.array(gamma)

    def dispersion(self, position):
        """
        Return dispersion functions at specific locations given by position. 
        If no lattice has been loaded, local values are returned.

            Longitudinal position at which the dispersion functions are 
            returned.

        Returns
        -------
        dispersion : array
            Dispersion functions.
        """
            return np.outer(self.local_dispersion, np.ones_like(position))
        else:
            dispersion = [
                self.dispX(position),
                self.disppX(position),
                self.dispY(position),
                self.disppY(position)
            ]
            return np.array(dispersion)

    def mu(self, position):
        """
        Return phase advances at specific locations given by position. 
        If no lattice has been loaded, None is returned.

            Longitudinal position at which the phase advances are returned.

        Returns
        -------
        mu : array
            Phase advances.
        """
            return np.outer(np.array([0, 0]), np.ones_like(position))
        else:
            mu = [self.muX(position), self.muY(position)]
            return np.array(mu)

    def plot(self, var, option, n_points=1000):
        """
    
        var : {"beta", "alpha", "gamma", "dispersion", "mu"}
            Optical variable to be plotted.
        option : str
            If var = "beta", "alpha" and "gamma", option = {"x","y"} specifying
            the axis of interest.
            If var = "dispersion", option = {"x","px","y","py"} specifying the 
            axis of interest for the dispersion function or its derivative.
        n_points : int
            Number of points on the plot. The default value is 1000.
    
        """

            "beta": self.beta,
            "alpha": self.alpha,
            "gamma": self.gamma,
            "dispersion": self.dispersion,
            "mu": self.mu
        }

        if var == "dispersion":
            option_dict = {"x": 0, "px": 1, "y": 2, "py": 3}

            label = ["D$_{x}$ (m)", "D'$_{x}$", "D$_{y}$ (m)", "D'$_{y}$"]

            ylabel = label[option_dict[option]]

        elif var == "beta" or var == "alpha" or var == "gamma" or var == "mu":
            option_dict = {"x": 0, "y": 1}
            label_dict = {
                "beta": "$\\beta$",
                "alpha": "$\\alpha$",
                "gamma": "$\\gamma$",
                "mu": "$\\mu$"
            }

            if option == "x": label_sup = "$_{x}$"
            elif option == "y": label_sup = "$_{y}$"

            unit = {
                "beta": " (m)",
                "alpha": "",
                "gamma": " (m$^{-1}$)",
                "mu": ""
            }

            ylabel = label_dict[var] + label_sup + unit[var]

        else:
            raise ValueError("Variable name is not found.")

        if self.use_local_values is not True:
            position = np.linspace(0, self.lattice.circumference,
                                   int(n_points))
        else:
            position = np.linspace(0, 1)

        var_list = var_dict[var](position)[option_dict[option]]
        ax.plot(position, var_list)

        ax.set_xlabel("position (m)")

        return fig

class PhysicalModel:
    """
    Store the lattice physical parameters such as apperture and resistivity.
    
    x_right : float
        Default value for the right horizontal aperture in [m].
    y_top : float
        Default value for the top vertical aperture in [m].
    shape : str
        Default value for the shape of the chamber cross section 
        (elli/circ/rect).
    rho : float
        Default value for the resistivity of the chamber material [ohm.m].
    x_left : float, optional
        Default value for the left horizontal aperture in [m].
    y_bottom : float, optional
        Default value for the bottom vertical aperture in [m]. 
    n_points : int or float, optional
        Number of points to use in class arrays
        
    position : array of shape (n_points,)
        Longitudinal position in [m].
    x_right : array of shape (n_points,)
        Right horizontal aperture in [m].
    y_top : array of shape (n_points,)
        Top vertical aperture in [m].
    shape : array of shape (n_points - 1,)
        Shape of the chamber cross section (elli/circ/rect).
    rho : array of shape (n_points - 1,)
        Resistivity of the chamber material in [ohm.m].
    x_left : array of shape (n_points,)
        Left horizontal aperture in [m].
    y_bottom : array of shape (n_points,)
        Bottom vertical aperture in [m].
    length : array of shape (n_points - 1,)
        Length of each segments between two longitudinal positions in [m].
    center : array of shape (n_points - 1,)
        Center of each segments between two longitudinal positions in [m].
        
    Methods
    -------
    change_values(start_position, end_position, x_right, y_top, shape, rho)
        Change the physical parameters between start_position and end_position.
    taper(start_position, end_position, x_right_start, x_right_end, 
          y_top_start, y_top_end, shape, rho)
        Change the physical parameters to have a tapered transition between 
        start_position and end_position.
    plot_aperture()
        Plot horizontal and vertical apertures.
    resistive_wall_effective_radius(optics)
        Return the effective radius of the chamber for resistive wall 
        calculations.
    """

                 ring,
                 x_right,
                 y_top,
                 shape,
                 rho,
                 x_left=None,
                 y_bottom=None,
                 n_points=1e4):

        self.n_points = int(n_points)
        self.position = np.linspace(0, ring.L, self.n_points)
        self.x_right = np.ones_like(self.position) * x_right
        self.y_top = np.ones_like(self.position) * y_top

        if x_left is None:
            self.x_left = np.ones_like(self.position) * -1 * x_right
        else:
            self.x_left = np.ones_like(self.position) * x_left

        if y_bottom is None:
            self.y_bottom = np.ones_like(self.position) * -1 * y_top
        else:
            self.y_bottom = np.ones_like(self.position) * y_bottom

        self.length = self.position[1:] - self.position[:-1]
        self.center = (self.position[1:] + self.position[:-1]) / 2
        self.rho = np.ones_like(self.center) * rho

        self.shape = np.repeat(np.array([shape]), self.n_points - 1)

    def resistive_wall_effective_radius(self,
                                        optics,
                                        x_right=True,
                                        y_top=True):
        """
        calculations as defined in Eq. 27 of [1].
        
        Section with rho = 0 ohm.m are not considered in this calculation.

        optics : Optics object
        x_right : bool, optional
            If True, x_right is used, if Fasle, x_left is used.
        y_top : bool, optional
            If True, y_top is used, if Fasle, y_bottom is used.

        Returns
        -------
        L : float
            Total length considered in [m].
        rho_star : float
            Effective resistivity of the chamber material in [ohm.m].
        a1_L : float
            Effective longitudinal radius of the chamber of the first power in
            [m].
        a2_L : float
            Effective longitudinal radius of the chamber of the second power 
            in [m].
        a3_H : float
            Effective horizontal radius of the chamber of the third power in
            [m].
        a4_H : float
            Effective horizontal radius of the chamber of the fourth power in
            [m].
        a3_V : float
            Effective vertical radius of the chamber of the third power in [m].
        a4_V : float
            Effective vertical radius of the chamber of the fourth power in 
            [m].

        References  
        [1] Skripka, Galina, et al. "Simultaneous computation of intrabunch 
        and interbunch collective beam motions in storage rings." Nuclear 
        Instruments and Methods in Physics Research Section A: Accelerators, 
        Spectrometers, Detectors and Associated Equipment 806 (2016): 221-230.         
        """

        idx = self.rho != 0

        if x_right is True:
            a0 = (self.x_right[1:] + self.x_right[:-1]) / 2
        else:
            a0 = np.abs((self.x_left[1:] + self.x_left[:-1]) / 2)

        if y_top is True:
            b0 = (self.y_top[1:] + self.y_top[:-1]) / 2
        else:
            b0 = np.abs((self.y_bottom[1:] + self.y_bottom[:-1]) / 2)

        a0 = a0[idx]
        b0 = b0[idx]

        beta = optics.beta(self.center[idx])
        length = self.length[idx]

        L = length.sum()
        sigma = 1 / self.rho[idx]
        beta_H_star = 1 / L * (length * beta[0, :]).sum()
        beta_V_star = 1 / L * (length * beta[1, :]).sum()
        sigma_star = 1 / L * (length * sigma).sum()

        a1_H = (((length * beta[0, :] /
                  (np.sqrt(sigma) * (a0)**1)).sum())**(-1) * L * beta_H_star /
                np.sqrt(sigma_star))**(1 / 1)
        a2_H = (((length * beta[0, :] /
                  (np.sqrt(sigma) * (a0)**2)).sum())**(-1) * L * beta_H_star /
                np.sqrt(sigma_star))**(1 / 2)

        a1_V = (((length * beta[1, :] /
                  (np.sqrt(sigma) * (b0)**1)).sum())**(-1) * L * beta_V_star /
        a2_V = (((length * beta[1, :] /
                  (np.sqrt(sigma) * (b0)**2)).sum())**(-1) * L * beta_V_star /

        a3_H = (((length * beta[0, :] /
                  (np.sqrt(sigma) * (a0)**3)).sum())**(-1) * L * beta_H_star /
                np.sqrt(sigma_star))**(1 / 3)
        a4_H = (((length * beta[0, :] /
                  (np.sqrt(sigma) * (a0)**4)).sum())**(-1) * L * beta_H_star /
                np.sqrt(sigma_star))**(1 / 4)

        a3_V = (((length * beta[1, :] /
                  (np.sqrt(sigma) * (b0)**3)).sum())**(-1) * L * beta_V_star /
        a4_V = (((length * beta[1, :] /
                  (np.sqrt(sigma) * (b0)**4)).sum())**(-1) * L * beta_V_star /

        a1_L = min((a1_H, a1_V))
        a2_L = min((a2_H, a2_V))

        return (L, 1 / sigma_star, beta_H_star, beta_V_star, a1_L, a2_L, a3_H,
                a4_H, a3_V, a4_V)

    def change_values(self,
                      start_position,
                      end_position,
                      x_right=None,
                      y_top=None,
                      shape=None,
                      rho=None,
                      x_left=None,
                      y_bottom=None,
                      sym=True):
        """

        start_position : float
        end_position : float
        x_right : float, optional
            Right horizontal aperture in [m].
        y_top : float, optional
            Top vertical aperture in [m].
        shape : str, optional
            Shape of the chamber cross section (elli/circ/rect).
        rho : float, optional
            Resistivity of the chamber material [ohm.m].
        x_left : float, optional
            Left horizontal aperture in [m].
        y_bottom : float, optional
            Bottom vertical aperture in [m].
        sym : bool, optional
            If True, right/left and top/bottum symmetry is applied.
        """
        ind = (self.position > start_position) & (self.position < end_position)
        if x_right is not None:
            self.x_right[ind] = x_right
        if y_top is not None:
            self.y_top[ind] = y_top
        if x_left is not None:
            self.x_left[ind] = x_left
        elif (x_right is not None) and (sym is True):
            self.x_left[ind] = -1 * x_right
        if y_bottom is not None:
            self.y_bottom[ind] = y_bottom
        elif (y_top is not None) and (sym is True):
            self.y_bottom[ind] = -1 * y_top

        ind2 = ((self.position[:-1] > start_position) &
                (self.position[1:] < end_position))
        if rho is not None:
            self.rho[ind2] = rho
        if shape is not None:
            self.shape[ind2] = shape

    def taper(self,
              start_position,
              end_position,
              x_right_start=None,
              x_right_end=None,
              y_top_start=None,
              y_top_end=None,
              shape=None,
              rho=None,
              x_left_start=None,
              x_left_end=None,
              y_bottom_start=None,
              y_bottom_end=None,
              sym=True):
        """

        x_right_start : float
            Right horizontal aperture at the start of the taper in [m].
        x_right_end : float
            Right horizontal aperture at the end of the taper in [m].
        y_top_start : float
            Top vertical aperture at the start of the taper in [m].
        y_top_end : float
            Top vertical aperture at the end of the taper in [m].
        shape : str
        rho : float
        x_left_start : float, optional
            Left horizontal aperture at the start of the taper in [m].
        x_left_end : float, optional
            Left horizontal aperture at the end of the taper in [m].
        y_bottom_start : float, optional
            Bottom vertical aperture at the start of the taper in [m].
        y_bottom_end : float, optional
            Bottom vertical aperture at the end of the taper in [m].
        """
        if (x_right_start is not None) and (x_right_end is not None):
            self.x_right[ind] = np.linspace(x_right_start, x_right_end,
                                            sum(ind))
            if sym is True:
                self.x_left[ind] = -1 * np.linspace(x_right_start, x_right_end,
                                                    sum(ind))

        if (y_top_start is not None) and (y_top_end is not None):
            self.y_top[ind] = np.linspace(y_top_start, y_top_end, sum(ind))
                self.y_bottom[ind] = -1 * np.linspace(y_top_start, y_top_end,
                                                      sum(ind))

        if (x_left_start is not None) and (x_left_end is not None):
            self.x_left[ind] = np.linspace(x_left_start, x_left_end, sum(ind))
        if (y_bottom_start is not None) and (y_bottom_end is not None):
            self.y_bottom[ind] = np.linspace(y_bottom_start, y_bottom_end,
                                             sum(ind))

    def plot_aperture(self):
        """Plot horizontal and vertical apertures."""
        fig, axs = plt.subplots(2)
        axs[0].plot(self.position, self.x_right * 1e3)
        axs[0].plot(self.position, self.x_left * 1e3)
        axs[0].set(xlabel="Longitudinal position [m]",
                   ylabel="Horizontal aperture [mm]")
        axs[0].legend(["Right", "Left"])

        axs[1].plot(self.position, self.y_top * 1e3)
        axs[1].plot(self.position, self.y_bottom * 1e3)
        axs[1].set(xlabel="Longitudinal position [m]",
                   ylabel="Vertical aperture [mm]")
        axs[1].legend(["Top", "Bottom"])

        return (fig, axs)

    def get_aperture(self, s):
        self.xp = interp1d(self.position, self.x_right, kind='linear')
        self.xm = interp1d(self.position, self.x_left, kind='linear')
        self.yp = interp1d(self.position, self.y_top, kind='linear')
        self.ym = interp1d(self.position, self.y_bottom, kind='linear')
        aperture = np.array([self.xp(s), self.xm(s), self.yp(s), self.ym(s)])
        return aperture

    def plot_resistivity(self):
        """Plot resistivity along the ring."""
        fig, ax = plt.subplots(1)
        ax.plot(self.position[1:], self.rho)
        ax.set(xlabel="Longitudinal position [m]",
               ylabel="Resistivity [ohm.m]")

        return (fig, ax)

"""
Module where function used to import impedance and wakes from other codes are
defined.
"""

import os
import re
from tempfile import NamedTemporaryFile

def read_CST(file, component_type='long', divide_by=None, imp=True):
    """
    Read CST text file format into an Impedance or WakeFunction object.

    file : str
        Path to the text file to read.
    component_type : str, optional
        Type of the Impedance or WakeFunction object to load.
        Default is 'long'.
    divide_by : float, optional
        Divide the impedance by a value. Mainly used to normalize transverse
        impedance by displacement.
    imp : bool, optional.
        If True a Impedance object is loaded, if False a WakeFunction object is
        loaded.

    Returns
    -------
    result : Impedance or WakeFunction object.
        Data from file.
    """
    if imp:
        df = pd.read_csv(file,
                         comment="#",
                         header=None,
                         sep="\t",
                         names=["Frequency", "Real", "Imaginary"])
        df["Frequency"] = df["Frequency"] * 1e9
        if divide_by is not None:
            df["Real"] = df["Real"] / divide_by
            df["Imaginary"] = df["Imaginary"] / divide_by
        if component_type == "long":
            df["Real"] = np.abs(df["Real"])
        df.set_index("Frequency", inplace=True)
        result = Impedance(variable=df.index,
                           function=df["Real"] + 1j * df["Imaginary"],
                           component_type=component_type)
    else:
                         sep="\t",
                         names=["Distance", "Wake"])
        df["Time"] = df["Distance"] * 1e-3 / c
        df["Wake"] = df["Wake"] * 1e12
            df["Wake"] = df["Wake"] / divide_by
        df.set_index("Time", inplace=True)
        result = WakeFunction(variable=df.index,
                              function=df["Wake"],
                              component_type=component_type)
    return result

def read_IW2D(file, file_type='Zlong', output=False):
    """
    Read IW2D file format into an Impedance object or a WakeField object.

    file : str
        Path to the file to read.
    file_type : str, optional
        Type of the Impedance or WakeField object.
    output : bool, optional
        If True, print out the interpolated values.

    Returns
    -------
    result : Impedance or WakeField object
    """
    if file_type[0] == "Z":
                         delim_whitespace=True,
                         names=["Frequency", "Real", "Imaginary"],
                         skiprows=1)
        df = df[df["Real"].notna()]
        df = df[df["Imaginary"].notna()]
                           component_type=file_type[1:])
    elif file_type[0] == "W":
                         names=["Distance", "Wake"],
        df["Time"] = df["Distance"] / c
        if np.any(df.isna()):
            index = df.isna().values
            df = df.interpolate()
            if output:
                print("Nan values have been interpolated to:")
                print(df[index])
                              component_type=file_type[1:])
    else:
        raise ValueError("file_type should begin by Z or W.")

def read_IW2D_folder(folder, suffix, select="WZ", output=False):
    """
    Read IW2D results into a WakeField object.

    folder : str
        Path to the folder to read.
    suffix : str
        End of the name of each files. For example, in "Zlong_test.dat" the
        suffix should be "_test.dat".
    select : str, optional
        Select which object to load. "W" for WakeFunction, "Z" for Impedance
        and "WZ" or "ZW" for both.

    Returns
    -------
    result : WakeField object
        WakeField object with Impedance and WakeFunction objects from the
        different files.
    """
    if (select == "WZ") or (select == "ZW"):
        types = {"W": WakeFunction, "Z": Impedance}
    elif (select == "W"):
        types = {"W": WakeFunction}
    elif (select == "Z"):
        types = {"Z": Impedance}
    else:
        raise ValueError("select should be W, Z or WZ.")

    components = ["long", "xdip", "ydip", "xquad", "yquad"]

    data_folder = Path(folder)

    list_for_wakefield = []
    for key, item in types.items():
        for component in components:
            name = data_folder / (key+component+suffix)
            res = read_IW2D(file=name,
                            file_type=key + component,
                            output=output)
            list_for_wakefield.append(res)

    wake = WakeField(list_for_wakefield)

    return wake

def read_ABCI(file, azimuthal=False, output=False):
    """
    Read ABCI output files [1].

    file : str
        Path to ABCI .pot file.
    azimuthal : bool, optional
        If True, the transverse wake potential and impedance is loaded from the
        "AZIMUTHAL" data.
        If False, it is loaded from the "TRANSVERSE" data. In that case, a -1
        factor is applied on the wake to agree with mbtrack2 sign convention.
        If True, print out the loaded components header.

    Returns
    -------
    wake : WakeField
        Object where the ABCI computed impedance and wake are stored.

    [1] : ABCI - https://abci.kek.jp/abci.htm

    """

    if azimuthal:
        source = "AZIMUTHAL"
    else:
        source = "TRANSVERSE"

    def _read_temp(file, file_type, file2=None):
        if file_type[0] == "Z":
            df = pd.read_csv(file,
                             delim_whitespace=True,
                             names=["Frequency", "Real"])
            df["Real"] = df["Real"] * 1e3
            df["Frequency"] = df["Frequency"] * 1e9
            df2 = pd.read_csv(file2,
                              delim_whitespace=True,
                              names=["Frequency", "Imaginary"])
            df2["Imaginary"] = df2["Imaginary"] * 1e3
            df2["Frequency"] = df2["Frequency"] * 1e9
            df.set_index("Frequency", inplace=True)
            df2.set_index("Frequency", inplace=True)
            result = Impedance(variable=df.index,
                               function=df["Real"] + 1j * df2["Imaginary"],
                               component_type=file_type[1:])
        elif file_type[0] == "W":
                             names=["Time", "Wake"])
            df["Time"] = df["Time"] / c
            df["Wake"] = df["Wake"] * 1e12
            if (not azimuthal) and (file_type[-3:] == "dip"):
                df["Wake"] = df["Wake"] * -1
            df.set_index("Time", inplace=True)
            result = WakeFunction(variable=df.index,
                                  function=df["Wake"],
                                  component_type=file_type[1:])
        return result

    abci_dict = {
        '  TITLE: LONGITUDINAL WAKE POTENTIAL             \n':
        'Wlong',
        '  TITLE: REAL PART OF LONGITUDINAL IMPEDANCE                      \n':
        'Zlong_re',
        '  TITLE: IMAGINARY PART OF LONGITUDINAL IMPEDANCE                 \n':
        'Zlong_im',
        f'  TITLE: {source} WAKE POTENTIAL               \n':
        'Wxdip',
        f'  TITLE: REAL PART OF {source} IMPEDANCE                        \n':
        'Zxdip_re',
        f'  TITLE: IMAGINARY PART OF {source} IMPEDANCE                   \n':
        'Zxdip_im'
    }

    wake_list = []
    start = True
    with open(file) as f:
        while True:
            if start is True:
                # read the header
                header = [next(f) for _ in range(5)]
                start = False
            elif line == '':
                # check if file is over
                break
            else:
                header = [next(f) for _ in range(4)]
                header.insert(0, line)
            if output:
                print(header)
            # read the body until next TITLE field or end of file
            body = []
            while True:
                line = f.readline()
                try:
                    if line[:8] == "  TITLE:":
                        break
                    if line == '':
                        break
                except:
                    pass
                body.append(line)
            # write body in temp file and then process it into wake/imp
            impedance_type = int(re.split('\s+', header[-2])[-3])
            try:
                if abci_dict[header[0]][0] == "W":
                    tmp = NamedTemporaryFile(delete=False, mode="w+")
                    tmp.writelines(body)
                    tmp.flush()
                    tmp.close()
                    if abci_dict[
                            header[0]][1:] == "long" and impedance_type == 0:
                        comp = _read_temp(tmp.name, abci_dict[header[0]])
                        wake_list.append(comp)
                    elif abci_dict[
                            header[0]][1:] == "long" and impedance_type == 1:
                        pass
                        # Wlongxdip & Wlongydip not implemented yet
                        # comp = _read_temp(tmp.name, abci_dict[header[0]]+'dip')
                        # wake_list.append(comp)
                    else:
                        comp_x = _read_temp(tmp.name, "Wxdip")
                        comp_y = _read_temp(tmp.name, "Wydip")
                        wake_list.append(comp_x)
                        wake_list.append(comp_y)
                    os.unlink(tmp.name)
                elif (abci_dict[header[0]][0]
                      == "Z") and (abci_dict[header[0]][-2:] == "re"):
                    tmp1 = NamedTemporaryFile(delete=False, mode="w+")
                    tmp1.writelines(body)
                    tmp1.flush()
                    tmp1.close()
                      == "Z") and (abci_dict[header[0]][-2:] == "im"):
                    tmp2 = NamedTemporaryFile(delete=False, mode="w+")
                    tmp2.writelines(body)
                    tmp2.flush()
                    tmp2.close()
                            header[0]][1:-3] == "long" and impedance_type == 0:
                        comp = _read_temp(tmp1.name, "Zlong", tmp2.name)
                            header[0]][1:-3] == "long" and impedance_type == 1:
                        pass
                        # Zlongxdip & Zlongydip not implemented yet
                        # comp = _read_temp(tmp1.name, "Zlongdip", tmp2.name)
                    else:
                        comp_x = _read_temp(tmp1.name, "Zxdip", tmp2.name)
                        comp_y = _read_temp(tmp1.name, "Zydip", tmp2.name)
                    os.unlink(tmp1.name)
                    os.unlink(tmp2.name)
            except KeyError:
                pass

    wake = WakeField(wake_list)

def read_ECHO2D(file, component_type='long'):
    """
    Read ECHO2D text file format (after matlab post-processing) into a
    WakeFunction object.

    file : str
        Type of the WakeFunction object to load.

    Returns
    -------
    result : WakeFunction object.
    """

    df = pd.read_csv(file,
                     delim_whitespace=True,
                     header=None,
                     names=["Distance", "Wake"])
    df["Time"] = df["Distance"] / 100 / c
    df["Wake"] = df["Wake"] * 1e12
    if component_type != 'long':
        df["Wake"] = df["Wake"] * -1
    df.set_index("Time", inplace=True)
    result = WakeFunction(variable=df.index,
                          function=df["Wake"],
                          component_type=component_type)

"""
Module where bunch and beam spectrums and profile are defined.
"""

from scipy.special import jv, spherical_jn

def spectral_density(frequency, sigma, m=1, k=0, mode="Hermite"):
    """
    Compute the spectral density of different modes for various values of the
    head-tail mode number, based on Table 1 p238 of [1].

    frequency : list or numpy array
        sample points of the spectral density in [Hz]
    m : int, optional
    k : int, optional
        radial mode number (such that |q|=m+2k, where |q| is the head-tail mode number)
        -"Hermite" modes for Gaussian bunches (typical for electrons)
        -"Chebyshev" for airbag bunches
        -"Legendre" for parabolic bunches (typical for protons)
        -"Sacherer" or "Sinusoidal" simplifying approximation of Legendre modes from [3]

    Returns
    -------
    numpy array

    [2] : Ng, K. Y. (2005). Physics of intensity dependent beam instabilities. WORLD SCIENTIFIC. https://doi.org/10.1142/5835
    [3] : Sacherer, F. J. (1972). Methods for computing bunched beam instabilities. CERN Tech. rep. CERN/SI-BR/72-5 https://cds.cern.ch/record/2291670?ln=en
    """

    if mode == "Hermite":
        return 1 / (np.math.factorial(m) *
                    2**m) * (2 * np.pi * frequency * sigma)**(
                        2 * m) * np.exp(-(2 * np.pi * frequency * sigma)**2)
    elif mode == "Chebyshev":
        tau_l = 4 * sigma
        return (jv(m, 2 * np.pi * frequency * tau_l))**2
    elif mode == "Legendre":
        return (spherical_jn(m, np.abs(2 * np.pi * frequency * tau_l)))**2
    elif mode == "Sacherer" or mode == "Sinusoidal":
        y = 4 * 2 * np.pi * frequency * sigma / np.pi
        return (2 * (m+1) / np.pi * 1 / np.abs(y**2 - (m + 1)**2) *
                np.sqrt(1 + (-1)**m * np.cos(np.pi * y)))**2
    else:

def gaussian_bunch_spectrum(frequency, sigma):
    """
    Compute a Gaussian bunch spectrum [1].

        sample points of the beam spectrum in [Hz].
        RMS bunch length in [s].

    Returns
    -------
    bunch_spectrum : array
        Bunch spectrum sampled at points given in frequency.

    regime and ion cloud mitigation in ThomX. p86, Eq. 4.19
    """
    return np.exp(-1 / 2 * (2 * np.pi * frequency)**2 * sigma**2)

def gaussian_bunch(time, sigma):
    """
    Compute a Gaussian bunch profile.

    time : array
        sample points of the bunch profile in [s].

    Returns
    -------
    bunch_profile : array
        Bunch profile in [s**-1] sampled at points given in time.
    """
    return np.exp(-1 / 2 * (time**2 / sigma**2)) / (sigma * np.sqrt(2 * np.pi))

def beam_spectrum(frequency,
                  M,
                  bunch_spacing,
                  bunch_spectrum=None):
    """
    Compute the beam spectrum assuming constant spacing between bunches [1].

    M : int
        Time between two bunches in [s].
    sigma : float, optional
        If bunch_spectrum is None then a Gaussian bunch with sigma RMS bunch 
        length in [s] is assumed.
    bunch_spectrum : array, optional

    Returns
    -------
    beam_spectrum : array

    [1] Rumolo, G - Beam Instabilities - CAS - CERN Accelerator School: 
        Advanced Accelerator Physics Course - 2014, Eq. 9
    """

    if bunch_spectrum is None:
        bunch_spectrum = gaussian_bunch_spectrum(frequency, sigma)

    beam_spectrum = (bunch_spectrum *
                     np.exp(1j * np.pi * frequency * bunch_spacing * (M-1)) *
                     np.sin(M * np.pi * frequency * bunch_spacing) /
                     np.sin(np.pi * frequency * bunch_spacing))

    return beam_spectrum