#!python
"""
Script to plot solutions
"""
import argparse
from argparse import RawTextHelpFormatter
import logging
from losoto.h5parm import h5parm
from losoto.operations import plot
import matplotlib as mpl
mpl.use("Agg")
import matplotlib.pyplot as plt  # after setting "Agg" to speed up
from matplotlib.colors import ListedColormap
import numpy as np
import os
from rapthor.lib import miscellaneous as misc

log = logging.getLogger('plotrapthor')


def main(h5file, soltype, root=None, refstat=None, soltab=None, dir=None):
    """
    Plot solutions

    Parameters
    ----------
    h5file : str
        Name of solution h5parm file
    soltype : str
        Type of solution to plot: phase or amplitude
    root : str, optional
        Root name for output plots. If None, the soltype is used
    refstat : str, optional
        Name of referance station. If None, the first stations is used
    soltab : str, optional
        Name of soltab to use. If None, the default for the given soltype is used
    dir : str, optional
        Name of direction to use. If None, all directions are used
    """
    with h5parm(h5file) as h:
        ss = h.getSolset('sol000')
        if soltype == 'phase':
            if soltab is None:
                st = ss.getSoltab('phase000')
            else:
                st = ss.getSoltab(soltab)
            ref = st.ant[0]
            ncol = 0
            if 'pol' in st.axesNames:
                color = 'pol'
            else:
                color = ''
            minmax = [-3.2, 3.2]
        elif soltype == 'amplitude':
            if soltab is None:
                st = ss.getSoltab('amplitude000')
            else:
                st = ss.getSoltab(soltab)
            ref = ''
            ncol = 0
            if 'pol' in st.axesNames:
                color = 'pol'
            else:
                color = ''
            minmax = [0, 0]
        else:
            raise ValueError('Solution type "{}" not understood. Must be one of "phase" or '
                             '"amplitude"'.format(soltype))

        if root is None:
            if 'pol' not in st.axesNames:
                root = 'scalar' + soltype + '_'
            else:
                root = soltype + '_'
        if refstat is not None:
            ref = refstat
        log.info('Plotting {} solutions...'.format(soltype))
        if dir is not None:
            st.setSelection(dir=dir)
        if len(st.time) > 1 and len(st.freq) > 1:
            plot.run(st, ['time', 'freq'], axisInTable='ant', axisInCol=color, NColFig=ncol, refAnt=ref,
                     prefix=root, minmax=minmax, plotFlag=True, markerSize=4)
        elif len(st.time) > 1:
            plot.run(st, ['time'], axisInTable='ant', axisInCol=color, NColFig=ncol, refAnt=ref,
                     prefix=root, minmax=minmax, plotFlag=True, markerSize=4)
        elif len(st.freq) > 1:
            plot.run(st, ['freq'], axisInTable='ant', axisInCol=color, NColFig=ncol, refAnt=ref,
                     prefix=root, minmax=minmax, plotFlag=True, markerSize=4)
        else:
            log.warn('Solution table contains only a single time and frequency. No plots made.')


if __name__ == "__main__":
    descriptiontext = "Plot solutions.\n"
    parser = argparse.ArgumentParser(description=descriptiontext, formatter_class=RawTextHelpFormatter)
    parser.add_argument('h5file', help="Name of solution h5parm file")
    parser.add_argument('soltype', help="Type of solution to plot: 'phase' or 'amplitude'")
    parser.add_argument('--root', help="Root name for output plots (default: 'soltype_')", type=str, default=None)
    parser.add_argument('--refstat', help="Name of referance station (default: first)", type=str, default=None)
    parser.add_argument('--soltab', help="Name of solution table (default: default for soltype)", type=str, default=None)
    parser.add_argument('--dir', help="Name of direction (default: all)", type=str, default=None)

    args = parser.parse_args()
    try:
        main(args.h5file, args.soltype, args.root, args.refstat, args.soltab, args.dir)
    except Exception as e:
        log.error('{0}.'.format(e))
