#!/usr/bin/env
"""
Plots multiple lag-energy spectra (from different observations, simulations,
etc.) on the same plot.

Change lag_dir and sim_dir to suit your machine and setup.

Example call: python overplot_lag-energy.py cygx1
"""
import argparse
import subprocess
import numpy as np
from astropy.io import fits
from astropy.table import Table
import os.path
import matplotlib.pyplot as plt
import matplotlib.font_manager as font_manager
from matplotlib.ticker import MultipleLocator
from matplotlib.ticker import ScalarFormatter
from datetime import datetime
import matplotlib.colors as colors
import itertools

__author__ = 'Abigail Stevens <A.L.Stevens at uva.nl>'
__year__ = '2015-2016'

HOME_DIR = os.path.expanduser("~")
DAY = datetime.now().strftime("%y%m%d")

PLOT_EXT = "eps"
LO_FREQ = 4.0
UP_FREQ = 7.0
LO_ENERGY = 3.0
UP_ENERGY = 20.0

DETCHANS = 64

################################################################################
def pairwise(iterable):
    """
    s -> (s0,s1), (s1,s2), (s2, s3), ...
    From https://docs.python.org/2/library/itertools.html#recipes
    Used when reading lines in the file so I can peek at the next line.

    Parameters
    ----------
    an iterable, like a list or an open file

    Returns
    -------
    The next two items in an iterable, like in the example a few lines above.

    """
    a, b = itertools.tee(iterable)
    next(b, None)
    return itertools.izip(a, b)

################################################################################
def make_plot(in_file_list, labels, plot_root, prefix="GX339-BQPO"):
    """
    Plots multiple lag plots in one figure.

    Parameters
    ----------
    in_file_list : list of str
        1-D list of the lag files to plot.

    labels : list of str
        1-D list of the line labels for each lag. Must be same length as
        in_file_list.

    plot_root : str
        Dir + base name of the plot file. "_lag-energy.[PLOT_EXT]" are appended
        to it.

    Returns
    -------
    nothing

    """

    colours = [colors.cnames['darkblue'],
               colors.cnames['darkviolet'],
               colors.cnames['coral'],
               "black"]

    # colours = [colors.cnames['darkred'],
    #            colors.cnames['darkcyan'],
    #            colors.cnames['deeppink'],
    #            "black"]

    markers = ['^',
               'x',
               's',
               'o']

    #######################
    ## Setting up the plot
    #######################

    font_prop = font_manager.FontProperties(size=20)
    energies = np.loadtxt(HOME_DIR + "/Reduced_data/" + prefix +
                          "/energies.txt")
    energy_list = [np.mean([x, y]) for x,y in pairwise(energies)]
    energy_err = [np.abs(a-b) for (a,b) in zip(energy_list, energies[0:-1])]

    plot_file = plot_root + "_lag-energy." + PLOT_EXT
    print "Lag-energy spectrum: %s" % plot_file

    e_chans = np.arange(0, DETCHANS)

    ###################
    ## Making the plot
    ###################

    fig, ax = plt.subplots(1, 1, figsize=(10,7.5), dpi=300, tight_layout=True)

    ax.hlines(0.0, 3, 21, linestyle='dashed', lw=2, color='black')

    ## Deleting the values at energy channel 10 for RXTE PCA event-mode data
    if DETCHANS == 64:
        e_chans = np.delete(e_chans, 10)
        energy_list = np.delete(energy_list, 10)
        energy_err = np.delete(energy_err, 10)
    i = 0

    for in_file in in_file_list:

        # try:
        #     fits_hdu = fits.open(in_file)
        # except IOError:
        #     print "\tERROR: File does not exist: %s" % in_file
        #     exit()
        # tlag = fits_hdu[2].data.field('TIME_LAG')
        # tlag_err = fits_hdu[2].data.field('TIME_LAG_ERR')

        try:
            lag_table = Table.read(in_file, format='fits', hdu=2)  # HDU 2 for energy lags
        except IOError:
            print "\tERROR: File does not exist: %s" % in_file
            exit()
        tlag = lag_table['TIME_LAG']
        tlag_err = lag_table['TIME_ERR']

        ## Deleting the values at energy channel 10 for RXTE PCA event-mode data
        if DETCHANS == 64:
            tlag = np.delete(tlag, 10)
            tlag_err = np.delete(tlag_err, 10)

            if labels[i].lower() == "data":
                ax.errorbar(energy_list[2:26], tlag[2:26],
                            xerr=energy_err[2:26], yerr=tlag_err[2:26],
                            ls='none', marker='o', ms=10, mew=2, mec='black',
                            mfc='black', ecolor='black', elinewidth=3,
                            capsize=0, label=labels[i])

            else:
                if markers[i] == 'x':
                    ax.errorbar(energy_list[2:26], tlag[2:26],
                                xerr=energy_err[2:26], yerr=tlag_err[2:26],
                                lw=3, drawstyle='steps-mid', ls='none', ms=11,
                                marker=markers[i], mec=colours[i], mew=2,
                                color=colours[i], fillstyle='none',
                                ecolor=colours[i], elinewidth=3, capsize=0,
                                label=labels[i])
                else:
                    ax.errorbar(energy_list[2:26], tlag[2:26],
                                xerr=energy_err[2:26], yerr=tlag_err[2:26],
                                lw=3, drawstyle='steps-mid', ls='none', ms=8,