# ==========================================================================
# This script provides a number of functions that are useful for handling
# CTA observations.
#
# Copyright (C) 2011-2014 Juergen Knoedlseder
#
# 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, see <http://www.gnu.org/licenses/>.
#
# ==========================================================================
import ctools
import gammalib
import numpy

# ===================== #
# Simulate observations #
# ===================== #
def sim(obs, log=False, debug=False, edisp=False, seed=0, nbins=0,
        binsz=0.05, npix=200, proj="TAN", coord="GAL"):
    """
    Simulate events for all observations in the container.
    
    Parameters:
     obs   - Observation container
    Keywords:
     log   - Create log file(s)
     debug - Create console dump?
     edisp - Apply energy dispersion?
     seed  - Seed value for simulations (default: 0)
     nbins - Number of energy bins (default: 0=unbinned)
     binsz - Pixel size for binned simulation (deg/pixel)
     npix  - Number of pixels in X and Y for binned simulation
    """
    # Allocate ctobssim application and set parameters
    sim = ctools.ctobssim(obs)
    sim["seed"].integer(seed)
        
    # Optionally open the log file
    if log:
        sim.logFileOpen()

    # Optionally switch-on debugging model
    if debug:
        sim["debug"].boolean(True)
    
    # Run ctobssim application. This will loop over all observations in the
    # container and simulation the events for each observation. Note that
    # events are not added together, they still apply to each observation
    # separately.
    sim.run()
    
    # Binned option?
    if nbins > 0:
    
        # Determine common energy boundaries for observations
        emin = None
        emax = None
        for run in sim.obs():
            run_emin = run.events().ebounds().emin().TeV()
            run_emax = run.events().ebounds().emax().TeV()
            if emin == None:
                emin = run_emin
            elif run_emin > emin:
                emin = run_emin
            if emax == None:
                emax = run_emax
            elif run_emax > emax:
                emax = run_emax
    
        # Allocate ctbin application and set parameters
        bin = ctools.ctbin(sim.obs())
        bin["ebinalg"].string("LOG")
        bin["emin"].real(emin)
        bin["emax"].real(emax)
        bin["enumbins"].integer(nbins)
        bin["usepnt"].boolean(True) # Use pointing for map centre
        bin["nxpix"].integer(npix)
        bin["nypix"].integer(npix)
        bin["binsz"].real(binsz)
        bin["coordsys"].string(coord)
        bin["proj"].string(proj)
        
        # Optionally open the log file
        if log:
            bin.logFileOpen()

        # Optionally switch-on debugging model
        if debug:
            bin["debug"].boolean(True)

        # Run ctbin application. This will loop over all observations in
        # the container and bin the events in counts maps
        bin.run()
        
        # Make a deep copy of the observation that will be returned
        # (the ctbin object will go out of scope one the function is
        # left)
        obs = bin.obs().copy()

    else:
    
        # Make a deep copy of the observation that will be returned
        # (the ctobssim object will go out of scope one the function is
        # left)
        obs = sim.obs().copy()
        
    # Delete the simulation
    del sim
    
    # Return observation container
    return obs


# ================ #
# Fit observations #
# ================ #
def fit(obs, log=False, debug=False, edisp=False):
    """
    Perform maximum likelihood fitting of observations in the container.
    
    Parameters:
     obs   - Observation container
    Keywords:
     log   - Create log file(s)
     debug - Create screen dump
     edisp - Apply energy dispersion?
    """
    # Allocate ctlike application
    like = ctools.ctlike(obs)
    
    # Optionally open the log file
    if log:
        like.logFileOpen()
    
    # Optionally switch-on debugging model
    if debug:
        like["debug"].boolean(True)

    # Optionally apply energy dispersion
    like["edisp"].boolean(edisp)

    # Run ctlike application.
    like.run()
    
    # Return observations
    return like


# ================= #
# Create counts map #
# ================= #
def cntmap(obs, proj="TAN", coord="GAL", xval=0.0, yval=0.0, \
           binsz=0.05, nxpix=200, nypix=200, \
           outname="cntmap.fits"):
    """
    Creates a counts map by combining the events of all observations.
    The counts map will be a summed map over all energies.
    
    Parameters:
     obs     - Observation container
    Keywords:
     proj    - Projection type (e.g. TAN, CAR, STG, ...) (default: TAN)
     coord   - Coordinate type (GAL, CEL) (default: GAL)
     xval    - Reference longitude value [deg] (default: 0.0)
     yval    - Reference latitude value [deg] (default: 0.0)
     binsz   - Pixel size [deg/pixel] (default: 0.05)
     nxpix   - Number of pixels in X direction (default: 200)
     nypix   - Number of pixels in Y direction (default: 200)
     outname - Counts map FITS filename (default: cntmap.fits)
    """
    # Allocate counts map
    map = gammalib.GSkymap(proj, coord, xval, yval, -binsz, binsz, nxpix, nypix, 1)

    # Set maximum pixel number
    maxpixel = nxpix * nypix
    
    # Fill all observations
    for run in obs:
        
        # Loop over all events
        for event in run.events():
            
            # Determine sky pixel
            skydir = gammalib.GCTAInstDir(event.dir()).dir()
            pixel  = map.dir2inx(skydir)
            
            # Set pixel
            if pixel < maxpixel:
                map[pixel] += event.counts()
    
    # Save sky map. The clobber flag is set to True, so any existing FITS
    # file will be overwritten.
    map.save(outname, True)
    
    # Return counts map
    return map


# ================ #
# Create model map #
# ================ #
def modmap(obs, eref=0.1, proj="TAN", coord="GAL", xval=0.0, yval=0.0, \
           binsz=0.05, nxpix=200, nypix=200, \
           outname="modmap.fits"):
    """
    Make model map for a given reference energy by combining all observations.
    The model map will be evaluated for a given reference energy 'eref' and will
    be given in units of [counts/(sr MeV s)].
    
    Parameters:
     obs     - Observation container
    Keywords:
     eref    - Reference energy for which model is created [TeV] (default: 0.1)
     proj    - Projection type (e.g. TAN, CAR, STG, ...) (default: TAN)
     coord   - Coordinate type (GAL, CEL) (default: GAL)
     xval    - Reference longitude value [deg] (default: 0.0)
     yval    - Reference latitude value [deg] (default: 0.0)
     binsz   - Pixel size [deg/pixel] (default: 0.05)
     nxpix   - Number of pixels in X direction (default: 200)
     nypix   - Number of pixels in Y direction (default: 200)
     outname - Model map FITS filename (default: modmap.fits)
    """
    # Allocate model map
    map = gammalib.GSkymap(proj, coord, xval, yval, -binsz, binsz, nxpix, nypix, 1)
    
    # Set reference energy, time and direction. The time is not initialised and is
    # in fact not used (as the IRF is assumed to be time independent for now).
    # The sky direction is set later using the pixel values.
    energy  = gammalib.GEnergy()
    time    = gammalib.GTime()
    instdir = gammalib.GCTAInstDir()
    energy.TeV(eref)
    
    # Loop over all map pixels
    for pixel in range(map.npix()):
        
        # Get sky direction
        skydir = map.inx2dir(pixel)
        instdir.dir(skydir)
        
        # Create event atom for map pixel
        atom = gammalib.GCTAEventAtom()
        atom.dir(instdir)
        atom.energy(energy)
        atom.time(time)
        
        # Initialise model value
        value = 0.0
        
        # Loop over all observations
        for run in obs:
            value += obs.models().eval(atom, run)
        
        # Set map value
        map[pixel] = value
    
    # Save sky map
    map.save(outname, True)
    
    # Return model map
    return map


# ===================================== #
# Set one CTA observation (old version) #
# ===================================== #
def set(pntdir, tstart=0.0, duration=1800.0, deadc=0.95, \
        emin=0.1, emax=100.0, rad=5.0, \
        irf="cta_dummy_irf", caldb="dummy"):
    """
    Obsolete function, use set_obs instead.
    """
    # Print warning
    print("Warning: obsutils.set is obsolete, use obsutils.set_obs instead.")
    
    # Call new function
    obs_cta = set_obs(pntdir, tstart=tstart, duration=duration, deadc=deadc, \
                      emin=emin, emax=emax, rad=rad, \
                      irf=irf, caldb=caldb)
    
    # Return CTA observation
    return obs_cta
    

# ======================= #
# Set one CTA observation #
# ======================= #
def set_obs(pntdir, tstart=0.0, duration=1800.0, deadc=0.95, \
            emin=0.1, emax=100.0, rad=5.0, \
            irf="cta_dummy_irf", caldb="dummy", id="000000"):
    """
    Returns a single CTA observation containing an empty CTA event list.
    By looping over this function you can add CTA observations to the
    observation container.
    
    Parameters:
     pntdir   - Pointing direction [GSkyDir]
    Keywords:
     tstart   - Start time [seconds] (default: 0.0)
     duration - Duration of observation [seconds] (default: 1800.0)
     deadc    - Deadtime correction factor (default: 0.95)
     emin     - Minimum event energy [TeV] (default: 0.1)
     emax     - Maximum event energy [TeV] (default: 100.0)
     rad      - ROI radius used for analysis [deg] (default: 5.0)
     irf      - Instrument response function (default: cta_dummy_irf)
     caldb    - Calibration database path (default: "dummy")
    """
    # Allocate CTA observation
    obs_cta = gammalib.GCTAObservation()

    # Set calibration database
    db = gammalib.GCaldb()
    if (gammalib.dir_exists(caldb)):
        db.rootdir(caldb)
    else:
        db.open("cta", caldb)
    
    # Set pointing direction
    pnt = gammalib.GCTAPointing()
    pnt.dir(pntdir)
    obs_cta.pointing(pnt)
    
    # Set ROI
    roi     = gammalib.GCTARoi()
    instdir = gammalib.GCTAInstDir()
    instdir.dir(pntdir)
    roi.centre(instdir)
    roi.radius(rad)
    
    # Set GTI
    gti = gammalib.GGti()
    gti.append(gammalib.GTime(tstart), gammalib.GTime(tstart+duration))
    
    # Set energy boundaries
    ebounds = gammalib.GEbounds(gammalib.GEnergy(emin, "TeV"), \
                                gammalib.GEnergy(emax, "TeV"))

    # Allocate event list
    events = gammalib.GCTAEventList()
    events.roi(roi)
    events.gti(gti)
    events.ebounds(ebounds)
    obs_cta.events(events)
    
    # Set instrument response
    obs_cta.response(irf, db)
    
    # Set ontime, livetime, and deadtime correction factor
    obs_cta.ontime(duration)
    obs_cta.livetime(duration*deadc)
    obs_cta.deadc(deadc)
    obs_cta.id(id)
    
    # Return CTA observation
    return obs_cta


# ============================ #
# Set list of CTA observations #
# ============================ #
def set_obs_list(obsdeflist, tstart=0.0, duration=1800.0, deadc=0.95, \
        emin=0.1, emax=100.0, rad=5.0, \
        irf="cta_dummy_irf", caldb="dummy"):
    """
    Returns an observation container filled with a list of CTA observations.
    The list is defined by the obsdeflist parameter which is a dictionnary
    containing the mandatory keywords 'ra' and 'dec' that specify the
    pointing direction for a given observation. Optional keyword give control
    over other observation proposerties, such as duration, deadtime correction,
    energy range, etc. If an optional keyword is not specified, the function
    keyword is used instead.
    
    Parameters:
     obsdeflist - Observation definition list [{'ra': x.xx, 'dec': x.xx}]
                  The directory can take the following optional keywords:
                  - duration: Duration of CTA observation [seconds]
                  - deadc: Deadtime correction factor
                  - emin: Minimum event energy [TeV]
                  - emax: Maximum event energy [TeV]
                  - rad: ROI radius used for analysis [deg]
                  - irf: Instrument response function
                  - caldb: Calibration database path
                  Optional keywords overwrite keywords specified in the
                  function call.
    Keywords:
     tstart     - Start time [seconds] (default: 0.0)
     duration   - Duration of one CTA observation [seconds] (default: 1800.0)
     deadc      - Deadtime correction factor (default: 0.95)
     emin       - Minimum event energy [TeV] (default: 0.1)
     emax       - Maximum event energy [TeV] (default: 100.0)
     rad        - ROI radius used for analysis [deg] (default: 5.0)
     irf        - Instrument response function (default: cta_dummy_irf)
     caldb      - Calibration database path (default: "dummy")
    """
    # Initialise empty observation container
    obs = gammalib.GObservations()

    # Initialise first time and identifier
    obs_start = tstart
    obs_id    = 1

    # Loop over observation definition list
    for obsdef in obsdeflist:

        # Set pointing direction
        pntdir = gammalib.GSkyDir()
        pntdir.radec_deg(obsdef['ra'], obsdef['dec'])
        
        # Set duration (use default if not found in definition list)
        if 'duration' in obsdef:
            obs_duration = obsdef['duration']
        else:
            obs_duration = duration

        # Set emin (use default if not found in definition list)
        if 'emin' in obsdef:
            obs_emin = obsdef['emin']
        else:
            obs_emin = emin

        # Set emax (use default if not found in definition list)
        if 'emax' in obsdef:
            obs_emax = obsdef['emax']
        else:
            obs_emax = emax

        # Set radius (use default if not found in definition list)
        if 'rad' in obsdef:
            obs_rad = obsdef['rad']
        else:
            obs_rad = rad

        # Set deadc (use default if not found in definition list)
        if 'deadc' in obsdef:
            obs_deadc = obsdef['deadc']
        else:
            obs_deadc = deadc

        # Set caldb (use default if not found in definition list)
        if 'caldb' in obsdef:
            obs_caldb = obsdef['caldb']
        else:
            obs_caldb = caldb

        # Set irf (use default if not found in definition list)
        if 'irf' in obsdef:
            obs_irf = obsdef['irf']
        else:
            obs_irf = irf
        
        # Generate identifier string
        id = "%6.6d" % obs_id

        # Set CTA observation
        obs_cta = set_obs(pntdir, tstart=obs_start, duration=obs_duration, \
                          deadc=obs_deadc, emin=obs_emin, emax=obs_emax, \
                          rad=obs_rad, irf=obs_irf, caldb=obs_caldb, id=id)
        
        # Append to container
        obs.append(obs_cta)

        # Update start time and identifier
        obs_start += obs_duration
        obs_id    += 1
    
    # Return observation container
    return obs


# ============================ #
# Set CTA observation patterns #
# ============================ #
def set_obs_patterns(pattern, ra=83.6331, dec=22.0145, offset=1.5):
    """
    Sets a number of standard patterns.

    Parameters:
     pattern - Observation pattern. Possible options are:
               - "single": single pointing
               - "four": four pointings 'offset' around pattern centre
             
    Keywords:
     ra      - Right Ascension of pattern centre [deg] (default: 83.6331)
     dec     - Declination of pattern centre [deg] (default: 22.0145)
     offset  - Offset from pattern centre [deg] (default: 1.5)
    """
    # Initialise observation definition list
    obsdeflist = []
    
    # Add patterns
    if pattern == "single":
        obsdef = {'ra': ra, 'dec': dec}
        obsdeflist.append(obsdef)
    elif pattern == "four":
        # Set pattern centre
        centre = gammalib.GSkyDir()
        centre.radec_deg(ra, dec)
        
        # Append pointings
        for phi in [0.0, 90.0, 180.0, 270.0]:
            pntdir = centre.copy()
            pntdir.rotate_deg(phi, offset)
            obsdeflist.append({'ra': pntdir.ra_deg(), 'dec': pntdir.dec_deg()})
    else:
        print("Warning: Observation pattern '"+str(pattern)+"' not recognized.")
    
    # Return observation definition list
    return obsdeflist

# =================================================#
# Make spectrum for one Source in aCTA Observation #
# ================================================ #
def getlogbins(nbins=12,emin=0.1,emax=100):
    """
    Returns list with logarithmiscenergy bins in
    format [[emin1,emax1],[emin2,emax2],..]
    
    Parameters:
       nbins   Number of bins
       emin    Minimum energy
       emax    Maximum energy
    """
    lgemin=numpy.log10(emin)
    lgemax=numpy.log10(emax)
    lgbins=[]
    lgedges=numpy.arange(lgemin,lgemax,(lgemax-lgemin)/nbins)
    lgedges=numpy.append(lgedges,[lgemax])
    for ii in range(len(lgedges)-1):
        lgbins.append([lgedges[ii],lgedges[ii+1]])
    ebins=numpy.power(10,lgbins)
    return ebins
def specpoints(obs,source,ebins):
    """
    Returns spectral points in the given binning for
    the passed observations and source.
    
    Parameters:
     obs    - Pointing direction
     source - Source name for which spectrum should be done
     ebins  - Binning passed as a list in units of TeV as
              [[emin1,emax1],[emin2,emax2],..]
    Keywords:
    """
    #Fix all parameters except prefactor
    for model in obs.models():
        for par in range(model.size()):
            model[par].fix()
    
    opt_spectral = obs.models()[source].spectral()
    opt_spectral["Prefactor"].free()
    
    #Define spectrum dictionary, units are in format readable by astropy
    pars = {"value":[],"eu_value":[],"ed_value":[],"unit":[]}
    spec = {"energy":pars.copy(),"flux":pars.copy(),"TS":[]}
    
    #Loop over energy bins, fit and write result to spectrum
    for ebin in ebins:
        #Get logarithmic bin center and flux value of all energy fit 
        ecenter = numpy.power(10,(numpy.log10(ebin[1])+numpy.log10(ebin[0]))/2.)
        allflux = opt_spectral.eval(gammalib.GEnergy(ecenter,"TeV"),gammalib.GTime(0))

        #Clone observations and reset energy thresholds
        select = ctools.ctselect(obs)
        select["ra"].real(-1.0)
        select["dec"].real(-1.0)
        select["rad"].real(-1.0)
        select["tmin"].real(0)
        select["tmax"].real(0)
        select["emin"].real(ebin[0])
        select["emax"].real(ebin[1])
        select.run()
        obsclone = select.obs()
        
        #Refit and write out to spectrum
        binlike=fit(obsclone,log=False, debug=False, edisp=False)
        binopt_spectral = binlike.obs().models()[source].spectral()
        binflux=binopt_spectral.eval(gammalib.GEnergy(ecenter,"TeV"),gammalib.GTime(0))
        if binopt_spectral["Prefactor"].value()>0:
            relerr=binopt_spectral["Prefactor"].error()/binopt_spectral["Prefactor"].value()
        else:
            relerr=0

        #Store and convert to TeV from MeV
        spec["energy"]["value"].append(ecenter)
        spec["energy"]["ed_value"].append(ecenter-ebin[0])
        spec["energy"]["eu_value"].append(ebin[1]-ecenter)
        spec["energy"]["unit"].append("1e12 eV")
        
        spec["flux"]["value"].append(binflux*1e6)
        spec["flux"]["ed_value"].append(relerr*binflux*1e6)
        spec["flux"]["eu_value"].append(relerr*binflux*1e6)
        spec["flux"]["unit"].append("cm-2 1e-12 eV-1")
        
        spec["TS"].append(binlike.obs().models()[source].ts())
    
    #Return spectrum
    return spec