import sys
import numpy as np


from onlinemodel.core import Facility
from onlinemodel.code import MadX

class Model:
    def __init__(self, phase, debug=False):
        print('Initializing online model ...')
        self.phase = phase   # current planned future
        self.om = Facility(phase)
        self.madx = MadX()
        self.twiss = None

    def update(self, snap) -> None:
        self.om.updateModel(snap,True)

    def track(self, start,end,twiss):
        destination = 'ARAMIS'
        start0 = start[0:7]
        end0 = end[0:7]
        if 'SPO' in end:
            destination = 'PORTHOS'
        elif 'SAT' in end:
            destination = 'ATHOS'
        elif 'S10BD' in end:
            destination = 'INJECTOR'
        self.om.setBranch(destination,start0,end0)
        self.madx.setInitialCondition(twiss['betax'],twiss['betay'],twiss['alphax'],twiss['alphay'],150.)
        if len(start) <8:
            start = '#s'
        if len(end) < 8:
            end = '#e'
        self.madx.track(self.om,False,start,end)
        self.twiss = self.madx.getTwiss()