import os
import re
from onlinemodel.madx import CMadX


class MatchMaker:
    def __init__(self):
        self.matchlist={'Reference-SwissFEL':'Scripts/Reference-SwissFEL','SwissFEL+':'Scripts/SFPlus'}
        self.referencePoints={}

    def initScripts(self,target):
        self.scriptdir = self.matchlist[target]
        if os.path.exists(self.scriptdir+'/initTwiss.madx'):
            self.parseReferencePoints()
        self.scriptInjector =  os.path.exists(self.scriptdir+'/matchInjector.madx')
        self.scriptAramis = os.path.exists(self.scriptdir+'/matchAramis.madx')
        self.scriptAthos = os.path.exists(self.scriptdir + '/matchAthos.madx')
        self.scriptPorthos = os.path.exists(self.scriptdir + '/matchPorthos.madx')
        self.settings = os.path.exists(self.scriptdir + '/ReferenceSettings')


    def parseReferencePoints(self):
        file = self.scriptdir+'/initTwiss.madx'
        self.referencePoints.clear()
        with open(file,'r') as f:
            lines = f.readlines()
            for line in lines:
                if 'location:' in line:
                    tags = line.split('location:')
                    location = tags[1].strip()
                    fields_all = tags[0].split(',')
                    fields = [fld.split(';')[0].strip() for fld in fields_all if "=" in fld]
                    self.referencePoints[location] = {twiss.split('=')[0].strip().lower():float(twiss.split('=')[1].strip()) for twiss in fields}
        print('##### Reference Twiss values parsed')


    def match(self, om, Injector=True, Athos = True, Aramis = False, Porthos = False):

        if Athos:
            target = 'SATBD01'
        else:
            target = 'SARUN20'
        madx = CMadX()
        madx.E0 = 2000
        om.setBranch(target, 'SINLH01')
        madx.updateLattice(om, target)  # write lattice
        madx.commonHeader('SwissFEL', '#s/#e', None)  # sets header
        madx.madx.call(self.scriptdir + '/initTwiss.madx')
        if Injector:
            madx.madx.call(self.scriptdir+'/matchInjector.madx', chdir=True)
            self.updateOnlineModel(om,madx.madx, 's[i1].*k1|s20sy01.*k1')
        if Athos:
            madx.madx.call(self.scriptdir+'/matchAthos.madx', chdir=True)
            self.updateOnlineModel(om, madx.madx, 'sat.*k1|s20sy02.*m[kq]')
        if Aramis:
            if target=='SATBD01':
                target = 'SARUN20'
                om.setBranch(target, 'SINLH01')
                madx.updateLattice(om, target)  # write lattice
                madx.commonHeader('SwissFEL', '#s/#e', None)  # sets header
                madx.madx.call(self.scriptdir + '/initTwiss.madx')
            madx.madx.call(self.scriptdir + '/matchAramis.madx', chdir=True)
            self.updateOnlineModel(om, madx.madx, 's[3a][0r].*k1|s20[bm][ca].*k1')
        return madx.madx.table.twiss

    def updateOnlineModel(self,om,madx,filter):
        reg = re.compile(filter)
        for var in list(madx.globals):
            if reg.match(var):
                if '.k1' in var:
                    element = var.split('.k1')[0]
                    ele = om.getElement(element)
                    if not ele is None:
                        ele.k1 = madx.globals.get(var,0)
                        print('... Updating', ele.Name, 'to k1=',ele.k1)
                    else:
                        print('### Unidentified element', element)
                elif '.cory' in var:
                    element = var.split('.cory')[0]
                    ele = om.getElement(element)
                    if not ele is None:
                        ele.design_kick = madx.globals.get(var, 0)
                        print('... Updating', ele.Name, 'to design_kick=', ele.design_kick)
                    else:
                        print('### Unidentified element', element)