#! /usr/bin/env python
# ==========================================================================
# This script generates the TS distribution for a particular model based
# on Monte-Carlo simulations.
#
# 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/>.
#
# ==========================================================================
from ctools import *
from gammalib import *
import sys
import csv
import math


# ============== #
# cstsdist class #
# ============== #
class cstsdist(GApplication):
    """
    This class implements the TS distribution generation script. It derives
    from the GammaLib::GApplication class which provides support for parameter
    files, command line arguments, and logging. In that way the Python
    script behaves just as a regular ctool.
    """
    def __init__(self, *argv):
        """
        Constructor.
        """
        # Set name
        self.name    = "cstsdist"
        self.version = "0.1.0"
        
        # Initialise some members
        self.obs        = None
        self.bkg_model  = None
        self.full_model = None
        
        # Make sure that parfile exists
        file = self.parfile()

        # Initialise application
        if len(argv) == 0:
            GApplication.__init__(self, self.name, self.version)
        elif len(argv) ==1:
            GApplication.__init__(self, self.name, self.version, *argv)
        else:
            raise TypeError("Invalid number of arguments given.")

        # Set logger properties
        self.log_header()
        self.log.date(True)

        # Return
        return
    
    def __del__(self):
        """
        Destructor.
        """
        #  Write separator into logger
        if self.logTerse():
            self.log("\n")
        
        # Return
        return

    def parfile(self):
        """
        Check if parfile exists. If parfile does not exist then create a
        default parfile. This kluge avoids shipping the cscript with a parfile.
        """
        # Set parfile name
        parfile = self.name+".par"
        
        try:
            pars = GApplicationPars(parfile)
        except:
            # Signal if parfile was not found
            sys.stdout.write("Parfile "+parfile+" not found. Create default parfile.\n")
            
            # Create default parfile
            pars = GApplicationPars()
            pars.append(GApplicationPar("outfile","f","h","ts.dat","","","Output file name"))
            pars.append(GApplicationPar("ntrials","i","a","10","","","Number of trials"))
            pars.append(GApplicationPar("caldb","s","a","$GAMMALIB/share/caldb/cta","","","Calibration database"))
            pars.append(GApplicationPar("irf","s","a","cta_dummy_irf","","","Instrument response function"))
            pars.append(GApplicationPar("type","s","a","point","","","Source model type (point/gauss/shell/disk)"))
            pars.append(GApplicationPar("index","r","h","-2.48","","","Spectral index"))
            pars.append(GApplicationPar("offset","r","a","0.0","0.0","","Source offset angle (deg)"))
            pars.append(GApplicationPar("bkg","s","a","$GAMMALIB/share/models/bkg_dummy.txt","","","Background model file function (none=power law for E)"))
            pars.append(GApplicationPar("emin","r","a","0.1","0.0","","Lower energy limit (TeV)"))
            pars.append(GApplicationPar("emax","r","a","100.0","0.0","","Upper energy limit (TeV)"))
            pars.append(GApplicationPar("enumbins","i","a","0","","","Number of energy bins (0=unbinned)"))
            pars.append(GApplicationPar("duration","r","a","180000.0","","","Effective exposure time (s)"))
            pars.append(GApplicationPar("rad","r","h","5.0","","","Radius of ROI (deg)"))
            pars.append(GApplicationPar("npix","i","h","200","","","Number of pixels for binned"))
            pars.append(GApplicationPar("binsz","r","h","0.05","","","Pixel size for binned (deg/pixel)"))
            pars.append_standard()
            pars.save(parfile)
        
        # Return
        return
        
    def get_parameters(self):
        """
        Get parameters from parfile and setup the observation.
        """
        # Get parameters
        self.m_outfile  = self["outfile"].filename()
        self.m_ntrials  = self["ntrials"].integer()
        self.m_caldb    = self["caldb"].string()
        self.m_irf      = self["irf"].string()
        self.m_type     = self["type"].string()
        self.m_index    = self["index"].real()
        self.m_offset   = self["offset"].real()
        self.m_bkg      = self["bkg"].string()
        self.m_emin     = self["emin"].real()
        self.m_emax     = self["emax"].real()
        self.m_enumbins = self["enumbins"].integer()
        self.m_duration = self["duration"].real()
        self.m_rad      = self["rad"].real()
        self.m_npix     = self["npix"].integer()
        self.m_binsz    = self["binsz"].real()

        # Set some fixed parameters
        self.m_log   = False # Logging in client tools
        self.m_debug = False # Debugging in client tools
        
        # Setup observation
        self.obs = self.set_obs(emin=self.m_emin, emax=self.m_emax)
        
        # Initialise models. Note that we centre the point source at the Galactic
        # center as our observation is also centred at the Galactic centre, so
        # we're onaxis.
        self.bkg_model  = GModels()
        self.full_model = GModels()
        self.bkg_model.append(self.set_bkg_model())
        self.full_model.append(self.set_bkg_model())
        self.full_model.append(self.set_src_model(0.0, self.m_offset, \
                                                  flux=0.010, \
                                                  type=self.m_type, \
                                                  index=self.m_index))

        # Attach background model to observation container
        self.obs.models(self.bkg_model)
        
        # Return
        return
    
    def models(self, models):
        """
        Set model.
        """
        # Copy models
        self.model = models.copy()
    
        # Return
        return
        
    def execute(self):
        """
        Execute the script.
        """
        # Run the script
        self.run()
        
        # Return
        return

    def run(self):
        """
        Run the script.
        """
        # Switch screen logging on in debug mode
        if self.logDebug():
            self.log.cout(True)

        # Get parameters
        self.get_parameters()
        
        #  Write input parameters into logger
        if self.logTerse():
            self.log_parameters()
            self.log("\n")
        
        # Write observation into logger
        if self.logTerse():
            self.log("\n")
            self.log.header1("Observation")
            self.log(str(self.obs))
            self.log("\n")

        # Write models into logger
        if self.logTerse():
            self.log("\n")
            self.log.header1("Test model")
            self.log(str(self.full_model))
            self.log("\n")

        # Write header
        if self.logTerse():
            self.log("\n")
            self.log.header1("Generate TS distribution")

        # Loop over trials
        for seed in range(self.m_ntrials):
        
            # Make a trial
            result = self.trial(seed)
            
            # Write out result immediately
            if seed == 0:
                file = open(self.m_outfile, 'w')
                writer = csv.DictWriter(file, result['colnames'])
                writer.writerow(dict((_,_) for _ in result['colnames']))
            else:
                file = open(self.m_outfile, 'a')
            writer = csv.DictWriter(file, result['colnames'])
            writer.writerow(result['values'])
            file.close()
        
        # Return
        return
    
    def set_obs(self, lpnt=0.0, bpnt=0.0, emin=0.1, emax=100.0):
        """
        Returns an observation container with a single CTA observation.
        
        Keywords:
         lpnt - Galactic longitude of pointing [deg] (default: 0.0)
         bpnt - Galactic latitude of pointing [deg] (default: 0.0)
         emin - Minimum energy [TeV] (default: 0.1)
         emax - Maximum energy [TeV] (default: 100.0)
        """
        # Allocate observation container
        obs = GObservations()
    
        # Set single pointing
        pntdir = GSkyDir()
        pntdir.lb_deg(lpnt, bpnt)
        
        # Create CTA observation
        run = obsutils.set(pntdir, caldb=self.m_caldb, irf=self.m_irf, \
                           duration=self.m_duration, \
                           emin=emin, emax=emax, rad=self.m_rad)
        
        # Append observation to container
        obs.append(run)
    
        # Return observation container
        return obs
    
    def set_bkg_model(self, fitsigma=False):
        """
        Setup CTA background model.
        """
        # Define radial component
        radial = GCTAModelRadialGauss(3.0)
        if fitsigma:
            radial["Sigma"].free()
        else:
            radial["Sigma"].fix()
        
        # Define spectral component
        spectrum = GModelSpectralFunc(self.m_bkg, 1.0)
        
        # Create background model
        model = GCTAModelRadialAcceptance(radial, spectrum)
        model.name("Background")
        model.instruments("CTA")
    
        # Return background model
        return model
    
    def set_src_model(self, l, b, flux=1.0, index=-2.48, \
                      type="point", sigma=1.0, radius=1.0, width=0.1, \
                      fitpos=False, fitidx=False):
        """
        Returns a single source with Crab-like spectrum. The source flux
        can be scaled in Crab units. The Crab spectrum is based on MAGIC
        observations (Albert et al. 2008, ApJ, 674, 1037).

        Parameters:
         l      - Galactic longitude of source location [deg]
         b      - Galactic latitude of source location [deg]
        Keywords:
         flux   - Source flux [Crabs]
         index  - Spectral index
         type   - Source type ("point", "gauss", "disk", "shell")
         sigma  - Gaussian sigma (for type="gauss")
         radius - Disk or shell inner radius [deg] (for type="disk" and type="shell")
         width  - Shell width [deg] (for type="shell")
         fitpos - Fit position? (default: True)
         fitidx - Fit index? (default: True)
        """
        # Set source location
        location = GSkyDir()
        location.lb_deg(l, b)
    
        # Set source spectrum
        spectrum = GModelSpectralPlaw(flux*5.7e-16, index, GEnergy(0.3, "TeV"))
        if fitidx:
            spectrum["Index"].free()
        else:
            spectrum["Index"].fix() 

        # Set source
        if type == "point":
            spatial = GModelSpatialPointSource(location)
            if fitpos:
                spatial[0].free()
                spatial[1].free()
        elif type == "gauss":
            spatial = GModelSpatialRadialGauss(location, sigma)
            if fitpos:
                spatial[0].free()
                spatial[1].free()
        elif type == "disk":
            spatial = GModelSpatialRadialDisk(location, radius)
            if fitpos:
                spatial[0].free()
                spatial[1].free()
        elif type == "shell":
            spatial = GModelSpatialRadialShell(location, radius, width)
            if fitpos:
                spatial[0].free()
                spatial[1].free()
        else:
            self.log("ERROR: Unknown source type '"+type+"'.\n")
            return None
        source = GModelSky(spatial, spectrum)
    
        # Set source name
        source.name("Test")
    
        # Return source
        return source
    
    def trial(self, seed):
        """
        Create the TS for a single trial.
        
        Parameters:
         seed - Random number generator seed
        """
        # Write header
        if self.logExplicit():
            self.log.header2("Trial "+str(seed+1))

        # Simulate events
        sim = obsutils.sim(self.obs, \
                           nbins=self.m_enumbins, \
                           seed=seed, \
                           binsz=self.m_binsz, \
                           npix=self.m_npix, \
                           log=self.m_log, debug=self.m_debug)

        # Determine number of events in simulation
        nevents = 0.0
        for run in sim:
            nevents += run.events().number()

        # Write simulation results
        if self.logExplicit():
            self.log.header3("Simulation")
            self.log.parformat("Number of simulated events")
            self.log(nevents)
            self.log("\n")

        # Fit background only
        sim.models(self.bkg_model)
        like_bgm   = obsutils.fit(sim, log=self.m_log, debug=self.m_debug)
        result_bgm = like_bgm.obs().models()
        LogL_bgm   = like_bgm.opt().value()
        npred_bgm  = like_bgm.obs().npred()

        # Write background fit results
        if self.logExplicit():
            self.log.header3("Background model fit")
            self.log.parformat("log likelihood")
            self.log(LogL_bgm)
            self.log("\n")
            self.log.parformat("Number of predicted events")
            self.log(npred_bgm)
            self.log("\n")
            for model in result_bgm:
                self.log.parformat("Model")
                self.log(model.name())
                self.log("\n")
                for par in model:
                    self.log(str(par)+"\n")

        # Fit background and test source
        sim.models(self.full_model)
        like_all   = obsutils.fit(sim, log=self.m_log, debug=self.m_debug)
        result_all = like_all.obs().models()
        LogL_all   = like_all.opt().value()
        npred_all  = like_all.obs().npred()
        ts         = 2.0*(LogL_bgm-LogL_all)
        
        # Write background and test source fit results
        if self.logExplicit():
            self.log.header3("Background and test source model fit")
            self.log.parformat("Test statistics")
            self.log(ts)
            self.log("\n")
            self.log.parformat("log likelihood")
            self.log(LogL_all)
            self.log("\n")
            self.log.parformat("Number of predicted events")
            self.log(npred_all)
            self.log("\n")
            for model in result_all:
                self.log.parformat("Model")
                self.log(model.name())
                self.log("\n")
                for par in model:
                    self.log(str(par)+"\n")
                    
        # Write result
        elif self.logTerse():
            self.log.parformat("Trial "+str(seed))
            self.log("TS=")
            self.log(ts)
            self.log("  Prefactor=")
            self.log(result_all["Test"]["Prefactor"].value())
            self.log("+/-")
            self.log(result_all["Test"]["Prefactor"].error())
            self.log("\n")
        
        # Initialise results
        colnames = []
        values   = {}
        
        # Set TS value
        colnames.append("TS")
        values["TS"] = ts

        # Set logL for background fit
        colnames.append("LogL_bgm")
        values["LogL_bgm"] = LogL_bgm

        # Set logL for full fit
        colnames.append("LogL_all")
        values["LogL_all"] = LogL_all

        # Set Nevents
        colnames.append("Nevents")
        values["Nevents"] = nevents

        # Set Npred for background fit
        colnames.append("Npred_bkg")
        values["Npred_bkg"] = npred_bgm

        # Set Npred for full fit
        colnames.append("Npred_all")
        values["Npred_all"] = npred_all
        
        # Gather free full fit parameters
        for i in range(result_all.size()):
            model      = result_all[i]
            model_name = model.name()
            for k in range(model.size()):
                par = model[k]
                if par.is_free():
                
                    # Set parameter name
                    name = model_name+"_"+par.name()
                    
                    # Append value
                    colnames.append(name)
                    values[name] = par.value()
                    
                    # Append error
                    name = "Unc_"+name
                    colnames.append(name)
                    values[name] = par.error()
        
        # Bundle together results
        result = {'colnames': colnames, 'values': values}
        
        # Return
        return result


# ======================== #
# Main routine entry point #
# ======================== #
if __name__ == '__main__':
    """
    Generates TS distribution.
    """
    # Create instance of application
    app = cstsdist(sys.argv)
    
    # Open logfile
    app.logFileOpen()
    
    # Execute application
    app.execute()