#!python
# -*- coding: utf-8 -*-

__author__ = 'Alan Loh'
__copyright__ = 'Copyright 2019, nenupytf'
__credits__ = ['Alan Loh']
__maintainer__ = 'Alan Loh'
__email__ = 'alan.loh@obspm.fr'
__status__ = 'WIP'


import argparse
import matplotlib.pyplot as plt
import numpy as np

from nenupytf.read import ObsRepo, Lane


# ============================================================= #
# ============================================================= #
if __name__ == '__main__':
    # --------------------------------------------------------- #
    # -------------------- Parse Arguments -------------------- #
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        '-o',
        '--obs',
        type=str,
        default='/.',
        help='Observation directory',
        required=False
        )
    parser.add_argument(
        '-l',
        '--lane',
        type=int,
        default=0,
        help='Lane selected',
        required=False
        )
    parser.add_argument(
        '-b',
        '--beam',
        type=int,
        default=0,
        help='Beam selected',
        required=False
        )
    parser.add_argument(
        '-s',
        '--stokes',
        type=str,
        default='I',
        help='Stokes parameter (I, Q, U, V)',
        required=False
        )
    parser.add_argument(
        '-f',
        '--freq',
        type=float,
        default=None,
        help='Frequency range in MHz (put two space-separated values)',
        required=False,
        nargs='+'
        )
    parser.add_argument(
        '-t',
        '--time',
        type=str,
        default=None,
        help='Time range in ISOT (put two space-separated values)',
        required=False,
        nargs='+'
        )

    args = parser.parse_args()


    # --------------------------------------------------------- #
    # ------------------------- Plot -------------------------- #
    obs = ObsRepo(repo=args.obs)
    if not args.lane in obs.lanes:
        raise ValueError(
            'No such lane.'
            )
    else:
        lanefile = obs.files[tuple([obs.lanes == args.lane])][0]

    l = Lane(spectrum=lanefile)
    t, f, d = l.average(
        time=args.time,
        freq=args.freq,
        stokes=args.stokes,
        beam=args.beam
        )

    pcm = plt.pcolormesh(
        (t-t[0]).sec,
        f,
        d.T,
        vmin=np.percentile(d, 5),
        vmax=np.percentile(d, 95),
        cmap='YlGnBu_r')
    cbar = plt.colorbar(pcm)
    cbar.set_label('Stokes {}'.format(args.stokes))
    plt.ylabel('Frequency (MHz)')
    plt.xlabel('Time (sec since {})'.format(t[0].isot))
    plt.show()
# ============================================================= #

