#!/usr/bin/env python
##############################################################################
#
# Usage example for the procedure PPXF, which implements the
# Penalized Pixel-Fitting (pPXF) method originally described in
# Cappellari M., & Emsellem E., 2004, PASP, 116, 138
#     http://adsabs.harvard.edu/abs/2004PASP..116..138C
# and upgraded in Cappellari M., 2017, MNRAS, 466, 798
#     http://adsabs.harvard.edu/abs/2017MNRAS.466..798C
#
# This procedure illustrates (and tests) both
#   1. The simultaneous fit of two reflection-symmetric LOSVDs;
#   2. The fit of the sky spectrum together with the templates.
#
# An application of this approach is presented in Section 3 of
#   Weijmans, A-M et al., 2009, MNRAS, 398, 561
#   http://adsabs.harvard.edu/abs/2009MNRAS.398..561W
#
# MODIFICATION HISTORY:
#   V1.0.0: Written by Michele Cappellari, based on a previous IDL procedure.
#       Oxford, 20 April 2017
#   V1.0.1: Updated MILES file names. MC, Oxford, 29 November 2017
#   V1.0.2: Changed imports for pPXF as a package.
#       Make file paths relative to the pPXF package to be able to run the
#       example unchanged from any directory. MC, Oxford, 17 April 2018
#   V1.0.3: Dropped legacy Python 2.7 support. MC, Oxford, 10 May 2018
#
################################################################################

import glob
from os import path
from time import clock

from astropy.io import fits
import numpy as np
from scipy import signal
from numpy.polynomial import legendre
import matplotlib.pyplot as plt

import ppxf as ppxf_package
from ppxf.ppxf import ppxf
import ppxf.ppxf_util as util

################################################################################

def ppxf_example_sky_and_symmetric_losvd():

    ppxf_dir = path.dirname(path.realpath(ppxf_package.__file__))

    # Solar metallicity, Age=12.59 Gyr
    hdu = fits.open(ppxf_dir + '/miles_models/Mun1.30Zp0.00T12.5893_iPp0.00_baseFe_linear_FWHM_2.51.fits')
    ssp = hdu[0].data
    h = hdu[0].header

    lamRange = h['CRVAL1'] + np.array([0.,h['CDELT1']*(h['NAXIS1']-1)])
    velscale = 70.  # km/s
    star, logLam, velscale = util.log_rebin(lamRange, ssp, velscale=velscale)
    star /= np.mean(star)

    # Adopted input parameters =================================================

    vel = 200./velscale    # Velocity of 1st spectrum in pixels (2nd has -vel)
    sigma = 300./velscale   # Dispersion of both spectra in pixels
    h3 = 0.1                # h3 of 1st spectrum (2nd has -h3)
    h4 = 0.1
    sn = 40.
    moments = 4
    deg = 4
    vshift = 10                 # Adopted systemic velocity in pixels
    vsyst = vshift*velscale     # Adopted systemic velocity in km/s

    # Generate input Sky =======================================================
    # For illustration, the sky is modelled as two Gaussian emission lines

    n = star.size
    x = np.arange(n)
    sky1 = np.exp(-0.5*(x - 1000)**2/100)
    sky2 = np.exp(-0.5*(x - 2000)**2/100)

    # Generate input LOSVD =====================================================

    dx = int(abs(vel) + 5*sigma)
    v = np.linspace(-dx, dx, 2*dx + 1)
    w = (v - vel)/sigma
    w2 = w**2
    gauss = np.exp(-0.5*w2)
    gauss /= np.sum(gauss)
    h3poly = w*(2*w2 - 3)/np.sqrt(3)
    h4poly = (w2*(4*w2 - 12) + 3)/np.sqrt(24)
    losvd = gauss*(1 + h3*h3poly + h4*h4poly)

    # Generate first synthetic spectrum ========================================
    # The template is convolved with the LOSVD

    x = np.linspace(-1, 1, n)
    galaxy1 = signal.fftconvolve(star, losvd, mode="same")
    galaxy1 = np.roll(galaxy1, vshift)   # Mimic nonzero systemic velocity
    galaxy1 *= legendre.legval(x, np.append(1, np.random.uniform(-0.1, 0.1, deg-1)))  # Multiplicative polynomials
    galaxy1 += legendre.legval(x, np.random.uniform(-0.1, 0.1, deg))  # Additive polynomials
    galaxy1 += sky1 + 2*sky2    # Add two sky lines
    galaxy1 = np.random.normal(galaxy1, 1/sn)   # Add noise

    # Generate symmetric synthetic spectrum ====================================
    # The same template is convolved with a reversed LOSVD
    # and different polynomials and sky lines are included

    galaxy2 = signal.fftconvolve(star, np.flip(losvd, 0), mode="same")
    galaxy2 = np.roll(galaxy2, vshift)   # Mimic nonzero systemic velocity
    galaxy2 *= legendre.legval(x, np.append(1, np.random.uniform(-0.1, 0.1, deg-1)))  # Multiplicative polynomials
    galaxy2 += legendre.legval(x, np.random.uniform(-0.1, 0.1, deg))  # Additive polynomials
    galaxy2 += 2*sky1 + sky2    # Add two sky lines
    galaxy2 = np.random.normal(galaxy2, 1/sn)   # Add noise

    # Load spectral templates ==================================================

    vazdekis = glob.glob(ppxf_dir + '/miles_models/Mun1.30Z*.fits')
    templates = np.empty((n, len(vazdekis)))
    for j, file in enumerate(vazdekis):
        hdu = fits.open(file)
        ssp = hdu[0].data
        sspNew, logLam2, velscale = util.log_rebin(lamRange, ssp, velscale=velscale)
        templates[:, j] = sspNew/np.median(sspNew)  # Normalize templates

    # Do the fit ===============================================================

    # Input both galaxy spectra simultaneously to pPXF
    galaxy = np.column_stack([galaxy1, galaxy2])

    # Use two sky templates for each galaxy spectrum
    sky = np.column_stack([sky1, sky2])

    # Randomized starting guess
    vel0 = vel + np.random.uniform(-1, 1)
    sigma0 = sigma*np.random.uniform(0.8, 1.2)
    start = np.array([vel0, sigma0])*velscale  # Convert to km/s
    goodpixels = np.arange(50, n - 50)

    print("\nThe input values are: Vel=%0.0f, sigma=%0.0f, h3=%0.1f, h4=%0.1f\n" %
          (vel*velscale, sigma*velscale, h3, h4))

    t = clock()

    pp = ppxf(templates, galaxy, np.full_like(galaxy, 1/sn), velscale, start,
              goodpixels=goodpixels, plot=1, moments=moments,
              vsyst=vsyst, mdegree=deg, degree=deg, sky=sky)

    print('Elapsed time in pPXF: %.2f s' % (clock() - t))
    plt.pause(1)

################################################################################

if __name__ == '__main__':

    np.random.seed(123)  # For reprodcible results
    ppxf_example_sky_and_symmetric_losvd()
