#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Feb 20 20:07:11 2025

OSP2CSUv9.py

Es como el v8, pero lee los parámetros astrométricos del mismo fichero que la OSP.

Lee el fichero XML de la máscara, NO del proyecto, Y el fichero CSU.

Genera un dibujo con las posiciones de las rendijas a partir del fichero CSU
y los objetos, a partir del XML.

No saca un fichero CSU alternativo.

Añado que pinte un punto gordo verde en la posición (X,Y) de los pix virtuales de los objetos

No usa el fichero de las yv, que es solo aproximado.

@author: fgl
@version: 9
@date: 17/04/25
"""

import numpy as np
import xml.etree.ElementTree as ET
import astropy.io.ascii as at
from astropy.coordinates import SkyCoord
from astropy.coordinates import Angle
from astropy import units as u
import matplotlib.pyplot as plt
import sys


def findYPoint(xa,xb,ya,yb,xc):
    m = (ya - yb) / (xa - xb)
#    print (m)
    yc = (xc - xb) * m + yb
    return yc

nod=-1e10
if (len(sys.argv) - 1)<1:
#    sys.exit('Este script necesita 2 argumentos: XML  y CSU de entrada')
    print('This script takes 2 arguments: XML and CSU input files')
    print('If just 1 is given then files are <name>.xml & <name>.csu')
    print('Option: a second or third argument beginning with + or -')
    print('   specifying the nodding throw in arcsecs')
    sys.exit('Try again')
elif (len(sys.argv) - 1)==1:
    fxml=sys.argv[1]+'.xml'
    fcsu=sys.argv[1]+'.csu'    
elif (len(sys.argv) - 1)==2:
    if sys.argv[2][0] not in ['+','-']:
        fxml=sys.argv[1]
        fcsu=sys.argv[2]
    else:
        fxml=sys.argv[1]+'.xml'
        fcsu=sys.argv[1]+'.csu'    
        nod=float(sys.argv[2])
elif (len(sys.argv) - 1)==3:
    nod=float(sys.argv[3])
    
fichresu=sys.argv[1]+'.summary'
tree=ET.parse(fxml)
data_csu=at.read(fcsu)
with open(fcsu,'r') as f1:
    idcsu=next(f1)
    idcsu=next(f1).split('\n')[0].split(': ')[1]
#print ('UUID: ',idcsu)

parname=sys.path[0]+"/calibration/EMIR_parameters.dat"
with open(parname,'r') as f_in:
    for line in f_in:
        if 'PLTSC' in line: pltsc=float(line.split('#')[0])
        elif 'PXSZ' in line: pxsz=float(line.split('#')[0])
        elif 'OFFSET_IPA' in line: paoff=float(line.split('#')[0])

nx=2047
ny=2047
off=1               # offset de origen 0 en python
# paoff=180.0552      # offset entre PA mask y EMIR IPA

fmt='{:03d}{:8.3f}{:2d}'+'{:8.2f}'*4
nbars=110
nslits=nbars//2
fov=400.				 #fov in arcs
hs=fov/nslits
hsr=hs*np.pi/180./3600.
# pxsz=0.166			  # 166 micras por pixel en CSU
# pltsc=1.1717        # GTC plate scale at EMIR CSU
hs2=hs/2/pltsc/pxsz   # semialtura de la rendija en px.
pxsc=pxsz*pltsc       # px projection on sky, arcsec

xarrow=-.032
yarrow=.4
dxarrow=0
dyarrow=.1


vpx_x=np.zeros(nslits)
vpx_y=np.zeros(nslits)
x0=1024.5
y0=1024.5
slwi=np.zeros(nslits)
slfl=np.zeros(nslits,dtype=int)
sljo=np.zeros(nslits,dtype=int)
slobjlist=[]
join_typ=['Double','Triple','Quadruple','>4']
njoin=len(join_typ)
objco=[]
slobjid=[]
slobjcom=[]

#calname="/Users/fgl/trabajo/emir/afinado/CSU/cal/csu_cal_tot_abr18.dat"
calname=sys.path[0]+"/calibration/EMIR_CSUcal.dat"
data=at.read(calname)
csup=data['csup']

yv=np.arange(nslits)*hs2*2 + hs2

root=tree.getroot()

#print (root.attrib)
ra=float(root.find('alpha').text)
dec=float(root.find('delta').text)
pa=Angle(float(root.find('angle').text),u.deg)
maskcoord=SkyCoord(ra,dec,unit='deg')
idmask=root.find('maskUUID').text
print (maskcoord.ra.to(u.hourangle),maskcoord.dec,pa.deg)
if idcsu!=idmask :
    print ('****Warning: CSU and XML IDs mismatch: CSU: ',idcsu,' MASK: ',idmask)
else:
    print ('UUID: ',idmask)
slitlist=root.find('slitsList')
nobj=0
for slit in slitlist.findall('slit'):
  slid=int(slit.get('id'))-1
#  print(slid+1)      
  slwi[slid]=float(slit.find('width').text)
  slfl[slid]=int(data_csu[slid*2+1][2])
  if slit.find('isJoined')!= None:
#      if slit.find('isJoined').text == 'true':
      sljo[slid]=int(slit.find('joinedRefSlit').text)-1
#      print (slid,sljo[slid])
  slobj=slit.find('affectedObject')
  if slobj != None:
      sljo[slid]=slid
      slobjlist=np.append(slobjlist,slid)
      slobjid.append(slobj.find('ObjectId').text)
      slobjcom.append(slobj.find('Comments').find('libelleFrom').text)
      slobjco=slobj.find('objCoord')
      objra=slobjco.find('objCoord').find('ra').text
      objdec=slobjco.find('objCoord').find('dec').text
      # print ('\tObject:',objra,objdec,' slit:',slid+1)
      objco.append(SkyCoord(objra+objdec, unit=(u.hourangle, u.deg)))
      q=np.arctan(np.tan(objco[-1].dec.rad)/np.cos(objco[-1].ra.rad-maskcoord.ra.rad))
      eta0=np.tan(q - maskcoord.dec.rad)
      xi0=np.cos(q)*np.tan(objco[-1].ra.rad-maskcoord.ra.rad) / \
          np.cos(q - maskcoord.dec.rad)
      xi=Angle(xi0*np.cos(pa.rad) - eta0*np.sin(pa.rad), u.rad)
      eta=Angle(xi0*np.sin(pa.rad) + eta0*np.cos(pa.rad), u.rad)
      # print(('\t\tCoord. estandar: '+'{:9.3f}'*2).format(xi.arcsec,eta.arcsec))

# inversión horizontal del detector por la cámara. Por eso +eta
      vpx_x[nobj]=x0+eta.arcsec/pltsc/pxsz
      vpx_y[nobj]=y0-xi.arcsec/pltsc/pxsz
      # print(('\t\tVirPix X Y:'+'{:9.3f}'*2).format(vpx_x[nobj],vpx_y[nobj]))
      nobj+=1


f_resu = open(fichresu,'w')
nsl_obj=np.ones(nobj,dtype=int)
slobjlist=slobjlist.astype(int)
print('Mask name: ',fcsu.split('/')[-1],file=f_resu)
print('Mask-id: ',idmask,file=f_resu)
print('\n','*'*10,' Summary ','*'*10,file=f_resu)
print(' '*4,'Pointing coordinates:',maskcoord.ra.to_string\
       (u.hourangle,precision=3,pad=True),\
       maskcoord.dec.to_string(precision=2, alwayssign=True, pad=True),\
       file=f_resu)
print(' '*4,'Mask position angle [deg]:',\
              str('%22.3f' % pa.deg),file=f_resu)
print(' '*4,'EMIR IPA [deg]:',str('%33.3f' % (pa.deg+paoff)),file=f_resu)
print(' '*4,'No. of assigned objects: ',str('%23d'%nobj),file=f_resu)
nass=len(sljo[sljo>0])
ind=np.where(slfl==2)[0]
nrefs=len(np.unique(sljo[ind]))
d = {x:sljo[sljo>0].tolist().count(x) for x in sljo[sljo>0].tolist()}
nsing=sum(value == 1 for value in d.values())
njoi=sum(value > 1 for value in d.values())
print(' '*4,'No. of open/assigned slits:',str('%21d'%nass),file=f_resu)
print(' '*4,'No. of reference slits:',str('%25d'%nrefs),file=f_resu)
print(' '*4,'No. of single slits:',str('%28d'%nsing),file=f_resu)
print(' '*4,'No. of joined slits:',str('%28d'%njoi),file=f_resu)
for j in range(njoin):
    nval=sum(value == (j+2) for value in d.values())
#    print(' '*26,join_typ[j],str('%17d'%nval),file=f_resu)
    print(' '*26,str('%-9s'%join_typ[j]),str('%13d'%nval),file=f_resu)

for i in range (nobj):
    print('',file=f_resu)
    slid=slobjlist[i]
    if slfl[slid]==2:
        ref=" ****Ref. Object****"
    else:
        ref=''
    if len(sljo[sljo==slid])==1:
        print(' '*4,'Slit: ',slid+1,' width:',slwi[slid],'arcsec',ref,file=f_resu)
    else:
        ind=np.where(sljo==slid)[0]
        print(' '*4,'Slits: ',ind[0]+1,'-',ind[-1]+1,\
              ' width:',slwi[slid],'arcsec',ref,file=f_resu)
    print(' '*10,'Assigned object:',slobjid[i],file=f_resu)
    print(' '*27,objco[i].ra.to_string(u.hourangle,precision=2,pad=True),\
          objco[i].dec.to_string(precision=2, alwayssign=True, pad=True),file=f_resu)
    print(' '*27,slobjcom[i],file=f_resu)
f_resu.close()          

print ("\n\t*****\n")

plt.close('all')     
for i in range(0,nbars,2):
    if i==0:
        plt.close('all')
        fig=plt.figure(figsize=(6,6),dpi=300)
        plt.axis('off')
        plt.plot([0,1,1,0,0],[0,0,1,1,0],'k-')
    j=i//2
    left=j+1
    px=data['B'+str(left)+'_VirXpx']
    id1=(np.abs(csup-data_csu[i][1])).argmin()
    if csup[id1]>data_csu[i][1]:
       id2=max(0,id1-1)
    else:
       id2=min(len(csup)-1,id1+1)

    xv_left=findYPoint(csup[id1],csup[id2],px[id1],px[id2],data_csu[i][1])
   
    right=left+nslits
    px=data['B'+str(right)+'_VirXpx']
    id1=(np.abs(csup-data_csu[i+1][1])).argmin()
    if csup[id1]>data_csu[i+1][1]:
        id2=max(0,id1-1)
    else:
        id2=min(len(px)-1,id1+1)

    xv_right=findYPoint(csup[id1],csup[id2],px[id1],px[id2],data_csu[i+1][1])

    xl=(xv_left-off)/nx
    x=np.array([0,xl,xl,0])
    y1=max(yv[j]-hs2-off,0)/ny
    y2=min(ny,yv[j]+hs2-off)/ny
    y=np.array([y1,y1,y2,y2])
    plt.plot(x,y,'k-')    
    xr=(xv_right-off)/nx
    x=np.array([1,xr,xr,1])
    plt.plot(x,y,'k-')

    xpcsu=(xv_left+xv_right)/2 
            
    if data_csu[i][3]<-99:
        xposp=1024.5
    else:
        xposp=data_csu[i][3]
        yposp=data_csu[i][4]
        plt.plot ((xposp-off)/nx,(yposp-off)/ny,'mo',markersize=3.5)
    
    if np.abs(xpcsu-xposp)>1:
        flag=' *****'
    else:
        flag=''

    print ('\tSlit:',str('%02d'%(i//2+1)),' VirPixX CSU OSP:',\
            str('%7.2f'% xpcsu),str('%7.2f'% xposp),flag)

#    print (left)
    if (j%2==0):
#        print(left,'left')
        plt.text(-.01,(y2+y1)/2-.007,str(left),fontsize=6,\
                 horizontalalignment='right')
    else:
#        print(left,'right')
#            plt.text(str('%d' % (i+1)),1.1,(y2+y1)/2,fontsize=6)
        plt.text(1.005,(y2+y1)/2-.007,str(left),fontsize=6,\
                 horizontalalignment='left')


for i in range (nobj):
    x=(vpx_x[i]-off)/nx
    y=(vpx_y[i]-off)/ny
    slid=slobjlist[i]
    if slfl[slid]==2:
        plt.plot (x,y,'rx',markersize=10)
    else:
        plt.plot (x,y,'bx',markersize=10)


plt.title(idmask)
plt.suptitle (sys.argv[1].split('/')[-1])
if nod > -1e9:
    if nod>0:
        yarrow+=1.5*dyarrow
        dyarrow=-dyarrow
    plt.arrow(xarrow, yarrow, dxarrow, dyarrow, width=.01,color='red')

    plt.text (xarrow-.005,yarrow+dyarrow/2+.005,str('%.1f'%nod)+'"',\
         fontsize=8,horizontalalignment='right')

fig.savefig(sys.argv[1].split('.xml')[0]+'.png')
plt.close('all')
#plt.show()