from epics import PV,caput_many
import numpy as np

class RF:
    def __init__(self, beamline = 'ARAMIS',debug=False, signal = None):
        self.beamline = beamline
        self.debug=debug
        self.signal = signal
        if self.beamline == 'ARAMIS':
            self.stations = ['S30CB%2.2d' % d for d in range(1,14)]
            self.RFamp = [PV('%s-RSYS:SET-ACC-VOLT' % station)  for station in self.stations] 
            self.RFphs = [PV('%s-RSYS:SET-BEAM-PHASE' % station)  for station in self.stations] 
            self.RFon =  [PV('%s-RSYS:REQUIRED-OP' % station)  for station in self.stations] 
            self.ErgRange = [PV('SATSY01-MBND200:P-SET'),PV('SARCL02-MBND100:P-SET')]
            self.EGain = PV('S30:SET-E-GAIN-OP')
            self.EPhase= PV('S30:SET-BEAM-PHASE-OP')
            self.Ns = 13
        else:
            self.EGain = ['S20:SET-E-GAIN-OP']
            self.Ns = 0       

    def energyProfile(self):
        if self.Ns == 0:
            return np.array([0,0])
        valAmp = [pv.value for pv in self.RFamp]
        valPhs = [pv.value for pv in self.RFphs]
        valOn =  [pv.value for pv in self.RFon]
        valErg = [pv.value for pv in self.ErgRange]

        prof = np.zeros((self.Ns+1))
        prof[0] = valErg[0]
        for i in range(self.Ns):
            dE = 0
            if valOn[i] == 1:
                dE = valAmp[i]*np.sin(np.pi*valPhs[i]/180.)
            prof[i+1] = prof[i]+dE
        prof = np.array(prof)
        Egain = prof - prof[0]
        scl = (valErg[1]-valErg[0])/np.max(Egain)
        prof = Egain*scl+valErg[0]
        return prof

    def changeEnergyGain(self,dE):
        if self.debug:
            print('DEBUG: Changing energy gain of linac by:',dE)
            return
        self.EGain.value = dE
        if dE > 0:
            self.EPhase.value = 90.
        else:
            self.EPhase.value = -90.

    def setRF(self,chns,vals):
        if self.debug:
            for i,ele in enumerate(chns):
                print('DEBUG: RF-Settings %s: %f' % (ele,vals[i]))
            return
        caput_many(chns,vals)