#!/usr/bin/env python
"""
Create searched area, sky confidence level, and PP plot using information
stored in N areas.dat files, output from run_sky_area.py. If those files are
not given, will look in current_dir/I/areas.dat with I in {1..N}.
"""

from __future__ import print_function
import bz2
import glob
import numpy as np
import matplotlib as mpl
mpl.use('Agg')

import matplotlib.pyplot as plt
from optparse import OptionParser
import plotutils.plotutils as pu
import scipy.stats as ss
import os

parser = OptionParser(
    usage='%prog [options] ev1_areas.dat ev2_areas.dat ... evN_areas.dat',
    description=__doc__)

parser.add_option('--prefix', default='', help='output file prefix')
parser.add_option('--noinj', action='store_true', default=False,
                  help='disable injection-dependent processing')

options, args = parser.parse_args()
cls = np.array([0.5, 0.75, 0.9])
cls_header = ['area({0:d})'.format(int(round(100.0*cl))) for cl in cls]

data = []
dtype = np.dtype([('simulation_id', np.str, 250),
                  ('p_value', np.float),
                  ('searched_area', np.float),
                  ('area50', np.float),
                  ('area75', np.float),
                  ('area90', np.float)])
if args is None or len(args) == 0:
    for file in glob.glob('*/areas.dat'):
        data.append(np.loadtxt(file, dtype=dtype, skiprows=1))
else:
    for file in args:
        data.append(np.loadtxt(file, dtype=dtype, skiprows=1))

new_data = np.zeros(len(data), dtype=data[0].dtype)
for i in range(len(data)):
    new_data[i] = data[i][()]
data = new_data

options.prefix = os.path.realpath(options.prefix)

if not os.path.isdir(options.prefix):
    os.makedirs(options.prefix)

with bz2.BZ2File(os.path.join(options.prefix, 'areas.dat.bz2'), 'w') as out:
    print('simulation_id', 'p_value', 'searched_area',
          *cls_header, sep='\t', file=out)
    for d in data:
        print(*d, sep='\t', file=out)

if not options.noinj:
    ks_stat, ks_p = ss.kstest(data['p_value'], lambda x: x)

    plt.clf()
    pu.plot_cumulative_distribution(data['p_value'], '-k')
    plt.plot(np.linspace(0, 1, 10), np.linspace(0, 1, 10), '--k')
    plt.xlabel(r'$p_\mathrm{inj}$')
    plt.ylabel(r'$P(p_\mathrm{inj})$')
    plt.title('K-S p-value {0:g}'.format(ks_p))
    plt.savefig(os.path.join(options.prefix, 'p-p.pdf'))
    plt.savefig(os.path.join(options.prefix, 'p-p.png'))

    plt.clf()
    pu.plot_cumulative_distribution(data['searched_area'], '-k')
    plt.xscale('log')
    plt.xlabel(r'Searched Area (deg$^2$)')
    plt.savefig(os.path.join(options.prefix, 'searched-area.pdf'))
plt.clf()
pu.plot_cumulative_distribution(data['area50'], label=str('50\%'))
pu.plot_cumulative_distribution(data['area75'], label=str('75\%'))
pu.plot_cumulative_distribution(data['area90'], label=str('90\%'))
plt.xscale('log')
plt.xlabel(r'Credible Area (deg$^2$)')
plt.legend(loc='upper left')
plt.savefig(os.path.join(options.prefix, 'credible-area.pdf'))
