# -*- coding: utf-8 -*-
"""
Tools and utilities functions for the Wakefield and impedances classes.

@author: Alexis Gamelin
@date: 14/01/2020
"""

import pandas as pd
import numpy as np
from scipy.integrate import trapz
from CollectiveEffect.wakefield import Wakefield, Impedance, WakeFunction
#from ..Tracking.synchrotron import Synchrotron

def read_CST(file, impedance_type='long'):
    """Read CST file format into an Impedance object."""
    df = pd.read_csv(file, comment="#", header = None, sep = "\t", 
                    names = ["Frequency","Real","Imaginary"])
    df["Frequency"] = df["Frequency"]*1e9 
    df.set_index("Frequency", inplace = True)
    result = Impedance(variable = df.index,
                       function = df["Real"] + 1j*df["Imaginary"],
                       impedance_type=impedance_type)
    return result

def read_IW2D(file, impedance_type='long'):
    """Read IW2D file format into an Impedance object."""
    df = pd.read_csv(file, sep = ' ', header = None, 
                     names = ["Frequency","Real","Imaginary"], skiprows=1)
    df.set_index("Frequency", inplace = True)
    result = Impedance(variable = df.index,
                       function = df["Real"] + 1j*df["Imaginary"],
                       impedance_type=impedance_type)
    return result

def loss_factor(wake, sigma):
    """
    Compute the loss factor or the kick factor of a Wakefield, Impedance or
    WakeFunction for a Gaussian bunch. Formulas from Eq. (2) p239 and 
    Eq.(7) p240 of [1].
    
    Parameters
    ----------
    wake: Wakefield, Impedance or WakeFunction object.
    sigma: RMS bunch length in [s]
    
    Returns
    -------
    kloss: loss factor in [V/C] or kick factor in [V/C/m] depanding on the 
    impedance type.
        
    [1] Handbook of accelerator physics and engineering, 3rd printing.
    """
    
    if isinstance(wake, Impedance):
        
        positive_index = wake.data.index > 0
        omega = 2*np.pi*wake.data.index[positive_index]
        gaussian_spectrum = np.exp(-1*omega**2*sigma**2)
        
        if(wake.impedance_type == "long"):
            kloss = trapz(x = omega, 
                          y = 1/np.pi*wake.data["real"][positive_index]*gaussian_spectrum)
        if(wake.impedance_type == "xdip" or wake.impedance_type == "ydip"):
            kloss = trapz(x = omega, 
                          y = 1/np.pi*wake.data["imag"][positive_index]*gaussian_spectrum)
    if isinstance(wake, WakeFunction):
        # To be implemented.
        pass
    
    if isinstance(wake, Wakefield):
        # To be implemented. Create field in wakefield with values ?
        pass    
    
    return kloss

def spectral_density(frequency, sigma, m = 1, 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].
    
    Parameters
    ----------
    frequency : list or numpy array
        sample points of the spectral density in [Hz]
    sigma : float
        RMS bunch length in [s]
    m : int, optional
        head-tail (or azimutal/synchrotron) mode number
    mode: str, optional
        type of the mode taken into account for the computation:
        -"Hermite" modes for Gaussian bunches

    Returns
    -------
    numpy array
        
    [1] Handbook of accelerator physics and engineering, 3rd printing.
    """
    
    if mode == "Hermite":
        return (2*np.pi*frequency*sigma)**(2*m)*np.exp(
                -1*(2*np.pi*frequency*sigma)**2)
    else:
        raise NotImplementedError("Not implemanted yet.")
    
def effective_impedance(ring, imp, m, mu, sigma, xi=None, mode="Hermite"):
    """
    Compute the effective (longitudinal or transverse) impedance. 
    Formulas from Eq. (1) and (2) p238 of [1].
    
    Parameters
    ----------
    ring : Synchrotron object
    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
    sigma : float
        RMS bunch length in [s]
    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:
        raise ValueError("A double sided impedance spectrum, with positive and"
                         " negative frequencies, is needed to compute the "
                         "effective impedance.")
        
    if mode == "Hermite":
        def h(f):
            return spectral_density(frequency=f, sigma=sigma, m=m,
                                    mode="Hermite")
    else:
        raise NotImplementedError("Not implemanted yet.")
    
    M = 416
    
    pmax = fmax/(ring.f0 * M) - 50
    pmin = fmin/(ring.f0 * M) + 50
    
    p = np.arange(pmin,pmax+1)
    
    if imp.impedance_type == "long":
        fp = ring.f0*(p*M + mu + m*ring.tune[2])
        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.impedance_type == "xdip" or imp.impedance_type == "ydip":
        if imp.impedance_type == "xdip":
            tuneXY = ring.tune[0]
            xi = ring.chro[0]
        elif imp.impedance_type == "ydip":
            tuneXY = ring.tune[1]
            xi = ring.chro[1]
        fp = ring.f0*(p*M + mu + tuneXY + m*ring.tuneS)
        f_xi = xi/ring.eta*ring.f0
        num = np.sum( imp(fp) * h(fp - f_xi) )
        den = np.sum( h(fp - f_xi) )
        Zeff = num/den
    else:
        raise TypeError("Effective impedance is only defined for long, xdip"
                        " and ydip impedance type.")
        
    return Zeff