#! /usr/bin/python3.6
# -*- coding: utf-8 -*-

"""
    nenuskybeam -s /Users/aloh/Desktop/TEST_SKYBEAM.png -t '2019-02-25 15:34:00' -a 180 -e 75 -f 60 -m 0
"""

import argparse
import healpy as hp
import matplotlib as mpl
import pylab as plt
import matplotlib.ticker as mtick
import numpy as np

from nenupy.hpx import Digibeam, Skymodel
from nenupy.astro import getAltaz, getSrc, getTime

__author__ = 'Alan Loh'
__copyright__ = 'Copyright 2019, nenupy'
__credits__ = ['Alan Loh']
__license__ = 'MIT'
__version__ = '0.0.1'
__maintainer__ = 'Alan Loh'
__email__ = 'alan.loh@obspm.fr'
__status__ = 'WIP'

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-t', '--time', type=str, help="UTC Time (e.g. 2016-09-30 14:00:00)", required=True)
    parser.add_argument('-s', '--store', type=str, help="Filename (needs to end by .fits/png)", required=True)
    parser.add_argument('-f', '--freq', type=float, default=50, help="Frequency in MHz", required=False)
    parser.add_argument('-a', '--azimuth', type=float, default=180, help="Azimuth in degrees", required=False)
    parser.add_argument('-e', '--elevation', type=float, default=90, help="Elevation in degrees", required=False)
    parser.add_argument('-p', '--polar', type=str, default='NW', help="Polarization (NW or NE)", required=False)
    parser.add_argument('-m', '--ma', type=int, default=[-1], help="Mini-Array name indices (e.g. -m -1: all mini-arrays; -m 10 20)", required=False, nargs='+')
    args = parser.parse_args()

    if args.ma == [-1]:
        args.ma = None

    db = Digibeam(azana=args.azimuth,
                azdig=args.azimuth,
                elana=args.elevation,
                eldig=args.elevation,
                miniarrays=args.ma,
                freq=args.freq,
                polar=args.polar,
                resol=0.2)

    beam = db.get_digibeam()

    sm = Skymodel(nside=db._nside,
                freq=args.freq)

    skybeam = beam * sm.get_skymodel(time=args.time, model='GSM')

    skybeam = hp.cartview(skybeam, return_projected_map=True)
    plt.close('all')
    skybeam = skybeam[int(skybeam.shape[0]/2):, :] 

    theta = np.linspace(0., 90., skybeam.shape[0])
    phi   = np.radians( np.linspace(-180., 180., skybeam.shape[1]) )

    fig = plt.figure(figsize=(15/2.54, 15/2.54))
    ax  = fig.add_subplot(111, projection='polar')
    normcb = mpl.colors.LogNorm(vmin=skybeam.max() * 1.e-5, vmax=skybeam.max())
    p = ax.pcolormesh(phi, theta, np.flipud(skybeam), norm=normcb, rasterized=True, cmap='bone')
    ax.grid(linestyle='-', linewidth=0.5, color='white', alpha=0.4)
    plt.setp(ax.get_yticklabels(), rotation='horizontal', color='white')

    g = lambda x,y: r'%d'%(90-x)
    ax.yaxis.set_major_formatter(mtick.FuncFormatter( g ))

    # Sources
    ateam = ['Vir A', 'Cyg A', 'Cas A', 'Tau A', 'Her A', 'Hyd A', 'Sun', 'Moon', 'Jupiter', 'Saturn']
    for at in ateam:
        try:
            altaz = getAltaz(source=at, time=args.time, loc='NenuFAR')
            az, el = altaz.az.deg, altaz.alt.deg
            if el > 0:
                ax.scatter(np.radians(az), 90-el, s=150, facecolor='#d62728', edgecolor='#d62728', alpha=0.3)
                ax.text(np.radians(az), 90-el, '   '+at, color='#d62728')
        except:
            pass

    ax.set_theta_zero_location('N')
    ax.set_theta_direction(-1)
    plt.title('pol={}, freq={:.2f}MHz, az={}, el={}'.format(args.polar, args.freq, args.azimuth, args.elevation))
    ax.set_ylim(0, 90)

    if args.store.lower().endswith('png'):
        plt.savefig(args.store, facecolor='none',
            edgecolor='none', transparent=True, bbox_inches='tight')
    else:
        plt.show()
    plt.close('all')