#! python
""" Code to simulate Stokes parameters
    
Context : SRP
Module  : SRPTNGPAOLOStokesSim
Author  : Stefano Covino
Date    : 16/02/2017
E-mail  : stefano.covino@brera.inaf.it
URL:    : http://www.merate.mi.astro.it/utenti/covino
Purpose : Simulate Stokes parameters for an observation.

Usage   : SRPTNGPAOLOStokesSim [-h] [-a n k] -c RA DEC [-d detoff] -e Exp
            [-f file] [-l lambda angle] -o file [-p pos] -s Q
            U V -t 'YYYY/MM/DD HH:MM:SS' [-v] [--version]
            [-w wave] [-z q0 u0 v0]
    
            -a Aluminium refractive and extinction coefficient multiplicative factors
            -c Object coordinates
            -d Detector offset (deg)
            -e Observation length (sec)
            -f Input FITS file with fit parameters.
            -l lambda2 or lambda4 (2, 4) and angle
            -o Output FITS file
            -p Detector position angle
            -s Object normalized Stokes parameters
            -t Observation time
            -w Observation wavelength (micron)
            -z Normalized instrumental polarization
    
History : (26/02/2012) First version.
        : (20/09/2012) Hours angles phased.
        : (29/11/2012) Correct sign for position angle.
        : (31/03/2013) Wavelength in output file.
        : (25/09/2013) Update for the latest atpy release.
        : (13/02/2017) Python3 porting.
        : (16/02/2017) Porting to astropy.
"""

__version__ = '0.2.3'


import argparse, math, sys
import ephem, numpy
import SRPTNG.PAOLO as STP
from SRP.SRPMath.PhaseAngle import PhaseAngle
from SRP.SRPPolarimetry.AluminiumRefractiveIndex import AluminiumRefractiveIndex
from SRPTNG.PAOLO.TNGMuellerMatrix import TNGMuellerMatrix
from SRPTNG.PAOLO.TNGMuellerMatrixPlate2 import TNGMuellerMatrixPlate2
from SRPTNG.PAOLO.TNGMuellerMatrixPlate4 import TNGMuellerMatrixPlate4
from SRPTNG.GetObj import GetObj
from SRPTNG.GetTNGSite import GetTNGSite
from SRP.SRPSky.HourAngle import HourAngle
from SRP.SRPSky.ParallacticAngle import ParallacticAngle
from SRPTNG.PAOLO.StokesOffsetVector import StokesOffsetVector
from astropy.table import Table, Column




parser = argparse.ArgumentParser()
parser.add_argument("-a", "--alum", action="store", nargs=2, type=float, help="Aluminium refractive and extinction coefficient multiplicative factors", metavar=('n','k'),default=(1.0,1.0))
parser.add_argument("-c", "--coord", action="store", nargs=2, help="Object coordinates", metavar=('RA','DEC'),required=True)
parser.add_argument("-d", "--detoff", action="store", type=float, help="Detector offset (deg)", metavar='detoff',default=0.0)
parser.add_argument("-e", "--exptime", action="store", type=float, help="Observation length (sec)", metavar='Exp',required=True)
parser.add_argument("-f", "--fitsfitfile", action="store", help="Input FITS file with fit parameters.", metavar="file")
parser.add_argument("-l", "--lamina", action="store", nargs=2, type=float, help="lambda2 or lambda4 (2, 4) and angle", metavar=('lambda','angle'))
parser.add_argument("-o", "--outfile", action="store", help="Output FITS file", required=True, metavar='file')
parser.add_argument("-p", "--posangle", action="store", type=float, help="Detector position angle", default=0.0, metavar='pos')
parser.add_argument("-s", "--stokes", action="store", type=float, nargs=3, help="Object normalized Stokes parameters", metavar=('Q', 'U', 'V'), required=True)
parser.add_argument("-t", "--obsepoch", action="store", help="Observation time", required=True, metavar="'YYYY/MM/DD HH:MM:SS'")
parser.add_argument("-v", "--verbose", action="store_true", help="Fully describe operations")
parser.add_argument("--version", action="version", version=__version__)
parser.add_argument("-w", "--wave", action="store", type=float, help="Observation wavelength (micron)", default=0.55, metavar='wave')
parser.add_argument("-z", "--zero", action="store", nargs=3, type=float, help="Normalized instrumental polarization", metavar=('q0','u0','v0'),default=(0.,0.,0.,))
options = parser.parse_args()


#
if options.lamina:
    if options.lamina[0] != 2 and options.lamina[0] != 4:
        parser.error("Only lambda2 and lambda4 plates are availbale.")
#
if (options.alum or options.detoff or options.posangle or options.zero) and options.fitsfitfile:
    if options.verbose:
        print("Table data supersede command line data")
#
if options.fitsfitfile:
    try:
        tf = Table.read(options.fitsfitfile,format='fits')
    except IOErrror:
        parser.error("Incorrect fit FITS file: %s" % options.fitsfitfile)
#
site = GetTNGSite()
obj = GetObj(options.coord[0],options.coord[1])
site.date = options.obsepoch
obj.compute(site)
if options.exptime <= 0.0:
    parser.error("Exposure time must be positive.")
if options.verbose:
    print("Object coordinates    : %s %s" % (obj.a_ra, obj.a_dec))
    print("Observation epoch     : %s" % site.date)
    print("Observation length    : %.0f sec" % options.exptime)
    print("Stokes parameters     : I=%.3f Q=%.3f U=%.3f V=%.3f" % (1.0, options.stokes[0], options.stokes[1], options.stokes[2]))
#
if options.fitsfitfile:
    try:
        nn = tf[STP.N][0]
        kk = tf[STP.K][0]
        offoff = tf[STP.DETOFF][0]
        q0q0 = tf[STP.Q0][0]
        u0u0 = tf[STP.U0][0]
        v0v0 = tf[STP.V0][0]
    except IndexError:
        parser.errpr("FITS Stokes parameter file without the expected entries.")
else:
    nn = options.alum[0]
    kk = options.alum[1]
    offoff = options.detoff
    q0q0 = options.zero[0]
    u0u0 = options.zero[1]
    v0v0 = options.zero[2]
#
nf, kf = AluminiumRefractiveIndex()
n = nf(options.wave)*nn
k = kf(options.wave)*kk
#
if q0q0 > 1. or u0u0 > 1 or v0v0 > 1 or (q0q0**2 + v0v0**2 + u0u0**2) > 1:
    parser.error("Unrealistic instrumental polarization.")
if nn <= 0.0 or kk <= 0:
    parser.error("Multiplicative factors must be positive.")
#
if options.verbose:
    print("Observation wavelength        : %.3f" % options.wave)
    print("Adopted refractive index      : %.2f" % n)
    print("Adopted extinction coefficient: %.2f" % k)
    print("Detector position angle (deg) : %.2f" % options.posangle)
    print("Detector offset (deg)         : %.1f" % offoff)
    print("Instrumental polarization     : Q0=%.1g, U0=%.1g, V0=%.1g" % (q0q0, u0u0, v0v0))
    if options.lamina:
        if options.lamina[0] == 2:
            print("Lambda2 plate with angle %.1f" % options.lamina[1])
        elif options.lamina[0] == 4:
            print("Lambda4 plate with angle %.1f" % options.lamina[1])
#
sto = [1.0, options.stokes[0], options.stokes[1], options.stokes[2]]
Stokes = numpy.matrix(sto).transpose()
#
orgdate = site.date
date = []
time = []
az = []
alt = []
hourang = []
parang = []
i = []
q = []
u = []
v = []
for en in numpy.linspace(0.0,options.exptime/60.0,math.ceil(options.exptime/60.0)):
    site.date = orgdate + en*ephem.minute
    sdate = str(site.date)
    date.append(sdate.split()[0])
    time.append(sdate.split()[1])
    obj.compute(site)
    az.append(math.degrees(float(obj.az)))
    alt.append(math.degrees(float(obj.alt)))
    hang = PhaseAngle(HourAngle(math.degrees(float(obj.a_ra)),site),-180.0,180.0)
    hourang.append(hang)
    pang = ParallacticAngle(obj,site)
    parang.append(pang)
    if options.lamina:
        if options.lamina[0] == 2:
            s = TNGMuellerMatrixPlate2(pang,n,k,rot=options.lamina[1],offset=(offoff-options.posangle))*Stokes+StokesOffsetVector(q0q0,u0u0,v0v0)
        elif options.lamina[0] == 4:
            s = TNGMuellerMatrixPlate4(pang,n,k,rot=options.lamina[1],offset=(offoff-options.posangle))*Stokes+StokesOffsetVector(q0q0,u0u0,v0v0)
    else:
        s = TNGMuellerMatrix(pang,n,k,offset=(offoff-options.posangle))*Stokes+StokesOffsetVector(q0q0,u0u0,v0v0)
    i.append(s[0,0])
    q.append(s[1,0]/s[0,0])
    u.append(s[2,0]/s[0,0])
    v.append(s[3,0]/s[0,0])
#    print i[-1],q[-1],u[-1],v[-1]
#
tout = Table()
tout.add_column(Column(numpy.array(date),STP.DATE))
tout.add_column(Column(numpy.array(time),STP.TIME))
tout.add_column(Column(numpy.array(az),STP.AZ))
tout.add_column(Column(numpy.array(alt),STP.ALT))
tout.add_column(Column(numpy.array(hourang),STP.HOURANG))
tout.add_column(Column(numpy.array(parang),STP.PARANG))
tout.add_column(Column(numpy.array(i),STP.I))
tout.add_column(Column(numpy.array(q),STP.Q))
tout.add_column(Column(numpy.array(u),STP.U))
tout.add_column(Column(numpy.array(v),STP.V))
tout.add_column(Column([options.wave]*len(i),STP.WAVE))
#
tout.write(options.outfile,format='fits',overwrite=True)
if options.verbose:
    print("%d entried saved in file %s" % (len(tout), options.outfile))
else:
    print("%d %s" % (len(tout), options.outfile))    
#


