# -*- coding: utf-8 -*-
"""
This module defines general classes to describes wakefields, impedances and 
wake functions.
"""

import pickle
import warnings
from copy import deepcopy

import numpy as np
import pandas as pd
import scipy as sc
from scipy.constants import c
from scipy.integrate import trapz
from scipy.interpolate import interp1d


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

    Parameters
    ----------
    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
    """

    def __init__(self,
                 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.
        
        Parameters
        ----------
        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':
            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())

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

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

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

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

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

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

    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 = ComplexData()
            result_structure.data = self.data * factor
            return result_structure
        else:
            warnings.warn(('The multiplication factor is not a float '
                           'or an int.'), UserWarning)
            return self

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

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

    def __call__(self, values, interp_kind="cubic"):
        """
        Interpolation of the data by calling the class to have a function-like
        behaviour.
        
        Parameters
        ----------
        values : list or numpy array of complex, int or float
            values to be interpolated.
        interp_kind : str, optional
            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"],
                             kind=interp_kind)
        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,
                'plane': 'x'
            },
            'ydip': {
                'a': 0,
                'b': 1,
                'c': 0,
                'd': 0,
                'plane': 'y'
            },
            'xydip': {
                'a': 0,
                'b': 1,
                'c': 0,
                'd': 0,
                'plane': 'x'
            },
            'yxdip': {
                'a': 1,
                'b': 0,
                'c': 0,
                'd': 0,
                'plane': 'y'
            },
            'xquad': {
                'a': 0,
                'b': 0,
                'c': 1,
                'd': 0,
                'plane': 'x'
            },
            'yquad': {
                'a': 0,
                'b': 0,
                'c': 0,
                'd': 1,
                'plane': 'y'
            },
            'xyquad': {
                'a': 0,
                'b': 0,
                'c': 0,
                'd': 1,
                'plane': 'x'
            },
            'yxquad': {
                'a': 0,
                'b': 0,
                'c': 1,
                'd': 0,
                'plane': 'y'
            },
        }

        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.
    
    Parameters
    ----------
    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", ...
    
    Attributes
    ----------
    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.
    """

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

    def __repr__(self):
        """Return representation of data"""
        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.
        """

        if isinstance(structure_to_add, (int, float, complex)):
            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

    def __radd__(self, structure_to_add):
        return self.add(structure_to_add)

    def __add__(self, structure_to_add):
        return self.add(structure_to_add)

    def multiply(self, factor):
        """
        Multiply a WakeFunction object with a float or an int.
        If the multiplication is done with something else, throw a warning.
        """
        result = super().multiply(factor)
        wake_to_return = WakeFunction(result.data.index,
                                      result.data.real.values,
                                      self.component_type)
        return wake_to_return

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

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

    def to_impedance(self, freq_lim, nout=None, sigma=None, mu=None):
        """
        Return an Impedance object from the WakeFunction data.
        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.

        Parameters
        ----------
        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.
        sigma : float
            RMS of the Gaussian distribution in [s].
            The default is None.
        mu : float
            Offset of the Gaussian distribution in [s].
            The default is None.

        """
        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):
        """
        Compute a wake function from wake potential data.
                
        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".

        Parameters
        ----------
        freq_lim : float
            The maximum frequency for calculating the impedance [Hz].
        sigma : float
            RMS of the Gaussian distribution in [s].
        mu : float
            Offset of the Gaussian distribution in [s].
        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.

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

    def plot(self):
        """
        Plot the wake function data.

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

        """
        ax = self.data["real"].plot()
        if self.component_type == "long":
            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].
        
        Parameters
        ----------
        sigma : float
            RMS bunch length in [s]
        
        Returns
        -------
        kloss: float
            Loss factor in [V/C] or kick factor in [V/C/m] depanding on the 
            impedance type.
            
        References
        ----------
        [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.
    
    Parameters
    ----------
    variable : array-like
        Frequency domain structure of the impedance in [Hz].
    function : array-like
        Impedance values in [Ohm].
    component_type : str, optinal
        Type of the impedance: "long", "xdip", "xquad", ...
    
    Attributes
    ----------
    data : DataFrame
    impedance_type : str
    
    Methods
    -------
    loss_factor(sigma)
        Compute the loss factor or the kick factor for a Gaussian bunch.
    to_wakefunction()
        Return a WakeFunction object from the impedance data.
    plot()
        Plot the impedance data.
    """

    def __init__(self,
                 variable=np.array([-1e15, 1e15]),
                 function=np.array([0, 0]),
                 component_type='long'):
        super().__init__(variable, function)
        self._component_type = component_type
        self.data.index.name = 'frequency [Hz]'
        self.initialize_coefficient()

    def __repr__(self):
        """Return representation of data"""
        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
        compared so that the addition is self-consistent.
        Beta functions can be precised as well.
        """

        if isinstance(structure_to_add, (int, float, complex)):
            result = super().add(structure_to_add,
                                 method=method,
                                 index_name="frequency [Hz]")
        elif isinstance(structure_to_add, Impedance):
            if (self.component_type == structure_to_add.component_type):
                weight = (beta_x**self.power_x) * (beta_y**self.power_y)
                result = super().add(structure_to_add * weight,
                                     method=method,
                                     index_name="frequency [Hz]")
            else:
                warnings.warn(('The two Impedance objects do not have the '
                               'same coordinates or plane or type. '
                               'Returning initial Impedance object.'),
                              UserWarning)
                result = self

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

    def __radd__(self, structure_to_add):
        return self.add(structure_to_add)

    def __add__(self, structure_to_add):
        return self.add(structure_to_add)

    def multiply(self, factor):
        """
        Multiply a Impedance object with a float or an int.
        If the multiplication is done with something else, throw a warning.
        """
        result = super().multiply(factor)
        impedance_to_return = Impedance(
            result.data.index,
            result.data.real.values + 1j * result.data.imag.values,
            self.component_type)
        return impedance_to_return

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

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

    def loss_factor(self, sigma):
        """
        Compute the loss factor or the kick factor for a Gaussian bunch. 
        Formulas from Eq. (2) p239 and Eq.(7) p240 of [1].
        
        Parameters
        ----------
        sigma : float
            RMS bunch length in [s]
        
        Returns
        -------
        kloss: float
            Loss factor in [V/C] or kick factor in [V/C/m] depanding on the 
            impedance type.
            
        References
        ----------
        [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"):
            kloss = trapz(x=frequency,
                          y=2 * self.data["imag"][positive_index] * sd)
        elif (self.component_type == "xquad"
              or self.component_type == "yquad"):
            kloss = trapz(x=frequency,
                          y=2 * self.data["imag"][positive_index] * sd)
        else:
            raise TypeError("Impedance type not recognized.")

        return kloss

    def to_wakefunction(self, nout=None, trim=False):
        """
        Return a WakeFunction object from the impedance data.
        The impedance data is assumed to be sampled equally.
    
        Parameters
        ----------
        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.
        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]

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

        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

    def plot(self):
        """
        Plot the impedance data.

        Returns
        -------
        ax : class:`matplotlib.axes.Axes`
            Impedance data.

        """
        ax = self.data.plot()
        if self.component_type == "long":
            unit = " [ohm]"
        else:
            unit = " [ohm/m]"
        label = r"$Z_{" + self.component_type + r"}$" + unit
        ax.set_ylabel(label)
        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.
    
    Parameters
    ----------
    structure_list : list, optional
        list of Impedance/WakeFunction objects to add to the WakeField.
    name : str, optional
    
    Attributes
    ----------
    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(file)
        Save WakeField element to file.
    load(file)
        Load WakeField element from file.
    """

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

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

    def __radd__(self, structure_to_add):
        return self._add(structure_to_add)

    def __add__(self, structure_to_add):
        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.

        Parameters
        ----------
        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)
        elif isinstance(structure_to_add, WakeFunction):
            attribute_name = "W" + 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)
        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.

        Parameters
        ----------
        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"
        ]
        return np.array([comp for comp in dir(self) if comp in valid])

    @property
    def components(self):
        """
        Return an array of the component names for the element.
        """
        valid = [
            "Wlong", "Wxdip", "Wydip", "Wxquad", "Wyquad", "Wxcst", "Wycst",
            "Zlong", "Zxdip", "Zydip", "Zxquad", "Zyquad", "Zxcst", "Zycst"
        ]
        return np.array([comp for comp in dir(self) if comp in valid])

    def drop(self, component):
        """
        Delete a component or a list of component from the WakeField.

        Parameters
        ----------
        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.")

    def save(self, file):
        """
        Save WakeField element to file.
        
        The same pandas version is needed on both saving and loading computer
        for the pickle to work.

        Parameters
        ----------
        file : str
            File where the WakeField element is saved.

        Returns
        -------
        None.

        """
        with open(file, "wb") as f:
            pickle.dump(self, f)

    @staticmethod
    def load(file):
        """
        Load WakeField element from file.
        
        The same pandas version is needed on both saving and loading computer
        for the pickle to work.

        Parameters
        ----------
        file : str
            File where the WakeField element is saved.

        Returns
        -------
        wakefield : WakeField
            Loaded wakefield.

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

        return wakefield

    @staticmethod
    def add_wakefields(wake1, beta1, wake2, beta2):
        """
        Add two WakeFields taking into account beta functions.

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

        Returns
        -------
        wake_sum : WakeField
            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)
            except AttributeError:
                setattr(wake_sum, component_name2, weight2 * comp2)

        return wake_sum

    @staticmethod
    def add_several_wakefields(wakefields, beta):
        """
        Add a list of WakeFields taking into account beta functions.
        
        Parameters
        ----------
        wakefields : list of WakeField
            WakeFields to add.
        beta : array of shape (len(wakefields), 2)
            Beta functions at the WakeField postions.

        Returns
        -------
        wake_sum : WakeField
            WakeField with summed components.

        """
        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.")
