"""
set of functions for calculating the transit time factor and the input phase resulting in max energy gain in a field map
-- M. Eshraqi 2018 July 11
"""

import Batch_Convert_and_Cut_Binary_Fieldmap as bccb
import SP_Relativity as myrel
import numpy as np


def field_on_axis(Field, fieldmap_dim, Nz, Nrx, Ny):
    """
    Take a 2D or 3D fieldmap as a 1D numpy array (as from TraceWin) in MV/m,
    dimension of the fieldmap,
    number of steps (No points - 1) in z direction,
    number of steps in x/r direction, and
    number of steps in the y direction (for 3D fieldmap) and
    returns the field on-axis in the z direction
    example:
    field_on_axis(Field, 3, 100, 10, 10)
    """
    if fieldmap_dim == 2:
        Fieldmapmatrix = np.reshape(Field, (int(Nz + 1), int(Nrx + 1)))
        midpoint = int(Nrx / 2)
    elif fieldmap_dim == 3:
        Fieldmapmatrix = np.reshape(Field, (int(Nz + 1), int((Nrx + 1) * (Ny + 1))))
        midpoint = int(Nrx * Ny / 2)
    # print(midpoint)
    return Fieldmapmatrix[:, midpoint]


def Field_TTF(field_1d, step_dz, freq, E_in, mass, RetMaxPhase):
    """
    takes a 1D field on axis along the acceleration axis in MV/m,
    the mesh size (length of each step),
    frequency in MHz,
    input energy in MeV,
    mass in MeV and
    RetMaxPhase as boolean
    returns the TTF and at the given energy and
    if RetMaxPhase==True returns also input phase resulting in max energy gain
    """
    omega = freq * 1e6 * 2 * np.pi
    TTF_Phase = [0] * 2
    dE = E_in
    Emax = E_in
    phase_max = 0
    counter = 0
    phase_start = 0
    phase_end = 360
    d_phase = 90

    for iteration in np.arange(1, 10):
        for phase in np.arange(int(phase_start / d_phase), int(phase_end / d_phase)):
            counter += 1
            t = 0
            dE = E_in
            for fval in field_1d:
                dE += fval[0] * step_dz * np.cos(omega * t + d_phase * np.radians(phase))
                # print(dE)
                if dE < 0:
                    print("Please adjust the field level, particle decelrated to zero energy!")
                t += dz / (myrel.c() * myrel.beta(dE, mass))
            if dE > Emax:
                Emax = dE
                # print('Emax', Emax)
                phase_max = phase * d_phase
        phase_start = phase_max - d_phase
        phase_end = phase_max + d_phase
        d_phase = d_phase / 4
        if d_phase * 4 < 0.1:
            break

    # d_phase*=4
    # print(counter, d_phase, phase_max, (Emax-E_in)/Field_integral(field_1d, dz))
    TTF_Phase[0] = (Emax - E_in) / Field_integral(field_1d, dz)
    if RetMaxPhase:
        TTF_Phase[1] = phase_max
    return TTF_Phase


def Field_integral(field_1d, dz):
    """
    takes a 1D field on axis along the acceleration axis in MV/m,
    the mesh size (length of each step) and
    returns the rectangular integral of field
    """
    field_int = 0
    for fval in field_1d:
        # print(abs(fval[0]))
        field_int += abs(fval[0])
    return dz * field_int


def TTF_beta(field_1d, step_dz, freq, mass):
    """
    takes a 1D field on axis along the acceleration axis in MV/m,
    the mesh size (length of each step),
    freq in MHz, and mass of particle in MeV and
    returns a 2D array containig the TTF vs. beta. [Not very fast]
    example
    TTFarray = TTF_beta(one_d_field_on_axis, step_z, 704.42, 938)
    beta_list = TTFarray[0]
    TTF_list = TTFarray[1]
    """
    TTFb = np.zeros((2, 50))
    i = 0
    beta_in = 0.2
    beta_out = 1
    TTFb[0, :] = np.arange(beta_in, beta_out, (beta_out - beta_in) / 50)
    for bt in TTFb[0, :]:
        dummy = Field_TTF(field_1d, step_dz, freq, myrel.energy(bt, mass), mass, False)
        TTFb[1, i] = dummy[0]
        i += 1
    return TTFb


# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - #
#                                                                                     #
#                                    E X A M P L E S                                  #
#                                                                                     #
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - #

Header, Field = bccb.read_binary_fieldmap("Data/FM/HB_W_coupler.edz", 3)
fz = field_on_axis(Field, 3, Header[0], Header[2], Header[5])
z = np.arange(0, int(Header[0] + 1))
dz = Header[1] / Header[0]
z_points = dz * z

print(
    "Int_E:", Field_integral(fz, dz), "[TTF, Phase_Max_Energy]: ", Field_TTF(fz, dz, 704.42, 700, 938, True),
)
TTFb = TTF_beta(fz, dz, 704.42, 938)

# plot(z_points, fz);
# plot(TTFb[0,:], 50*TTFb[1,:]);