#!/usr/bin/env python
""" Plot variation in PSD
"""
import matplotlib
matplotlib.use('Agg')
import h5py, numpy, argparse, pylab, pycbc.results, sys
import pycbc.psd
import pycbc.version


parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--version", action="version",
                    version=pycbc.version.git_verbose_msg)
parser.add_argument("--psd-files", nargs='+', required=True,
                    help='HDF file(s) containing the PSDs to plot')
parser.add_argument("--output-file", required=True, help='Output file name')
pycbc.psd.insert_psd_option_group(parser, output=False)
args = parser.parse_args()

fig = pylab.figure(0)
ax = fig.gca()
ax.grid(which='both', ls='solid', alpha=0.2, lw=0.3)
ax.set_ylabel('Amplitude Spectral Density (Strain / $\\sqrt{\\rm Hz}$)')
ax.set_xlabel('Frequency (Hz)')

y_min = None

for psd_file in args.psd_files:
    f = h5py.File(psd_file, 'r')
    ifo = f.keys()[0]
    df = f[ifo + '/psds/0'].attrs['delta_f']
    keys = f[ifo + '/psds'].keys()
    psds = [f[ifo + '/psds/' + key][:] for key in keys]

    flow = f.attrs['low_frequency_cutoff']
    kmin = int(flow / df)

    fac = 1.0 / pycbc.DYN_RANGE_FAC
    high = numpy.percentile(psds, 95, axis=0)[kmin:] ** 0.5 * fac
    low = numpy.percentile(psds, 5, axis=0)[kmin:] ** 0.5 * fac
    middle = numpy.percentile(psds, 50, axis=0)[kmin:] ** 0.5 * fac
    samples = numpy.arange(0, len(psds[0]))[kmin:] * df

    if y_min is None or y_min > low.min():
        y_min = low.min()

    color = pycbc.results.ifo_color(ifo)

    ax.fill_between(samples, low, high, alpha=0.4, linewidth=0, color=color)
    ax.loglog(samples, middle, linewidth=0.3, color=color, label=ifo)
    ax.set_xlim(flow, samples[-1])

reference_psd = pycbc.psd.from_cli(args, 2048, 1., 10., None)
if reference_psd is not None:
    ax.loglog(reference_psd.sample_frequencies, reference_psd ** 0.5,
              '-k', lw=0.3, label='Reference')

ax.set_ylim(y_min * 0.5, y_min * 100)

ax.legend(loc='upper right', fontsize='small')
pycbc.results.save_fig_with_metadata(fig, args.output_file, 
    title="Noise Amplitude Spectral Density",
    caption="Median amplitude spectral density plotted with a shaded region " 
              "between the 5th and 95th perentiles. ",
    cmd=' '.join(sys.argv),
    fig_kwds={'dpi': 200})
