import matplotlib
matplotlib.use('AGG')
from matplotlib import pyplot as plt
import numpy as np

import glob
import os
from collections import namedtuple
import argparse

parser = argparse.ArgumentParser(description="Process .loss files generated by OPAL and extract RF phase", 
                                 formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('searchdir', nargs="+", type=str, help="Directory or directories to search for PROBE*.loss files")
args = parser.parse_args()

nominal_freq = 116.4e6  # in Hz
PI = np.pi
h = 25
N_sectors = 6

class ProbeData():
    def __init__(self, fn):
        """ 
        Read data from a .loss file generated by a PROBE element in OPAL
        
        x, y, z are given in mm
        px, py, pz are given in dimensionless (beta*gamma) units
        id, turn are dimensionless integers
        time is given in s
        """
        data = np.loadtxt(fn, usecols=range(1, 10)).reshape((-1, 9))
        self.x = data[:, 0]*1e-3
        self.y = data[:, 1]*1e-3
        self.z = data[:, 2]*1e-3
        self.px = data[:, 3]
        self.py = data[:, 4]
        self.pz = data[:, 5]
        self.id = data[:, 6]
        self.turn = data[:, 7]
        self.time = data[:, 8]*1e-9
        

def get_cavity_phases(fn):
    """ Given an OPAL input file, extract phi1, phi2, phi3, phi4 (rad) """
    with open (fn, 'r') as f:
        getvars = {'t_offset': None, 
                   'dphi': None, 
                   'phi_correction': None, 
                   'phi0': None, 
                   'phi1': None, 
                   'phi2': None, 
                   'phi3': None, 
                   'phi4': None}
        locals().update(getvars)
        for line in f:
            if not line.startswith("//") and any(line.startswith(v) for v in getvars.keys()):
                varname, varvalue = [l.strip().strip(';') for l in line.split('=')]
                varvalue = eval(varvalue)
                if varname == 'phi_correction':
                    varvalue *= 1e-6  # MHz vs Hz
                locals()[varname] = varvalue
    
    return [locals()['phi%d' % i] for i in range(1, 5)]

def extract_phase(fn, phase):
    """
    Given a path to a PROBE.loss file and an initial phase (rad), 
    calculate the RF phase for each crossing
    """
    probe = ProbeData(fn)
    Ncrossings = probe.time.size

    RFphases = ((probe.time*nominal_freq*2*np.pi + phase)*180./np.pi)%360 - 360 
    return RFphases

if __name__ == "__main__":
    for cdir in args.searchdir:
        search = os.path.join(cdir, 'PROBE*.loss')
        for fn in glob.glob(search):
           simdir = os.path.dirname(fn)
           simfile = glob.glob(os.path.join(simdir, '*.in'))[0]
           phases = get_cavity_phases(simfile)
           probenum = int(fn.split("PROBE")[1][0])
           cavphase = extract_phase(fn, phases[probenum])
           print('%s: %s deg' % (fn, cavphase))
