#-- Lib

import numpy as np

from ess import lib_tw


#-- Input

path_name_lat    ='lattice.dat'
path_name_lat_tmp='lattice_tmp.dat'

path_name_res='tracewin.out'

#-- How to use the LATTICE class

lat=lib_tw.LATTICE(path_name_lat)

#-- Pick up the line # (I'm calling "idx") and elem # (I'm calling "idx_elem") of SPK field maps

idx_spk     =[]
idx_elem_spk=[]
for lat_i in lat.lst:
    if lat_i.typ=='FIELD_MAP':
        if lat_i.name_fmap=='Spoke_W_coupler':
            idx_spk.append(lat_i.idx)
            idx_elem_spk.append(lat_i.idx_elem)            

## # This does the same thing...
## for i in range(len(lat.lst)):
##     if lat.lst[i].typ=='FIELD_MAP':
##         if lat.lst[i].name_fmap=='Spoke_W_coupler':
##             idx_spk.append(i)
##             idx_elem_spk.append(lat.lst[i].idx_elem)
            
print 'idx (line #) of SPK field maps (starting from 0):'
print idx_spk,'\n'
print 'idx_elem (elem # of) SPK field maps (starting from 0):'
print idx_elem_spk,'\n'

#-- Give a name "SPK2" to the 2nd SPK and pick up its idx and idx_elem from the name

lat.lst[idx_spk[1]].name='SPK2'

for lat_i in lat.lst:
    if lat_i.name=='SPK2':
        idx_spk2     =lat_i.idx
        idx_elem_spk2=lat_i.idx_elem
        print 'idx (line #) of 2nd SPK:'     ,idx_spk2
        print 'idx_elem (elem #) of 2nd SPK:',idx_elem_spk2,'\n'
                
#-- Introduce an error command

err_comm_name='MY_ERROR_COMMAND'
err_comm_typ ='ERROR_CAV_NCPL_STAT'  # I know it's redundant but we also need this for now...
err_comm_para=[0,0,0,0,0,0,0,0,0]    # This also works: err_comm_para=[0]*100

err_comm=lib_tw.ERROR_CAV_NCPL_STAT(err_comm_name,err_comm_typ,err_comm_para)

print 'This is how it looks in the TraceWin syntax:'
print err_comm.get_tw(),'\n'

#-- Change parameters of the error command

err_comm.typ_dist='0'           # To make sure the dist type is constant
err_comm.E       =2e-2          # To introduce 2% field error
err_comm.phs_rf  =3.0*np.pi/180.0  # To introduce 3 deg phase error

print 'This is how the command looks after the changes:'
print err_comm.get_tw(),'\n'

#-- Insert the command in front of the 2nd SPK with the index "idx_spk2"

print 'This is how the lattice looks around the 2nd SPK looks originally:'
for i in range(idx_spk2-2,idx_spk2+3):
    print lat.lst[i].get_tw()
print ''

lat.lst.insert(idx_spk[1],err_comm)  # This does the insertion
lat.update_idx()                     # "magic word" to update idx and idx_elem    

# Update idx and idx_elem of SPK2
for lat_i in lat.lst:
    if lat_i.name=='SPK2':
        idx_spk2     =lat_i.idx
        idx_elem_spk2=lat_i.idx_elem
        print 'idx (line #) of 2nd SPK after the insertion:'     ,idx_spk2
        print 'idx_elem (elem #) of 2nd SPK after the insertion:',idx_elem_spk2,'\n'
print 'Note idx is +1 but no change for idx_elem (since we inserted a command).\n'

print 'This is how the lattice looks around the 2nd SPK after the change:'
for i in range(idx_spk2-2,idx_spk2+3):
    print lat.lst[i].get_tw()
print ''

#-- Save the new lattice to a file
            
lat.get_tw(path_name_lat_tmp)

#--

exit()












#-- Steerers

# Indices of steerers (can be done with names??) and append max to THIN_STEERING
idx_st_x=[]; idx_st_y=[]
for i in range(len(lat.lst)):
    # STEERER
    if lat.lst[i].typ=='STEERER':
        for j in range(i)[::-1]:
            if lat.lst[j].typ=='ADJUST_STEERER'   :
                idx_st_x.append(i); idx_st_y.append(i); break
            if lat.lst[j].typ=='ADJUST_STEERER_BY':
                idx_st_x.append(i); break
            if lat.lst[j].typ=='ADJUST_STEERER_BX':
                idx_st_y.append(i); break                                
            if lat.lst[j].idx_elem!=lat.lst[i].idx_elem:
                break
    # THIN_STEERING (assuming the same for x and y for the dual plane ones)
    if lat.lst[i].typ=='THIN_STEERING':
        for j in range(i)[::-1]:
            if lat.lst[j].typ=='ADJUST':
                lat.lst[i].max=lat.lst[j].max
                if lat.lst[j].var==1: idx_st_y.append(i)
                if lat.lst[j].var==2: idx_st_x.append(i)
            if lat.lst[j].idx_elem!=lat.lst[i].idx_elem-1:
                break

# Indices of physical steerer locations for STEERER (later to extract L and s)
idx_st_x_elem=[]
for i in idx_st_x:
    if lat.lst[i].typ=='THIN_STEERING': idx_st_x_elem.append(i)
    if lat.lst[i].typ=='STEERER'      :
        for j in range(i+1,len(lat.lst)):
            if lat.lst[j].idx_elem==lat.lst[i].idx_elem+1:
                idx_st_x_elem.append(j); break
idx_st_y_elem=[]
for i in idx_st_y:
    if lat.lst[i].typ=='THIN_STEERING': idx_st_y_elem.append(i)
    if lat.lst[i].typ=='STEERER'      :
        for j in range(i+1,len(lat.lst)):
            if lat.lst[j].idx_elem==lat.lst[i].idx_elem+1:
                idx_st_y_elem.append(j); break
                                
# Check max is defined
for i in set(idx_st_x+idx_st_y):
    if lat.lst[i].max<=0:
        print 'Max B/BL not defined for elem #'+str(lat.lst[i].idx_elem+1)+'. Exiting...'; exit()

#-- BPMs

idx_bpm_x=[i for i in range(len(lat.lst)) if lat.lst[i].typ=='DIAG_POSITION']
idx_bpm_y=idx_bpm_x[:]

###########################################################################

# Remove redundant BPMs to avoid over-constraint (not necessarily generic...)
k=0
while k < len(idx_bpm_x):
    if len([i for i in idx_st_x if i<idx_bpm_x[k]])<k+1: del idx_bpm_x[k]
    else                                               : k+=1
k=0
while k < len(idx_bpm_y):
    if len([i for i in idx_st_y if i<idx_bpm_y[k]])<k+1: del idx_bpm_y[k]
    else                                               : k+=1
    
###########################################################################
    
idx_elem_bpm_x=[lat.lst[i].idx_elem for i in idx_bpm_x]
idx_elem_bpm_y=[lat.lst[i].idx_elem for i in idx_bpm_y]

###########################################################################

# To do 1-to-1 with the "twist" pattern

#idx_elem_bpm[7],idx_elem_bpm[8]=idx_elem_bpm[8],idx_elem_bpm[7]

###########################################################################
        
#-- Define indices for blocks of R-matrix from "step" (may be too much...)

block_x=[]
for k in range(len(step_x)+1):
    if int(sum(step_x[:k]))>=len(idx_bpm_x): block_x.append(len(idx_bpm_x)      ); break
    else                                   : block_x.append(int(sum(step_x[:k])))
if block_x[-1]<len(idx_bpm_x): block.append(len(idx_bpm_x))
block_y=[]
for k in range(len(step_y)+1):
    if int(sum(step_y[:k]))>=len(idx_bpm_y): block_y.append(len(idx_bpm_y)      ); break
    else                                   : block_y.append(int(sum(step_y[:k])))
if block_y[-1]<len(idx_bpm_y): block.append(len(idx_bpm_y))

#-- R matrix

Rx=Pool(Ncpu).map(Rx_column,range(len(idx_st_x)))
Ry=Pool(Ncpu).map(Ry_column,range(len(idx_st_y)))

Rx=array(Rx).transpose()
Ry=array(Ry).transpose()

if len(Rx)>len(Rx[0]): print '[# of BPMs] > [# of steerers] for x, exiting...'; exit()
if len(Ry)>len(Ry[0]): print '[# of BPMs] > [# of steerers] for y, exiting...'; exit()
if len(Rx)<len(Rx[0]): print '[# of BPMs] < [# of steerers] for x, exiting...'; exit()
if len(Ry)<len(Ry[0]): print '[# of BPMs] < [# of steerers] for y, exiting...'; exit()
            
#-- Main part

data=Pool(Ncpu).map(job,range(Nrun))
data=array(data).transpose()

s  =data[0][0]
x  =array(data[1].tolist()).transpose()  # array not applied to the 3rd level ??
y  =array(data[2].tolist()).transpose()  # array not applied to the 3rd level ??
BLy=array(data[3].tolist()).transpose()  # array not applied to the 3rd level ??
BLx=array(data[4].tolist()).transpose()  # array not applied to the 3rd level ??

#-- Writing

# x
with open('x.out','w') as file:
    # Header
    file.write('s  ')
    for n in range(Nrun): file.write('x%03d  '%n)
    file.write('\n')
    # Data
    for k in range(len(s)):
        file.write('%.4f  '%s[k])
        for n in range(Nrun): file.write('%.5f  '%x[k][n])
        file.write('\n')

# y
with open('y.out','w') as file:
    # Header
    file.write('s  ')
    for n in range(Nrun): file.write('y%03d  '%n)
    file.write('\n')
    # Data
    for k in range(len(s)):
        file.write('%.4f  '%s[k])
        for n in range(Nrun): file.write('%.5f  '%y[k][n])
        file.write('\n')
        
# x steering (BLy [Gm])
with open('st.x.out','w') as file:
    # Header
    file.write('##  s  ')
    for n in range(Nrun): file.write('BLy%03d  '%n)
    file.write('\n')
    # Data
    for k in range(len(idx_st_x)):
        file.write('%03d  '%(k+1)+'%.4f  '%lat.lst[idx_st_x_elem[k]].s)
        for n in range(Nrun): file.write('%.4f  '%(1e4*BLy[k][n]))
        file.write('\n')

# y steering (BLx [Gm])
with open('st.y.out','w') as file:
    # Header
    file.write('##  s  ')
    for n in range(Nrun): file.write('BLx%03d  '%n)
    file.write('\n')
    # Data
    for k in range(len(idx_st_y)):
        file.write('%03d  '%(k+1)+'%.4f  '%lat.lst[idx_st_y_elem[k]].s)
        for n in range(Nrun): file.write('%.4f  '%(1e4*BLx[k][n]))
        file.write('\n')
        
#-- Ending

exit()

#-- Obsolete

## def job(n):
##     '''
##         Arbitrary x step correction for the random seed seed[n].
##     '''
##     # Indices of sub-matrices, e.g., step=[8,7] => block=[8,15]
##     block=[]
##     for k in range(len(step)+1):
##         if int(sum(step[:k]))>=len(idx_elem_bpm): block.append(len(idx_elem_bpm )); break
##         else                                    : block.append(int(sum(step[:k])))    
##     if block[-1]<len(idx_elem_bpm): block.append(len(idx_elem_bpm))
                
##     # Define child calc dir
##     path_cal_n=path_cal+'/tmp_'+str(n)

##     # Set-up TraceWin
##     opt_tw=setup_tw(path_cal_n,seed[n])
##     lat_n =LATTICE(path_cal_n+'/'+file_name_lat[::-1].partition('/')[0][::-1],[])

##     # Loop for steps
##     for b in range(len(block)-1):

##         # Borders of the sub-matrix
##         k0=block[b]; k1=block[b+1]
        
##         # Intermediate trajectory
##         lat_n.get_tw(path_cal_n+'/'+file_name_lat[::-1].partition('/')[0][::-1])
##         call(opt_tw,shell=True)
##         tw=PARTRAN(path_cal_n+'/tracewin.out')

##         # Temp R matrices, using only the [k0,k1-1] block
##         Rx_k=zeros((len(Rx),len(Rx[0]))); Rx_k[k0:k1,k0:k1]=Rx[k0:k1,k0:k1]  # Can't do this with a list
##         Ry_k=zeros((len(Ry),len(Ry[0]))); Ry_k[k0:k1,k0:k1]=Ry[k0:k1,k0:k1]  # Can't do this with a list
        
##         # Correction
##         st_x_nom=lstsq(Rx_k,[-tw.x[i] for i in idx_elem_bpm])[0]  # Can include BPM error here
##         st_y_nom=lstsq(Ry_k,[-tw.y[i] for i in idx_elem_bpm])[0]  # Can include BPM error here
        
##         # Apply correction
##         for k in range(k0,k1):
##             i=idx_st_x[k]; st_x=st_x_nom[k]*lat.lst[i].max
##             if lat_n.lst[i].typ=='STEERER'      : lat_n.lst[i].By =st_x
##             if lat_n.lst[i].typ=='THIN_STEERING': lat_n.lst[i].BLy=st_x
##             i=idx_st_y[k]; st_y=st_y_nom[k]*lat.lst[i].max
##             if lat_n.lst[i].typ=='STEERER'      : lat_n.lst[i].Bx =st_y
##             if lat_n.lst[i].typ=='THIN_STEERING': lat_n.lst[i].BLx=st_y                
        
##     # Trajectory after correction
##     lat_n.get_tw(path_cal_n+'/'+file_name_lat[::-1].partition('/')[0][::-1])
##     call(opt_tw,shell=True)
##     tw=PARTRAN(path_cal_n+'/tracewin.out')

##     # Save steering strengths
##     BLy=[]; BLx=[]
##     for k in range(len(idx_st_x)):
##         i=idx_st_x[k]; j=idx_st_x_elem[k]
##         if lat_n.lst[i].typ=='STEERER'      : BLy.append(lat_n.lst[i].By *lat_n.lst[j].L)
##         if lat_n.lst[i].typ=='THIN_STEERING': BLy.append(lat_n.lst[i].BLy               )
##     for k in range(len(idx_st_y)):
##         i=idx_st_y[k]; j=idx_st_y_elem[k]
##         if lat_n.lst[i].typ=='STEERER'      : BLx.append(lat_n.lst[i].Bx *lat_n.lst[j].L)
##         if lat_n.lst[i].typ=='THIN_STEERING': BLx.append(lat_n.lst[i].BLx               )            
        
##     # Clean
##     call('rm -rf '+path_cal_n,shell=True)
##     print 'job #'+str(n)+' done.'
    
##     return [tw.s,tw.x,tw.y,BLy,BLx]