#!/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 astropy, astroplan

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.tiles, gwemopt.segments_astroplan

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("--doDatabase",  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("--catalogDir", help="catalog directory",default="../catalogs")
    parser.add_option("--galaxy_catalog", help="Source catalog",default="GLADE")
 
    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=1.0,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("--doSingleExposure",  action="store_true", default=False)
    parser.add_option("--filters",default="r,g,r")
    parser.add_option("--exposuretimes",default="30.0,30.0,30.0")
    parser.add_option("--doReferences",  action="store_true", default=False)

    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),comments='%')

        location = astropy.coordinates.EarthLocation(params["config"][telescope]["longitude"],params["config"][telescope]["latitude"],params["config"][telescope]["elevation"])
        observer = astroplan.Observer(location=location)
        params["config"][telescope]["observer"] = observer

        # 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["catalogDir"] = opts.catalogDir
    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["doDatabase"] = opts.doDatabase

    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

    params["galaxy_catalog"] = opts.galaxy_catalog

    params["doSingleExposure"] = opts.doSingleExposure
    params["filters"] = opts.filters.split(",")
    params["exposuretimes"] = np.array(opts.exposuretimes.split(","),dtype=np.float)
    params["mindiff"] = 0.0
    params["doReferences"] = opts.doReferences

    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_astroplan.get_telescope_segments(params)
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.doCatalog:
    print("Generating catalog...")
    map_struct = gwemopt.catalog.get_catalog(params, map_struct)

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.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, 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)
