import json

def configMatching(label_in = 'Reference',var={}):
    path = "MatchingConfig/"
    label = label_in
    print(var)
    dependence = generateDependences()
    reference = generateReferences()
    order = generateOrder()
    parameter = generateMatchingSteps(reference,var)


    settings={'Label':label,'Order':order,'Dependence':dependence,'Reference':reference,'Parameter':parameter}
    fname=path+label+'.json'
    with open(fname, 'w', encoding='utf-8') as f:
        json.dump(settings, f, ensure_ascii=False, indent=4)
    return fname
#######
# some internal functions
def MatchReference(loc,twiss,limit=None):  # limit 1 -> smaller than, 2 -> larger than (default is euqal)
    match ={'Location': loc, 'Twiss': twiss}
    if not limit is None:
        match['Limit'] = limit
    return match

def MatchStep(ref,id,end, dest,sequence=None,saveTwiss=None):
    match={'Reference':ref,'MatchID':id,'End':end, 'Destination':dest,'Sequence':sequence,'Save':saveTwiss}
    return match

def MatchParameter(var0,tar0,user0=None):
    if user0 is None:
        return {'Variable':var0,'Target':tar0}
    return {'Variable':var0,'Target':tar0,'User':user0}

############
# generating functions

def generateReferences():
    # list of references with fixed twiss parameters
    reference = {}
    reference['Start'] = MatchReference('start', {'betax': 29.971, 'alphax': 0.003, 'betay': 26.00, 'alphay': -0.288})
    reference['Laser Heater'] = MatchReference('sinlh02.mqua410$start',
                                               {'betax': 11.2, 'alphax': 3.2, 'betay': 0.8, 'alphay': 0.75, })
    reference['Bunch Compressor 1'] = MatchReference('sindi02.mqua020$start',
                                                     {'betax': 11.1557, 'alphax': -1.17, 'betay': 50, 'alphay': 0.})
    reference['Bunch Compressor 2'] = MatchReference('s10ma01.mqua020$start',
                                                     {'betax': 6.14, 'alphax': -0.83, 'betay': 22.86, 'alphay': -1.18})
    reference['Switchyard'] = MatchReference('s20sy02$start',
                                             {'betax': 48.26, 'alphax': 7.322, 'betay': 14.293, 'alphay': -3.513})
    reference['Linac3'] = MatchReference('s30cb10.mqua430$start',
                                         {'betax': 8.83, 'alphax': 0.548, 'betay': 30.44, 'alphay': -1.866})
    reference['Aramis Stopper'] = MatchReference('sarma01.mqua010$start',
                                                 {'betax': 2.304, 'alphax': -1.4045, 'betay': 25.929, 'alphay': -6.363})
    reference['HERO'] = MatchReference('satdi01.mqua250$start',
                                       {'betax': 35.76, 'alphax': -1.166, 'betay': 24.087, 'alphay': 0.648})
    reference['Athos Deflector'] = MatchReference('satbd01.mqua010$start',
                                                  {'betax': 50, 'alphax': 0, 'betay': 50, 'alphay': 0})
    reference['EEHG'] = MatchReference('satma01.mqua610$start',
                                       {'betax': 10., 'alphax': 0.027, 'betay': 16.85, 'alphay': 0.027})
    return reference


def generateDependences():
    dependence = {}
    dependence['SINSB04-MQUA230']=['SINSB03-MQUA230']
    dependence['S10CB02-MQUA230']=['S10DI01-MQUA120'] + ['S10CB%2.2d-MQUA430' % i for i in range(3,9)]
    dependence['S10CB02-MQUA430']=['S10CB%2.2d-MQUA230' % i for i in range(3,9)]
    dependence['S20CB01-MQUA430']=['S20CB03-MQUA430']
    dependence['SATCB01-MQUA230']=['SATCL02-MQUA430']
    dependence['SATCB01-MQUA430']=['SATCL02-MQUA430']
    dependence['SATUN06-MQUA080']=['SATUN%2.2d-MQUA080' % (2*i+6) for i in range(1,4)] +['SATUN14-MQUA420']+['SATUN%d-MQUA080' % (2*i+16) for i in range(3)]
    dependence['SATUN07-MQUA080']=['SATUN%2.2d-MQUA080' % (2*i+7) for i in range(1,8)]
    dependence['S30CB01-MQUA430'] = ['S30CB%2.2d-MQUA430' % (i*2 + 3)  for i in range(7)]
    dependence['S30CB02-MQUA430'] = ['S30CB%2.2d-MQUA430' % (i*2 + 4) for i in range(6)]
    dependence['SARUN03-MQUA080'] = ['SARMA02-MQUA050','SARUN01-MQUA080'] + ['SARUN%2.2d-MQUA080' % (2*i + 5) for i in range(8)]
    dependence['SARUN04-MQUA080'] = ['SARMA02-MQUA120','SARUN02-MQUA080'] + ['SARUN%2.2d-MQUA080' % (2*i + 6) for i in range(8)]
    return dependence

def generateOrder():
    order = []
    # order of elements:
    # - Match starting point with Twiss, None has no initial condition and starts with #s
    # - ID to find additional parameteters
    # - Last section in tracking/matching. If none it tracks to the end (#e)
    # - destination: selection of branch by giving a unique section name (e.g. SARBD01 for Aramis)
    # - sequence: sequence to use. If it is a list of sequence then a temporary sequence is used. Default to swissfel
    # - flag to last twiss parameters, needed for other routines to refer to.
    order.append(MatchStep(None, 'Preset', None, 'SATBD01'))
    order.append(MatchStep('Start', 'Laser Heater', 'SINLH02', 'SATBD01'))
    order.append(MatchStep(None, 'SINSB04', None, 'SATBD01', 'SINSB04', True))
    order.append(MatchStep('Laser Heater', 'Match to SINSB04', 'SINSB04', 'SATBD01'))
    order.append(MatchStep('Laser Heater', 'Bunch Compressor 1', 'SINDI02', 'SATBD01'))
    order.append(MatchStep(None, 'Linac 1', None, 'SATBD01', 'S10CB02', True))
    order.append(MatchStep('Bunch Compressor 1', 'Match to Linac 1', 'S10CB02', 'SATBD01'))
    order.append(MatchStep('Bunch Compressor 1', 'End of Linac 1', 'S10BC01', 'SATBD01'))
    order.append(MatchStep('Bunch Compressor 1', 'Bunch Compressor 2', 'S10MA01', 'SATBD01'))
    order.append(MatchStep(None, 'Linac 2', None, 'SATBD01', ['S20CB01', 'S20CB02'], True))
    order.append(MatchStep('Bunch Compressor 2', 'Match to Linac 2', 'S20CB04', 'SATBD01'))
    order.append(MatchStep('Bunch Compressor 2', 'Match to Switchyard', 'S20SY02', 'SATBD01'))
    order.append(MatchStep('Switchyard', 'Resonant Kicker', 'S20SY02', 'SATBD01'))
    order.append(MatchStep('Switchyard', 'Switchyard', 'SATCL01', 'SATBD01'))
    order.append(MatchStep('Switchyard', 'Athos Diagnostics', 'SATDI01', 'SATBD01'))
    order.append(MatchStep(None, 'Athos Linac', None, 'SATBD01', 'SATCB01', True))
    order.append(MatchStep('HERO', 'Athos Stopper', 'SATCL02', 'SATBD01'))
    order.append(MatchStep('HERO', 'Athos EEHG', 'SATMA01', 'SATBD01'))
    order.append(MatchStep(None, 'Athos Undulator', None, 'SATBD01', ['SATUN06', 'SATUN07'], True))
    order.append(MatchStep('EEHG', 'Match to Athos Undulator', 'SATUN09', 'SATBD01'))
    order.append(MatchStep('EEHG', 'Match to Athos Deflector', 'SATBD01', 'SATBD01'))
    order.append(MatchStep(None, 'Linac 3', None, 'SARBD01',['S30CB01','S30CB02'], True))
    order.append(MatchStep('Switchyard', 'Match to Linac 3', 'S30CB03', 'SARBD01'))
    order.append(MatchStep('Switchyard', 'End of Linac 3', 'S30CB10', 'SARBD01'))
    order.append(MatchStep(None, 'Aramis Undulator', None, 'SARBD01', ['SARUN03', 'SARUN04'], True))
    order.append(MatchStep('Linac3', 'Aramis - ECOL', None, 'SARBD01'))
    order.append(MatchStep('Linac3', 'Match to Aramis Undulator', 'SARUN20', 'SARBD01'))
    order.append(MatchStep('Linac3', 'Aramis Dump', 'SARBD02', 'SARBd02'))
    return order


def generateMatchingSteps(reference,var_ext):

    parameter={}
    var=[]
    tar = [{'Preset':{'sinlh01.mqsk030.k1':0,'sinbc02.mqsk110.k1':0,'sinbc02.mqsk350.k1':0,
                      'sindi01.mqsk030.k1':0,'s10bc02.mqsk110.k1':0,'s10bc02.mqsk350.k1':0,
                      'sarcl02.mqsk160.k1':0,'sarcl02.mqsk300.k1':0,'sarcl02.mqsk420.k1':0,
                      'satsy01.mqsk030.k1':0,'satsy01.mqsk270.k1':0,
                      'sinbc02.mqua120.k1':0,'sinbc02.mqua340.k1':0,
                      's10bc02.mqua120.k1':0,'s10bc02.mqua340.k1':0,
                      'sarcl02.mqua150.k1':0,'sarcl02.mqua430.k1':0,
                      'satbd01.mqua010.k1':0.50697,'satbd01.mqua030.k1':-0.618466,
                      'satbd01.mqua050.k1':0.13249,'satbd01.mqua070.k1':0.47403,
                      'satbd01.mqua090.k1':-0.74426,'satbd02.mqua030.k1':-1.05081}}]

    parameter['Preset'] = MatchParameter(var,tar)

    var = ['SINLH01-MQUA020', 'SINLH01-MQUA040','SINLH01-MQUA050','SINLH01-MQUA070','SINLH02-MQUA010']
    tar = [reference['Laser Heater']]
    parameter['Laser Heater']  = MatchParameter(var,tar)

    var = ['SINSB04-MQUA130','SINSB04-MQUA230']
    tar =[MatchReference('#e',{'MUX':0.2,'MUY':0.2})]
    parameter['SINSB04']  = MatchParameter(var,tar)

    var = ['SINLH02-MQUA410','SINLH03-MQUA030','SINLH03-MQUA040','SINLH03-MQUA060','SINLH03-MQUA080','SINSB03-MQUA130']
    tar =[{'Location':'sinsb04$start','SavedID':'SINSB04'}]
    parameter['Match to SINSB04']=MatchParameter(var,tar)

    var = ['SINBC01-MQUA020','SINBC01-MQUA050','SINBC01-MQUA070','SINBC01-MQUA090','SINBC01-MQUA110','SINDI01-MQUA020','SINDI01-MQUA070']
    tar =[reference['Bunch Compressor 1'],
          MatchReference('sinbc02.mbnd400',{'betax':5.1,'betay':100},1),
          MatchReference('sinbc01.mqua110$end',{'betax':100},1)]
    parameter['Bunch Compressor 1']=MatchParameter(var, tar)

    var = ['S10CB02-MQUA230','S10CB02-MQUA430']
    tar = [MatchReference('#e',{'MUX':0.1883,'MUY':0.1883})]
    parameter['Linac 1'] = MatchParameter(var, tar)


    var = ['S10CB01-MQUA230','S10CB01-MQUA430','SINDI02-MQUA020','SINDI02-MQUA030','SINDI02-MQUA050','SINDI02-MQUA060','SINDI02-MQUA090']
    tar =[{'Location':'s10cb02$start','SavedID':'Linac 1'}]
    parameter['Match to Linac 1']=MatchParameter(var,tar)


    var = ['S10CB07-MQUA430','S10CB08-MQUA230','S10CB08-MQUA430','S10CB09-MQUA230']
    tar = [MatchReference('s10bc01.mqua020$start',{'betax':4.88,'betay':16.11,'alphax':0.5546,'alphay':-1.81})]
    parameter['End of Linac 1']=MatchParameter(var, tar)


    var = ['S10BC01-MQUA020','S10BC01-MQUA040','S10BC01-MQUA060','S10BC01-MQUA080','S10BC01-MQUA100']
    tar = [reference['Bunch Compressor 2']]
    parameter['Bunch Compressor 2']=MatchParameter(var,tar)

    var = ['S20CB01-MQUA430','S20CB02-MQUA430']
    tar = [MatchReference('#e', {'MUX': 0.2, 'MUY': 0.2})]
    parameter['Linac 2'] = MatchParameter(var, tar)


    var = ['S10MA01-MQUA020','S10MA01-MQUA050','S10MA01-MQUA070','S10MA01-MQUA110','S10MA01-MQUA130']
    tar =[{'Location':'s20cb01$start','SavedID':'Linac 2'}]
    parameter['Match to Linac 2']=MatchParameter(var, tar)

    var = ['S20SY01-MQUA020','S20SY01-MQUA030','S20SY01-MQUA050','S20SY01-MQUA080']
    tar = [reference['Switchyard']]
    parameter['Match to Switchyard']=MatchParameter(var, tar)


    var = ['S20SY02-MQUA070','S20SY02-MQUA100','S20SY02-MQUA140','S20SY02-MQUA180','S20SY02-MKDC010','S20SY02-MKAC020','S20SY02-MKDC030','S20SY02-MKAC040','S20SY02-MKDC050']
    tar=[{'Script':'Scripts/switchyard_kicker.madx'}]
    parameter['Resonant Kicker']=MatchParameter(var, tar)


    var=['SATSY01-MQUA020','SATSY01-MQUA040','SATSY01-MQUA070','SATSY01-MQUA090','SATSY01-MQUA210','SATSY01-MQUA230','SATSY01-MQUA260','SATSY01-MQUA280',
         'SATSY01-MQUA300','SATSY02-MQUA010','SATSY02-MQUA110','SATSY02-MQUA120','SATSY02-MQUA230','SATSY03-MQUA010','SATSY03-MQUA040','SATSY03-MQUA070',
         'SATSY03-MQUA100','SATSY03-MQUA130','SATCL01-MQUA120','SATCL01-MQUA130','SATCL01-MQUA180','SATCL01-MQUA190']
    tar=[{'Script':'Scripts/switchyard.madx'}]
    val = 0.0
    if 'SYR56' in var_ext:
        val = var_ext['SYR56']
    uservar = {'LEAKDISP':val}
    parameter['Switchyard']=MatchParameter(var, tar, uservar)


    var=['SATDI01-MQUA025','SATDI01-MQUA080','SATDI01-MQUA220','SATDI01-MQUA230']
    tar = [reference['HERO']]
    parameter['Athos Diagnostics']  = MatchParameter(var,tar)

    var = ['SATCB01-MQUA230','SATCB01-MQUA430']
    tar =[MatchReference('#e',{'MUX':0.053,'MUY':0.053})]
    parameter['Athos Linac']  = MatchParameter(var,tar)

    var=['SATDI01-MQUA250','SATDI01-MQUA260','SATDI01-MQUA280','SATDI01-MQUA300']
    tar =[{'Location':'satcl01$start','SavedID':'Athos Linac'}]
    parameter['Athos Stopper']  = MatchParameter(var,tar)


    var=['SATMA01-MQUA050','SATMA01-MQUA120','SATMA01-MQUA140','SATMA01-MQUA160','SATMA01-MQUA180','SATMA01-MQUA230','SATMA01-MQUA250']
    tar = [reference['EEHG'],
           MatchReference('satma01.mqua140$start', {'betax': 70, 'betay': 70}, 1),
           MatchReference('satma01.mqua160$start', {'betax': 70, 'betay': 70}, 1),
           MatchReference('satma01.mqua180$start', {'betax': 70, 'betay': 70}, 1)]
    parameter['Athos EEHG'] = MatchParameter(var,tar)

    var = ['SATUN06-MQUA080','SATUN07-MQUA080']
    tar =[MatchReference('#e',{'MUX':0.15,'MUY':0.15})]
    parameter['Athos Undulator']  = MatchParameter(var,tar)

    var=['SATMA01-MQUA610','SATMA01-MQUA630','SATMA01-MQUA640','SATUN04-MQUA060','SATUN05-MQUA420']
    tar =[{'Location':'satun06$start','SavedID':'Athos Undulator'},
          MatchReference('satun07$end', {'x': 10, 'y': 10,'px':10}, 1)]
    parameter['Match to Athos Undulator'] = MatchParameter(var,tar)

    var = ['SATUN22-MQUA080','SATMA02-MQUA010','SATMA02-MQUA020','SATMA02-MQUA040','SATMA02-MQUA050','SATMA02-MQUA070']
    tar = [reference['Athos Deflector'],
           MatchReference('satun22$end', {'x':0, 'px': 0})]
    parameter['Match to Athos Deflector'] = MatchParameter(var,tar)

    var = ['S30CB01-MQUA430', 'S30CB02-MQUA430']
    tar = [MatchReference('#e', {'MUX': 0.1883, 'MUY': 0.1883})]
    parameter['Linac 3'] = MatchParameter(var, tar)

    var = ['S20SY03-MQUA020', 'S20SY03-MQUA030', 'S20SY03-MQUA050', 'S20SY03-MQUA060', 'S20SY03-MQUA100']
    tar = [{'Location': 's30cb01$start', 'SavedID': 'Linac 3'}]
    parameter['Match to Linac 3'] = MatchParameter(var, tar)

    var = ['S30CB01-MQUA430', 'S30CB01-MQUA430', 'S30CB01-MQUA430', 'S30CB01-MQUA430']
    tar = [reference['Linac3']]
    parameter['End of Linac 3'] = MatchParameter(var, tar)

    var = ['SARCL02-MQUA130','SARCL02-MQUA210','SARCL02-MQUA250',
           'SARCL02-MQUA310','SARCL02-MQUA350','SARCL02-MQUA460',
           'SARCL01-MQUA020', 'SARCL01-MQUA050', 'SARCL01-MQUA080',
           'SARCL01-MQUA100', 'SARCL01-MQUA140', 'SARCL01-MQUA190'
           ]
    tar = [{'Script': 'Scripts/ecol.madx'}]
    parameter['Aramis - ECOL'] = MatchParameter(var, tar)

    var = ['SARUN03-MQUA080', 'SARUN04-MQUA080']
    tar = [MatchReference('#e', {'MUX': 0.15, 'MUY': 0.15})]
    parameter['Aramis Undulator'] = MatchParameter(var, tar)

    var = ['SARMA01-MQUA010', 'SARMA01-MQUA060', 'SARMA01-MQUA080', 'SARMA01-MQUA120', 'SARMA01-MQUA140']
    tar = [{'Location': 'sarun03$start', 'SavedID': 'Aramis Undulator'},
           MatchReference('sarma01.mqua060', {'betax': 40}, 1)]
    parameter['Match to Aramis Undulator'] = MatchParameter(var, tar)

    var = ['SARUN19-MQUA080', 'SARUN20-MQUA080', 'SARBD01-MQUA020', 'SARBD02-MQUA030']
    tar = [MatchReference('SARBD02$END', {'betax': 250, 'betay': 250,'etax':100,'etay':1.2}, 1)]
    parameter['Aramis Dump'] = MatchParameter(var, tar)
    return parameter