# -*- coding: utf-8 -*-
"""
Running script to test getplot method in Bunch class, particles module

@author: Watanyu Foosang
@date: 27/03/2020
"""

import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
from tracking.particles import Bunch, Beam
from machines.soleil import soleil
from tracking.rf import *
from scipy.constants import e

ring = soleil() # Define the properties of the storage ring
mybunch = Bunch(ring,mp_number=1000,current=1e-3) # Define bunch parameters

mybunch.init_gaussian() # Initialize 6D bunch parameters


v_rf = 1.1e6 # RF peak voltage [v]
sync_phase = np.arccos(ring.U0/v_rf) # Determine the synchronous phase
    
rf = RFCavity(ring, m=1, Vc=v_rf, theta=sync_phase)

rev = []
mean_tau = []
mean_delta = []    
def track_long(bunch,ring,rf,turn):
    """
    Track longitudinal beam parameters for one turn
    """
    
    rf.track(bunch) # update "delta" for one turn
    bunch.particles["tau"] = bunch.particles["tau"] - bunch.particles["delta"]\
        *ring.T0*ring.ac
    turn += 1
    rev.append(turn)
    mean_tau.append(bunch.particles["tau"].mean())
    mean_delta.append(bunch.particles["delta"].mean())
    return turn

def get_meanplot(ini,fin,step,val):
        """
        Plot the evolution of a mean value over turns.
        (For now, this function only works subject to plotting module.)
        
        Parameters
        ----------
        ini: int, starting revolution
        fin: int, final revolution
        step: int, step size on the revolution axis
        val: str, "tau" for maen values of tau or "del" for mean values of delta

        """
        
        x = rev[ini-1:fin:step]
        
        if val == "tau":
            val = mean_tau[ini-1:fin:step]
        else: val = mean_delta[ini-1:fin:step]
        
        plt.plot(x,val)