# -*- coding: utf-8 -*-
"""
Created on Tue Mar 29 22:28:32 2016

For the specific sets of sources, Calculate the linear drift of both coordinate,
global spin and simulated global spin.

@author: Neo
"""

import numpy as np
from LinearFitting import APM
import time
from fun import ADepoA
from RotationFit import MeanWeiRotation, RotationFit
from GASimulation import simulation
cos = np.cos
sin = np.sin

cat_fil = ['icrf1', 'icrf2', 'MFV247', 'AMS260']
res_dir = '../results/'
dat_dir = '../data/opa/'
rot = 'SpecialSets.rot'
gli = 'SpecialSets.sim'
frot = open(res_dir + rot,'w')
fgli = open(res_dir + gli,'w')
## headers of output.
print(time.strftime('## This list was generated at %Y-%m-%d %H:%M:%S',\
        time.localtime(time.time())), file=frot)
print('''############################################################################
## Global Spin Vector of selected source subsets.
##data format
##Subsets   w_tol   w_x   w_y   w_z
##            (Expectation | Std)
##unit: uas/yr
############################################################################''', file=frot)
print(time.strftime('## This list was generated at %Y-%m-%d %H:%M:%S',\
        time.localtime(time.time())), file=fgli)
print('''############################################################################
## Simulated global Spin Vector of selected source subsets.
##data format
##Subsets   w_tol   w_x   w_y   w_z
##            (Expectation | Std)
##unit: uas/yr
############################################################################''', file=fgli)

for i in range(len(cat_fil)):
    cat = cat_fil[i]
###############################################################################
## Linear drift fitting.
    apm_fil = cat  + '.apm'
    soun = np.loadtxt('../list/' + cat + '.cat', dtype=str)
    
    ## output file. 
    fapm = open(res_dir + apm_fil, 'w')
    print(time.strftime('## This list was generated at %Y-%m-%d %H:%M:%S',\
        time.localtime(time.time())), file=fapm)
    print('## source list: %s'%cat, file=fapm)
    print('''##data format
##source   P.M.(uas/yr)  Err. (uas/yr)   R.A.(o)   Dec.(o)  Err. (mas)
##            (Total R.A.* Dec.)
##R.A. and Dec. corrsponding to JD2000.0
############################################################################''', file=fapm)
    
    for i in range(len(soun)):
        sou_name = soun[i]
    ##data format
    ##epo(Julian day)   R.A.(o)   Dec.(o)  Err. (mas)  Corr. Del.  IERS   IVS   
        epo, ra, dec, era, edec = \
        np.loadtxt(dat_dir+sou_name+'.dat', usecols=list(range(5)), unpack=True)
        ## Unit: Julian day, deg, deg, mas, mas
        
        if epo.size>1:
            epo = ADepoA(epo)                       #yr
            ra  = ra*3.6e6                          #mas
            dec = dec*3.6e6                         #mas 
    
            [pmRA, pmDE, RA0, DE0, epmRA, epmDE, eRA0, eDE0] = \
            APM(epo, ra, era, dec, edec)
            
            era0, ede0 = eRA0, eDE0
    ##  mas -> deg    
            ra0, de0 = RA0/3.6e6, DE0/3.6e6
    ##  mas/yr -> uas/yr
            mu_ra, mu_de = pmRA*1.0e3*cos(np.deg2rad(de0)), pmDE*1.0e3
            emu_ra, emu_de = epmRA*1.0e3*cos(np.deg2rad(de0)), epmDE*1.0e3
        
        else:
            mu_ra, mu_de = 0.0, 0.0
            emu_ra, emu_de = 1.0e6, 1.0e6
            ra0, era0 = ra,  era 
            de0, ede0 = dec, edec
            
        mu_tol  = np.sqrt( mu_ra**2 +  mu_de**2)
        emu_tol = np.sqrt(emu_ra**2 + emu_de**2)
        
#        earray = np.array([emu_ra, emu_de, emu_tol])
#        emu_ra, emu_de, emu_tol = np.where(earray>1.0e6, 1.0e6, earray)
    # Write the results into outpu file
    ##Unit: microarcsecond/year for proper motion, and degree for angles 
    #data format
    #source   P.M.(uas/yr)  Err. (uas/yr)   R.A.(o)   Dec.(o)  Err. (mas)
    #            (Total R.A. Dec.)
    #R.A. and Dec. corrspond to JD2000.0
        print(sou_name + ("  %12.4f"*6+"  %14.10f"*2+"  %12.4f"*2) \
        %(mu_tol,mu_ra, mu_de, emu_tol, emu_ra, emu_de,\
        ra0, de0, era0, ede0), file=fapm)
        
    fapm.close()
###############################################################################
## Global spin of fitting and simulated proper motion fields.
    V, V_RA, V_Dec, V_err, Va_E, Vd_E, RA,  Dec = \
        np.loadtxt(res_dir + apm_fil, usecols=list(range(1,9)), unpack=True) 
###  apm >= 1mas/yr are considered unreliable.
#    ind = np.where(V<1000)[0]
#    RA = np.take(RA, ind)
#    Dec= np.take(Dec, ind)
#    V_RA = np.take(V_RA, ind)
#    V_Dec= np.take(V_Dec, ind)    
## Unit: deg -> rad
    RA = np.deg2rad(RA)
    Dec = np.deg2rad(Dec)    
## Global spin of fitting proper motion fields.   
    w, sig, corrcoef = MeanWeiRotation(V_RA, V_Dec, RA, Dec)
#    w, sig, corrcoef = RotationFit(V_RA, V_Dec, Va_E, V_Dec, RA, Dec)
    wx, wy, wz = w
    sigwx, sigwy, sigwz = sig
    wt = np.sqrt(sum(w**2))
    sigw = np.sqrt(sum(sig**2))
    print("%6s"%cat + "    %8.4f"*8%(wt, sigw, wx, sigwx, wy, sigwy, wz, sigwz), file=frot)
## Global spin of simulated proper motion fields.    
    ua, ud = simulation(RA, Dec)
    w, sig, corrcoef = MeanWeiRotation(ua, ud, RA, Dec)
    wx, wy, wz = w
    sigwx, sigwy, sigwz = sig
    wt = np.sqrt(sum(w**2))
    sigw = np.sqrt(sum(sig**2))
    print("%6s"%cat + "    %8.4f"*8%(wt, sigw, wx, sigwx, wy, sigwy, wz, sigwz), file=fgli)
      
frot.close()  
fgli.close()  

print("Done!") 