#! /usr/bin/env python
# ==========================================================================
# This script simulates CTA surveys. It can be considered as the master
# script for CTA science simulations. Please modify as needed.
#
# 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 *
from math import *
import os
import glob
import sys
from ctools import obsutils


# =========== #
# Plot counts #
# =========== #
def plot_counts(obs):
    """
    Plot counts.
    """
    # Only proceed if matplotlib is available
    try:
        # Import matplotlib
        import matplotlib.pyplot as plt
        
        # Set legend fontsize
        params = {'legend.fontsize': 10}
        plt.rcParams.update(params)
        
        # Set plot styles
        styles = ['b-', 'g-', 'y-', 'n-']
        
        # Dump header
        print("")
        print("Make plots (using matplotlib):")
        print("==============================")
        
        # Create figure 1
        plt.figure(1,figsize=(12,6))
        plt.subplots_adjust(hspace=.7)
        
        # Create subplot 1
        plt.subplot(121)
        plt.title("Spectrum (summed over all pixels)")
        
        # Loop over observations
        for run in obs:
        
            # Get event list
            list = run.events()
            
            # Create energy axis
            ebounds = GEbounds()
            emin    = GEnergy()
            emax    = GEnergy()
            emin.TeV(0.1)
            emax.TeV(100.0)
            ebounds.setlog(emin, emax, 10)
            energy = [ebounds.elogmean(i).TeV() for i in range(ebounds.size())]
            
            # Create spectrum
            print("Extract data:")
            counts = [0.0 for i in range(ebounds.size())]
            print(list.size())
            for atom in list:
                index         = ebounds.index(atom.energy())
                counts[index] = counts[index] + 1.0
            
            # Create error bars
            error = [sqrt(c) for c in counts]
            
            # Plot spectrum
            plt.loglog(energy, counts, 'ro', label='data')
            #plt.errorbar(energy, counts, error, fmt=None, ecolor='r')
            
            # Extract models
            print("Extract models:")
            sum_model = [0.0 for i in range(ebounds.size())]
            for k, m in enumerate(obs.models()):
                print("- "+m.name())
                model = [0.0 for i in range(ebounds.size())]
                for atom in list:
                    index        = ebounds.index(atom.energy())
                    prob         = m.eval(atom, run)
                    model[index] = model[index] + prob * atom.size()
                for i in range(ebounds.size()):
                    sum_model[i] = sum_model[i] + model[i]
                #plt.loglog(energy, model, styles[k], label=m.name())
            #plt.loglog(energy, sum_model, 'r-', label='total')
        
        # Put labels
        plt.xlabel("Energy (TeV)")
        plt.ylabel("Counts")
        plt.legend(loc="lower left")
        
        # Create subplot 2
        plt.subplot(122)
        plt.title("Offset (summed over all energies)")
        
        # Set Crab direction
        crab = GSkyDir()
        crab.radec_deg(83.63, 22.01)
        
        # Loop over observations
        for run in obs:
        
            # Get event list
            list = run.events()
            
            # Create offset histogram
            print("Extract data:")
            nx       = 30
            doffset2 = 0.01
            offset2  = [(i+0.5)*doffset2 for i in range(nx)]
            counts   = [0.0  for i in range(nx)]
            for atom in list:
                off   = atom.dir().dist_deg(crab)
                off2  = off*off
                index = int(off2/doffset2)
                if index < nx:
                    counts[index] = counts[index] + 1.0
            
            # Create error bars
            error = [sqrt(c) for c in counts]
            
            # Plot distribution
            #plt.semilogy(offset2, counts, 'ro', label='data')
            
            # Extract models
            print("Extract models:")
            sum_model = [0.0 for i in range(nx)]
            for k, m in enumerate(obs.models()):
                print("- "+m.name())
                model = [0.0 for i in range(nx)]
                for atom in list:
                    off   = atom.dir().dist_deg(crab)
                    off2  = off*off
                    index = int(off2/doffset2)
                    if index < nx:
                        prob         = m.eval(atom, run)
                        model[index] = model[index] + prob * atom.size()
                for i in range(nx):
                    sum_model[i] = sum_model[i] + model[i]
                #plt.plot(offset2, model, styles[k], label=m.name())
            #plt.plot(offset2, sum_model, 'r-', label='total')
            #plt.ylim(ymin=0.1)
        
        # Put labels
        plt.xlabel("Offset (deg^2)")
        plt.ylabel("Counts")
        #plt.legend(loc="upper right")
        
        # Show counts spectra
        plt.show()
        
    except ImportError:
        print("Matplotlib is not (correctly) installed on your system.")

    # Return
    return


# ======================================== #
# Add CTA background model to observations #
# ======================================== #
def add_background_model(obs):
    """
    Add standard CTA background model to observations container.
    
    We use a simple power law here, scaled to Konrad's E configuration
    performance table. The model needs still to be validated.
    """
    # Recover models from observation
    models = obs.models()
    
    # Define background model
    bgd_radial   = GCTAModelRadialGauss(3.0)
    bgd_spectrum = GModelSpectralPlaw(61.8e-6, -1.85, GEnergy(1.0, "TeV"))
    bgd_model = GCTAModelRadialAcceptance(bgd_radial, bgd_spectrum)
    bgd_model.name("Background")
    bgd_model.instruments("CTA")
    
    # Add background model to container
    models.append(bgd_model)
    
    # Put container back in observation container
    obs.models(models)
    
    # Return observation container
    return obs


# ================= #
# Set Crab spectrum #
# ================= #
def crab_spec():
    """
    Set Crab spectrum based on MAGIC observations
    (Albert et al. 2008, ApJ, 674, 1037)
    """
    # Set parameters
    spectrum = GModelSpectralPlaw(5.7e-16, -2.48, GEnergy(0.3, "TeV"))
    
    # Return spectrum
    return spectrum


# =============================== #
# Setup single observation survey #
# =============================== #
def survey_single():
    """
    Creates a single observation survey for test purposes.
    """
    # Allocate observation container
    obs = GObservations()
    
    # Set single pointing at galactic centre
    pntdir = GSkyDir()
    pntdir.lb_deg(0.0, 0.0)
    run = obsutils.set_obs(pntdir)
    obs.append(run)
    
    # Define single point source with Crab flux at galactic centre
    center = GSkyDir()
    center.lb_deg(0.0, 0.0)
    point_spatial  = GModelSpatialPointSource(center)
    point_spectrum = crab_spec()
    point          = GModelSky(point_spatial, point_spectrum)
    point.name('GC source')
    
    # Create model container
    models = GModels()
    models.append(point)
    obs.models(models)
    
    # Return observation container
    return obs


# =========================== #
# Setup Galactic plane survey #
# =========================== #
def survey_gplane(lrange=10, lstep=2):
    """
    Creates a single observation survey for test purposes.
    
    Keywords:
     lrange - Longitude range (integer deg)
     lstep  - Longitude step size (integer deg)
    """
    # Allocate observation container
    obs = GObservations()
    
    # Loop over longitudes
    for l in range(-lrange,lrange+lstep,lstep):
        
        # Set pointing
        pntdir = GSkyDir()
        pntdir.lb_deg(l, 0.0)
        run = obsutils.set_obs(pntdir)
        run.id(str(l))
        obs.append(run)
    
    # Define single point source with Crab flux at galactic centre
    center = GSkyDir()
    center.lb_deg(0.0, 0.0)
    point_spatial  = GModelSpatialPointSource(center)
    point_spectrum = crab_spec()
    point          = GModelSky(point_spatial, point_spectrum)
    point.name('GC source')
    
    # Create model container
    models = GModels()
    models.append(point)
    obs.models(models)
    
    # Return observation container
    return obs


#==========================#
# Main routine entry point #
#==========================#
if __name__ == '__main__':
    """
    CTA survey simulation and analysis script.
    """
    # Initialise flags
    need_help = False
    
    # Test for command line arguments
    print(sys.argv[0])
    if (len(sys.argv) > 1):
        if sys.argv[1] == "-h":
            need_help = True
        else:
            need_help = True

    # Print help if needed and exit
    if need_help:
        print("Usage: example_survey.py [OPTIONS]")
        print("     -h       Display this usage message")
        sys.exit()
    
    # Dump header
    print("***********************")
    print("* Simulate CTA survey *")
    print("***********************")
    
    # Remove any existing result files
    list = [glob.glob("*.fits"), glob.glob("*.log"), glob.glob("*.xml")]
    for files in list:
        for file in files:
            os.remove(file)
    
    # Setup single observation survey
    #obs = survey_single()
    obs = survey_gplane()
    
    # Add background model
    obs = add_background_model(obs)
    
    # Simulate events
    print("Simulate events")
    obs = obsutils.sim(obs)
    
    # Make counts map
    print("Make counts map")
    cntmap = obsutils.cntmap(obs)
    
    # Fit observations
    print("Fit observations")
    like = obsutils.fit(obs)
    #print like.opt()
    #print like.obs().models()
    
    # Make model map
    print("Make model map (this step will take some time)")
    modmap = obsutils.modmap(obs)
    
    # Show fit results
    #plot_counts(like.obs())