#!/usr/bin/python3
# extract index of various  elements and print them
import at
import re
import os
#import matplotlib.pylab as mp
homedir=os.path.expanduser('~')
#lattname = "srElettra2_9_4_beta.mat"
lattname = "srElettra2_high_betax_long_straights_ring.m"
#lattname = "srElettra2_9_4_beta_test_digital_twins.mat"
latticefile = homedir + "/src/gitlab/dt/machine/lattice/elettra2/" + lattname
ring=at.load_lattice(latticefile, energy=2.4e9)
refpts = range(len(ring) + 1)


# find index of quadrupoles in ring

iquads = ring.uint32_refpts(at.Quadrupole)
print("Name, Length,K,PolynomA,PolynomB")
for i in iquads:
    q=ring[i]
    print(q.FamName,q.Length,q.K,q.PolynomA[1],q.PolynomB[1])
    
#find index of reverse bends Q24AB quads
print("Name, Length,K,PolynomA,PolynomB")
irb = ring.uint32_refpts("QAB*")
for i in irb:
    q=ring[i]
    print(q.FamName,q.Length,q.K,q.PolynomA[1],q.PolynomB[1])
    
# find index of sextupoles in ring

isexts = ring.uint32_refpts(at.Sextupole)
print("Name, Length,K,PolynomA,PolynomB,is_corr,is_skew")
for i in isexts:
    q=ring[i]
    dorder = q.DefaultOrder
    morder = q.MaxOrder
    if hasattr(q,'corID'):
        iscorr = 'is_corr'
    else:
        iscorr = '-------'
    if hasattr(q,'SkewID'):
        isskew = 'is_skew'
    else:
        isskew = '-------'
    
    print(q.FamName,q.Length,q.K,q.PolynomA[dorder],q.PolynomB[morder],iscorr,isskew)

#find index of octupoles in ring
iocts = ring.uint32_refpts(at.Octupole)
print("Name, Length,K,PolynomA,PolynomB,is_corr,is_skew")
for i in iocts:
    q=ring[i]
    dorder = q.DefaultOrder
    morder = q.MaxOrder
    if hasattr(q,'corID'):
        iscorr = 'is_corr'
    else:
        iscorr = '-------'
    if hasattr(q,'SkewID'):
        isskew = 'is_skew'
    else:
        isskew = '-------'
    print(q.FamName,q.Length,q.K,q.PolynomA[dorder],q.PolynomB[morder],iscorr,isskew)
    
#find index of b64 dipoles in ring
#this is triky..... We have to set name to indexes ... or modify something in the lattice.

ib64 = ring.uint32_refpts("b1")
for i in ib64:
    q=ring[i]
    dorder = q.DefaultOrder
    morder = q.MaxOrder
    print(q.FamName,q.Length,q.K,q.PolynomA[dorder],q.PolynomB[morder])
    
#find index of b80 dipoles in ring
#this is enven trikier! triky..... We have to set name to indexes and elemets are a composition of pieces

# TO DO

# get bpms and their s position

ibpms = ring.uint32_refpts(at.Monitor)
print("num bpm:", len(ibpms))
print("ind. ,name,  s[m]")
for i in ibpms:

    bpm=ring[i]
    sp=ring.get_s_pos(i)
    print(i, bpm.FamName,sp[0])
    
# get corectors  and their s position
ichv = ring.uint32_refpts("chv*")
print("ind. ,name,  s[m]")
for i in ichv:
    chv=ring[i]
    sp=ring.get_s_pos(i)
    print(i,chv.FamName,sp[0])

#get RF cavities
selector = "CAV"
irfcavs = ring.uint32_refpts(selector)
for i in irfcavs:
    rfcav=ring[i]
    srf=ring.get_s_pos(i)
    print(i,rfcav.FamName,srf[0])