#!/usr/bin/python

# Copyright (C) 2017 Michael Coughlin
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 3 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

""".
Gravitational-wave Electromagnetic Optimization

This script generates an optimized list of pointings and content for
reviewing gravitational-wave skymap likelihoods.

Comments should be e-mailed to michael.coughlin@ligo.org.

"""


import os, sys, glob, optparse, shutil, warnings
import numpy as np
np.random.seed(0)

import healpy as hp
import ephem

import gwemopt.utils, gwemopt.gracedb
import gwemopt.rankedTilesGenerator, gwemopt.waw
import gwemopt.lightcurve, gwemopt.coverage
import gwemopt.efficiency, gwemopt.plotting
import gwemopt.moc, gwemopt.catalog
import gwemopt.tiles, gwemopt.segments
import gwemopt.footprint, gwemopt.transients

if not os.getenv("DISPLAY", None):
    import matplotlib
    matplotlib.use("agg", warn=False)

__author__ = "Michael Coughlin <michael.coughlin@ligo.org>"
__version__ = 1.0
__date__    = "6/17/2017"

# =============================================================================
#
#                               DEFINITIONS
#
# =============================================================================

def parse_commandline():
    """@Parse the options given on the command-line.
    """
    parser = optparse.OptionParser(usage=__doc__,version=__version__)

    parser.add_option("-c", "--configDirectory", help="GW-EM config file directory.", default ="../config/")
    parser.add_option("-s", "--skymap", help="GW skymap.", default='../output/skymaps/G268556.fits')
    parser.add_option("-g", "--gpstime", help="GPS time.", default=1167559936.0, type=float)
    parser.add_option("--do3D",  action="store_true", default=False)

    parser.add_option("-o", "--outputDir", help="output directory",default="../output")
    parser.add_option("-n", "--event", help="event name",default="G268556")
    parser.add_option("--tilingDir", help="tiling directory",default="../tiling")

    parser.add_option("--doEvent",  action="store_true", default=False)
    parser.add_option("--doSkymap",  action="store_true", default=False)
    parser.add_option("--doSamples",  action="store_true", default=False)

    parser.add_option("--doCoverage",  action="store_true", default=False)

    parser.add_option("--doSchedule",  action="store_true", default=False)
    parser.add_option("--scheduleType", help="schedule type",default="greedy")
    parser.add_option("--timeallocationType", help="time allocation type",default="powerlaw")

    parser.add_option("--doPlots",  action="store_true", default=False)
    parser.add_option("--doMovie",  action="store_true", default=False)
    parser.add_option("--doTiles",  action="store_true", default=False)
    parser.add_option("--tilesType", help="tiling type",default="moc")

    parser.add_option("--doCatalog",  action="store_true", default=False)
    parser.add_option("--doObservability",  action="store_true", default=False)
    parser.add_option("--doSkybrightness",  action="store_true", default=False)

    parser.add_option("--doEfficiency",  action="store_true", default=False)
    parser.add_option("-e", "--efficiencyOutput", help="Output file of the efficiency.", 
                      default="efficiency.txt")
    #parser.add_option("-t", "--telescopes", help="Telescope names.",
    #                  default ="PS1")
    #parser.add_option("-d", "--coverageFiles", help="Telescope data files.",
    #                  default ="../data/PS1_GW170104.dat")
    parser.add_option("-t", "--telescopes", help="Telescope names.",
                      default ="ATLAS")
    parser.add_option("-d", "--coverageFiles", help="Telescope coverage files.",
                      default ="../data/ATLAS_GW170104.dat")
    parser.add_option("-l", "--lightcurveFiles", help="Lightcurve files.",
                      default ="../lightcurves/neutron_precursor3.dat,../lightcurves/rpft_m005_v2.dat,../lightcurves/APR4-1215_k1.dat")
    parser.add_option("--Ninj",default=1000,type=int)
    parser.add_option("--Ntiles",default=10,type=int)
    parser.add_option("--Ndet",default=1,type=int)
    parser.add_option("--nside",default=256,type=int)
    parser.add_option("--DScale",default=1.0,type=float)
    parser.add_option("--Tobs",default="0.0,1.0,1.0,2.0,2.0,3.0,3.0,4.0")
    #parser.add_option("--Tobs",default="0.0,1.0")

    parser.add_option("--powerlaw_cl",default=0.9,type=float)
    parser.add_option("--powerlaw_n",default=0.66,type=float)
    parser.add_option("--powerlaw_dist_exp",default=0,type=float)

    parser.add_option("--doFootprint", action="store_true", default=False)
    parser.add_option("--footprint_ra",default=30.0,type=float)
    parser.add_option("--footprint_dec",default=60.0,type=float)
    parser.add_option("--footprint_radius",default=10.0,type=float)

    parser.add_option("--doTransients",  action="store_true", default=False)
    parser.add_option("--transientsFile",default="../transients/ps1_objects.csv")
    parser.add_option("--dt",default=14.0,type=float)

    parser.add_option("-v", "--verbose", action="store_true", default=False,
                      help="Run verbosely. (Default: False)")

    opts, args = parser.parse_args()

    # show parameters
    if opts.verbose:
        print >> sys.stderr, ""
        print >> sys.stderr, "running gwemopt_run..."
        print >> sys.stderr, "version: %s"%__version__
        print >> sys.stderr, ""
        print >> sys.stderr, "***************** PARAMETERS ********************"
        for o in opts.__dict__.items():
          print >> sys.stderr, o[0]+":"
          print >> sys.stderr, o[1]
        print >> sys.stderr, ""

    return opts

def params_struct(opts):
    """@Creates gwemopt params structure
    @param opts
        gwemopt command line options
    """

    params = {}
    params["config"] = {}
    configFiles = glob.glob("%s/*.config"%opts.configDirectory)
    for configFile in configFiles:
        telescope = configFile.split("/")[-1].replace(".config","")
        params["config"][telescope] = gwemopt.utils.readParamsFromFile(configFile)
        if "tesselationFile" in params["config"][telescope]:
            if not os.path.isfile(params["config"][telescope]["tesselationFile"]):
                if params["config"][telescope]["FOV_type"] == "circle": 
                    gwemopt.tiles.tesselation_spiral(params["config"][telescope])
                elif params["config"][telescope]["FOV_type"] == "square":
                    gwemopt.tiles.tesselation_packing(params["config"][telescope])
            params["config"][telescope]["tesselation"] = np.loadtxt(params["config"][telescope]["tesselationFile"],usecols=(0,1,2))

        observer = ephem.Observer()
        observer.lat = str(params["config"][telescope]["latitude"])
        observer.lon = str(params["config"][telescope]["longitude"])
        observer.horizon = str(-12.0)
        observer.elevation = params["config"][telescope]["elevation"]
        params["config"][telescope]["observer"] = observer

    params["skymap"] = opts.skymap
    params["gpstime"] = opts.gpstime
    params["outputDir"] = opts.outputDir
    params["tilingDir"] = opts.tilingDir
    params["event"] = opts.event
    params["coverageFiles"] = opts.coverageFiles.split(",")
    params["telescopes"] = opts.telescopes.split(",")
    params["lightcurveFiles"] = opts.lightcurveFiles.split(",")
    params["tilesType"] = opts.tilesType
    params["scheduleType"] = opts.scheduleType
    params["timeallocationType"] = opts.timeallocationType
    params["Ninj"] = opts.Ninj
    params["Ndet"] = opts.Ndet
    params["Ntiles"] = opts.Ntiles
    params["DScale"] = opts.DScale
    params["nside"] = opts.nside
    params["Tobs"] = np.array(opts.Tobs.split(","),dtype=np.float)
    params["powerlaw_cl"] = opts.powerlaw_cl
    params["powerlaw_n"] = opts.powerlaw_n
    params["powerlaw_dist_exp"] = opts.powerlaw_dist_exp

    params["doPlots"] = opts.doPlots
    params["doMovie"] = opts.doMovie
    params["doObservability"] = opts.doObservability
    params["do3D"] = opts.do3D

    params["doFootprint"] = opts.doFootprint
    params["footprint_ra"] = opts.footprint_ra
    params["footprint_dec"] = opts.footprint_dec
    params["footprint_radius"] = opts.footprint_radius

    params["doTransients"] = opts.doTransients
    params["transientsFile"] = opts.transientsFile
    params["dt"] = opts.dt

    return params

# =============================================================================
#
#                                    MAIN
#
# =============================================================================

warnings.filterwarnings("ignore")

# Parse command line
opts = parse_commandline()
if not os.path.isdir(opts.outputDir): os.makedirs(opts.outputDir)

params = params_struct(opts)

if opts.doEvent:
    params["skymap"], eventinfo = gwemopt.gracedb.get_event(params)
    params["gpstime"] = eventinfo["gpstime"]
elif opts.doSkymap:
    params["skymap"]
    params["gpstime"]
elif opts.doFootprint:
    params["skymap"] = gwemopt.footprint.get_skymap(params)
else:
    print("Need to enable --doEvent, --doFootprint or --doSkymap")
    exit(0)

params = gwemopt.segments.get_telescope_segments(params)

print("Loading skymap...")
# Function to read maps
if opts.do3D:
    map_struct = gwemopt.utils.read_skymap(params, is3D=True)
else:
    map_struct = gwemopt.utils.read_skymap(params, is3D=False)

if opts.doPlots:
    print("Plotting skymap...")
    gwemopt.plotting.skymap(params,map_struct)

if opts.doObservability:
    print("Generating observability")
    observability_struct = gwemopt.utils.observability(params, map_struct)
    map_struct["observability"] = observability_struct
    if opts.doPlots:
        print("Plotting observability...")
        gwemopt.plotting.observability(params,map_struct)

if opts.doSamples:
    print("Generating samples from skymap...")
    if opts.do3D:
        samples_struct = gwemopt.utils.samples_from_skymap(map_struct,is3D=True)
    else:
        samples_struct = gwemopt.utils.samples_from_skymap(map_struct,is3D=False)

if opts.doTiles:
    if params["tilesType"] == "moc":
        print("Generating MOC struct...")
        moc_structs = gwemopt.moc.create_moc(params)
        tile_structs = gwemopt.tiles.moc(params, map_struct, moc_structs)
    elif params["tilesType"] == "ranked":
        print("Generating ranked struct...")
        tile_structs = gwemopt.tiles.rankedTiles(params, map_struct)
    elif params["tilesType"] == "hierarchical":
        print("Generating hierarchical struct...")
        tile_structs = gwemopt.tiles.hierarchical(params, map_struct)
    elif params["tilesType"] == "greedy":
        print("Generating greedy struct...")
        tile_structs = gwemopt.tiles.greedy(params, map_struct)
    else:
        print("Need tilesType to be moc, greedy, hierarchical, or ranked")
        exit(0)

    if opts.doPlots:
        print("Plotting tiles struct...")
        gwemopt.plotting.tiles(params, map_struct, tile_structs)

if opts.doCatalog:
    print("Generating catalog...")
    catalog_struct = gwemopt.catalog.get_catalog(params, map_struct)

if opts.doSchedule:
    if opts.doTiles:
        print("Generating coverage...")
        coverage_struct = gwemopt.coverage.timeallocation(params, map_struct, tile_structs)
    else:
        print("Need to enable --doTiles to use --doSchedule")
        exit(0)

elif opts.doCoverage:
    print("Reading coverage from file...")
    coverage_struct = gwemopt.coverage.read_coverage_files(params)

if opts.doSchedule or opts.doCoverage:
    print("Summary of coverage...")
    gwemopt.scheduler.summary(params,map_struct,coverage_struct)

    if opts.doPlots:
        print("Plotting coverage...")
        gwemopt.plotting.coverage(params, map_struct, coverage_struct)

if opts.doEfficiency:
    if opts.doSchedule or opts.doCoverage:
        print("Computing efficiency...")
        lightcurve_structs = gwemopt.lightcurve.read_files(params["lightcurveFiles"])
        efficiency_structs = {}
        for key in lightcurve_structs.iterkeys():
            lightcurve_struct = lightcurve_structs[key]
            efficiency_struct = gwemopt.efficiency.compute_efficiency(params,map_struct, eventinfo, lightcurve_struct, coverage_struct)
            efficiency_structs[key] = efficiency_struct
            efficiency_structs[key]["legend_label"] = lightcurve_struct["legend_label"]

        if opts.doPlots: 
            print("Plotting efficiency...")
            gwemopt.plotting.efficiency(params, map_struct, efficiency_structs)
    else:
        print("Need to enable --doSchedule or --doCoverage for --doEfficiency")
        exit(0)

if opts.doTransients:
    print("Loading transients list...")
    transients_struct = gwemopt.transients.read_transients(params, map_struct)

    transientsfile = os.path.join(params["outputDir"],'transients.dat')
    fid = open(transientsfile,'w')
    for data,name,classification in zip(transients_struct["data"],transients_struct["name"],transients_struct["classification"]):
        fid.write("%s %.5f %.5f %.5e\n"%(name,data[0],data[1],data[7]))
    fid.close()

    if opts.doPlots:
        print("Plotting transients struct...")
        gwemopt.plotting.transients(params, map_struct, transients_struct)
