import sys
import numpy as np


from onlinemodel.core import Facility
from onlinemodel.code import CMadX

class Model:
    def __init__(self, phase=0, parent=None):
        print('Initializing online model ...')
        self.phase = phase   # current planned future
        self.parent=parent
        self.om = Facility(phase)
        self.madx = CMadX()
        self.twiss = None
        # hook up events
        self.eventHandling()

    def eventHandling(self):
        self.parent.UITrack.clicked.connect(self.track)

    def track(self):
        start = str(self.parent.UITrackStart.text())
        end = str(self.parent.UITrackEnd.text())
        twiss = {}
        twiss['betax0'] = float(str(self.parent.UIBetax.text()))
        twiss['betay0'] = float(str(self.parent.UIBetay.text()))
        twiss['alphax0'] = float(str(self.parent.UIAlphax.text()))
        twiss['alphay0'] = float(str(self.parent.UIAlphay.text()))
        twiss['energy0'] = 150
        self.doTrack(start,end,'SINLH02.MQUA410$START',twiss)


    def doBackTrack(self,start=None,end=None,twiss=None):
        twiss['alphax0'] = -twiss['alphax0']
        twiss['alphay0'] = -twiss['alphay0']
        self.madx.updateVariables(twiss)
        self.madx.updateLattice(self.om)
        twiss = self.madx.track(start, end, reverse=True)
        #self.parent.plot.newData(twiss)
        twiss0 = {}
        twiss0['betax0']= twiss.betx[-1]
        twiss0['betay0'] = twiss.bety[-1]
        twiss0['alphax0'] = -twiss.alfx[-1]
        twiss0['alphay0'] = -twiss.alfy[-1]
        self.madx.clear()
        return twiss0


    def doTrack(self, start,end,refloc, twiss):
        self.setBranch(start,end)
        # check if refloc is in the range of start and end
        if not refloc == start:
            twiss = self.doBackTrack(refloc,start,twiss)

        self.madx.updateVariables(twiss)
        self.madx.updateLattice(self.om)

        twiss = self.madx.track()
        self.parent.plot.newData(twiss)

    def setBranch(self,start,end):
        start0 = start[0:7]
        end0 = end[0:7]
        destination = 'ARAMIS'
        if 'SPO' in end:
            destination = 'PORTHOS'
        elif 'SAT' in end:
            destination = 'ATHOS'
        elif 'S10BD' in end or 'SIN' in end:
            destination = 'INJECTOR'
        self.om.setBranch(destination,start0,end0)