#! /usr/bin/env python
# ==========================================================================
# This script displays one or several CTA response functions.
#
# Required 3rd party modules:
# - matplotlib
# - numpy
#
# Copyright (C) 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 sys
import math
import gammalib
import matplotlib.pyplot as plt


# ================================== #
# Compute Eq. (17) of Li & Ma (1983) #
# ================================== #
def sigma_lima(Non, Noff, alpha=0.2):
	"""
	Compute Eq. (17) of Li & Ma (1983).
	
	Parameters:
	 Non   - Number of on counts
	 Noff  - Number of off counts
	Keywords:
	 alpha - Ratio of on-to-off exposure
	"""
	# Compute sensitivity
	alpha1 = alpha + 1.0
	sum    = Non + Noff
	arg1   = Non / sum
	arg2   = Noff / sum
	term1  = Non  * math.log((alpha1/alpha)*arg1)
	term2  = Noff * math.log(alpha1*arg2)
	sigma  = math.sqrt(2.0 * (term1 + term2))
	
	# Return sensitivity
	return sigma


# ======================================== #
# Solve Eq. (17) of Li & Ma (1983) for Non #
# ======================================== #
def Non_lima(sigma, Noff, alpha=0.2):
	"""
	Solve Eq. (17) of Li & Ma (1983) for Non
	
	Parameters:
	 sigma - Required significance
	 Noff  - Number of off counts
	Keywords:
	 alpha - Ratio of on-to-off exposure
	"""
	# Set initial guess for Non
	S   = 25.0
	Non = S + alpha*Noff
	
	# Compute initial sigma estimate
	s     = sigma_lima(Non, Noff, alpha=alpha)
	delta = s - sigma
	
	# Iterate
	while abs(delta) > 1.0e-6:
		S    *= sigma/s
		Non   = S + alpha*Noff
		s     = sigma_lima(Non, Noff, alpha=alpha)
		delta = s - sigma
	
	# Return
	return Non


# ======================= #
# Show one effective area #
# ======================= #
def show_one_effective_area(rsp, name, color="r"):
    """
    Show one effective area.
    """
    # Set to figure 1
    plt.figure(1)
    
    # Generate logE(TeV) vector
    nbins = 20
    logE  = [-1.7+0.2*i for i in range(nbins)]
    E     = [math.pow(10.0, e) for e in logE]

    # Generate effective area vector
    aeff = [0.0 for i in range(nbins)]
    for i in range(nbins):
        aeff[i] = rsp.aeff()(logE[i])

    # Plot data
    plt.loglog(E, aeff, color+'-', label=name)
    plt.loglog(E, aeff, color+'o')

    # Set axes
    plt.xlabel("Energy [TeV]")
    plt.ylabel("Effective area [cm2]")

    # Set legend
    plt.legend(loc="lower right")

    # Return
    return


# ======================== #
# Show one background rate #
# ======================== #
def show_one_background_rate(rsp, name, color="r"):
    """
    Show one background rate.
    """
    # Set to figure 2
    plt.figure(2)
    
    # Generate logE(TeV) vector
    nbins = 20
    logE  = [-1.7+0.2*i for i in range(nbins)]
    E     = [math.pow(10.0, e) for e in logE]

    # Generate on-axis background rate vector
    bgrate = [0.0 for i in range(nbins)]
    for i in range(nbins):
        bgrate[i] = rsp.background()(logE[i], 0.0, 0.0)
        if bgrate[i] < 1.0e-10:
            bgrate[i] = 0.0
    
    # Plot data
    plt.loglog(E, bgrate, color+'-', label=name)
    plt.loglog(E, bgrate, color+'o')

    # Set axes
    plt.xlabel("Energy [TeV]")
    plt.ylabel("Background rate [events/s/MeV/sr]")

    # Set legend
    plt.legend(loc="lower left")

    # Return
    return


# ==================== #
# Show one sensitivity #
# ==================== #
def show_one_sensitivity(rsp, name, color="r", duration=180000.0, alpha=0.2, sigma=5.0):
    """
    Show one sensitivity.
    """
    # Set constants
    r68_to_sigma = 0.6624305

    # Set to figure 3
    plt.figure(3)
    
    # Generate logE(TeV) vector
    nbins = 20
    logE  = [-1.7+0.2*i for i in range(nbins)]
    E     = [math.pow(10.0, e) for e in logE]

    # Generate sensitivity vector
    flux = [0.0 for i in range(nbins)]
    for i in range(nbins):
        ewidth     = (math.pow(10.0, logE[i]+0.1) - \
                      math.pow(10.0, logE[i]-0.1))*1.0e6
        aeff       = rsp.aeff()(logE[i])
        r68        = rsp.psf().delta_max(logE[i]) / 5.0 / r68_to_sigma
        solidangle = gammalib.twopi * (1.0 - math.cos(r68))
        bgd_counts = rsp.background()(logE[i], 0.0, 0.0) * duration * ewidth * solidangle
        Noff       = bgd_counts/alpha
        if Noff > 0:
            Non        = Non_lima(sigma, Noff)
            src_counts = Non - bgd_counts
            flux[i]    = src_counts / (0.68*aeff * duration) * E[i]*1.0e6 * gammalib.MeV2erg
        print logE[i], flux[i]
    
    # Plot data
    plt.loglog(E, flux, color+'-', label=name)
    plt.loglog(E, flux, color+'o')

    # Set axes
    plt.xlabel("Energy [TeV]")
    plt.ylabel("Sensitivity [erg/cm2/s]")

    # Set legend
    plt.legend(loc="upper right")

    # Return
    return


# ==================== #
# Show one sensitivity #
# ==================== #
def show_one_response(rspname, dbname, name, color="r"):
    """
    Show one response.
    """
    # Set-up calibration database
    caldb = gammalib.GCaldb()
    if gammalib.dir_exists(dbname):
        caldb.rootdir(dbname)
    else:
        caldb.open("cta", dbname)

    # Load response function
    rsp = gammalib.GCTAResponse(rspname, caldb)

    # Show effective area
    show_one_effective_area(rsp, name, color=color)

    # Show background rate
    show_one_background_rate(rsp, name, color=color)

    # Show sensitivity
    show_one_sensitivity(rsp, name, color=color)
    
    # Return
    return
    
    
# ======================== #
# Main routine entry point #
# ======================== #
if __name__ == '__main__':
    """
    Display spectrum in PHA format.
    """
    # Print usage information
    #usage = "Usage: show_pha filename" \
    #        " [-n bins] [-c column] [-t title] [-p plot]"
    #if len(sys.argv) < 2:
    #    print(usage)
    #    sys.exit()

    # Extract parameters
    #filename = sys.argv[1]

    # Create figures
    plt.figure(1)
    plt.title("Effective area")
    plt.figure(2)
    plt.title("Background rate")

    # Set response dictionary
    rsps = [
            {'dbname':  "/usr/local/gamma/share/caldb/data/cta/dummy", \
             'rspname': "cta_dummy_irf", \
             'name':    "Prod1 (Array E, MPIK)", \
             'color':   "b"}, \
            {'dbname':  "e", \
             'rspname': "IFAE20120510_50h", \
             'name':    "Prod1 (Array E, IFAE)", \
             'color':   "g"}, \
            {'dbname':  "aar", \
             'rspname': "DESY20140105_50h", \
             'name':    "Prod2 (Aar, DESY)", \
             'color':   "r"} \
            ]

    # Loop over all responses
    for rsp in rsps:

        # Get attributes
        dbname  = rsp['dbname']
        rspname = rsp['rspname']
        color   = rsp['color']
        name    = rsp['name']
        
        # Show response
        show_one_response(rspname, dbname, name, color=color)

    # Show plot
    plt.show()