#!/usr/bin/python3
# extract index and data of various  elements and print them
"""
still missing: get TYPE of magnets etc... which must be got form another source
"""
import at
import os

def lattlister(ring,geo,selector):
    idlist = ring.uint32_refpts(selector)
    for i in idlist:
        elem = ring[i]
        ename = elem.FamName
        length = elem.Length
        x=geo[i].x
        y=geo[i].y
        angle=geo[i].angle
        if hasattr(elem,'corID'):
            iscorr = '1'
        else:
            iscorr = '0'
        if hasattr(elem,'SkewID'):
            isskew = '1'
        else:
            isskew = '0'
        spos = ring.get_s_pos(i)[0]
        print(i,'\t',ename,'\t',length,'\t',spos,'\t',iscorr,'\t',isskew,'\t',x,'\t',y,'\t',angle)
        #return i,ename,length,spos,iscorr,isskew

homedir=os.path.expanduser('~')
lattname = "srElettra2_9_4_beta.mat"
#lattname = "srElettra2_high_betax_long_straights_ring.m"
latticefile = homedir + "/src/gitlab/dt/machine/lattice/elettra2/" + lattname
ring=at.load_lattice(latticefile, energy=2.4e9)
"""
up to now the low_beta lattice file can be ignored. It will have the same structure of rElettra2_9_4_beta but different
magnet strenghts
ring=at.load_lattice("../../machine/lattice/elettra2/srElettra2_low_beta.mat",energy=2.4e9)
"""
#ring geometry
geo,radius=ring.get_geometry(centered=True)
print('i','\t','ename','\t','length','\t','spos','\t','iscorr','\t','isskew','\t','x','\t','y','\t','angle')


# find index and print of quadrupoles in ring
lattlister(ring,geo,at.Quadrupole)

    
#find index and print of reverse bends Q24R quads
#lattlister(ring,geo,"QAB*")
    
# find index and print of sextupoles in ring
lattlister(ring,geo,at.Sextupole)

#find index and print of octupoles in ring
lattlister(ring,geo,at.Octupole)
    
#find index and print of b64 dipoles in ring
#this is triky..... We have to set name to indexes ... or modify something in the lattice.
#lattlister(ring,geo,"b1")
    
#find index and print 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
lattlister(ring,geo,"BPM*")

# get corectors  C
#lattlister(ring,geo,"C*")

#get correctors c
#lattlister(ring,geo,"c*")

#get RF cavities
lattlister(ring,geo,"CAV")