# First pass redshift finder of redmonster.  Able to search entire
# parameter-space and redshift-space of models in increments
# of d(loglam)/d(pixel)
#
# Tim Hutchinson, University of Utah @ IAC, May 2014
# t.hutchinson@utah.edu
from __future__ import division
from os import environ, makedirs
from os.path import join, exists

import numpy as n
from astropy.io import fits
from scipy import linalg
from scipy.optimize import nnls

import multiprocessing
import time

from redmonster.datamgr.ssp_prep import SSPPrep
from redmonster.physics.misc import poly_array, two_pad
from redmonster.datamgr.io2 import read_ndArch, write_chi2arr
from redmonster.physics.misc import create_mask

# Assumes all templates live in $REDMONSTER_DIR/templates/


def _zchi2(arg):
    return zchi2_single_template(**arg)


def _zchi2_no_poly(arg):
    return zchi2_single_template_no_poly(**arg)


def zchi2_single_template(j, poly_fft, t_fft, t2_fft, data_fft, ivar_fft,
                          pmat_pol, bvec_pol, chi2_0, chi2_null, num_z,
                          npixstep, zminpix, flag_val_neg_model):
    npoly = poly_fft.shape[0]  # degree+1 of polynomial
    fftnaxis1 = t_fft.shape[0]  # number of redshift bins
    pmat = pmat_pol.copy()  # precomputed block for polynomial terms ,
    # n.zeros( (npoly+1, npoly+1, fftnaxis1), dtype=float)
    bvec = bvec_pol.copy()  # precomputed vector block for polynomial terms ,
    # n.zeros( (npoly+1, fftnaxis1), dtype=float)

    # fill matrix
    pmat[0, 0] = n.fft.ifft(t2_fft * ivar_fft.conj()).real
    bvec[0] = n.fft.ifft(t_fft * data_fft.conj()).real
    for ipos in range(npoly):
        pmat[ipos+1, 0] = pmat[0, ipos+1] = n.fft.ifft(t_fft *
                                                       poly_fft[ipos].conj()
                                                       ).real

    # solve for each z

    zchi2arr = n.zeros((num_z))
    zwarning = n.zeros((num_z))
    for l in n.arange(num_z)*npixstep:
        try:  # try to solve for this redshift
            f = linalg.solve(pmat[:, :, l+zminpix], bvec[:, l+zminpix])

            zchi2arr[(l//npixstep)] = chi2_0 - \
                n.dot(n.dot(f, pmat[:, :, l + zminpix]), f)  # is this true ??
            if f[0] < 0:
                zwarning[(l//npixstep)] = (int(zwarning[(l//npixstep)]) |
                                           flag_val_neg_model)
                zchi2arr[(l//npixstep)] = chi2_null
                try:
                    n.dot(n.dot(f, pmat[:, :, l+zminpix]), f)
                except Exception as e:
                    print("Except: %r" % e)
                    zchi2arr[(l//npixstep)] = chi2_null
        except Exception as e:  # failure
            print("Exception: %s" % e)
            zchi2arr[(l//npixstep)] = chi2_null
    return j, zchi2arr, zwarning


def zchi2_single_template_no_poly(j, t_fft, t2_fft, data_fft, ivar_fft, chi2_0,
                                  num_z, npixstep, zminpix,
                                  flag_val_neg_model):

    a = n.fft.ifft(t2_fft * ivar_fft.conj()).real
    b = n.fft.ifft(t_fft * data_fft.conj()).real

    ii = zminpix+n.arange(num_z)*npixstep
    f = (a[ii] != 0)*b[ii]/(a[ii]+(a[ii] == 0))
    zchi2arr = chi2_0 - a[ii]*f**2  # is this true ?????

    zwarning = n.zeros((num_z))
    zchi2arr[f < 0] = chi2_0
    zwarning[f < 0] = flag_val_neg_model
    return j, zchi2arr, zwarning


class ZFinder:
    def __init__(self, fname=None, group=[0], npoly=None, zmin=None, zmax=None,
                 nproc=1):
        self.fname = fname
        if type(group) == list:
            self.group = group
        elif type(group) == int:
            self.group = [group, ]
        else:
            self.group = eval(group)
        self.npoly = npoly if npoly is not None else 4
        self.zmin = float(zmin)
        self.zmax = float(zmax)
        self.nproc = nproc
        self.pixoffset = None
        self.zchi2arr = None

        try:
            self.templatesdir = environ['REDMONSTER_TEMPLATES_DIR']
        except KeyError as e:
            self.templatesdir = None
            print("Enviromental variable 'REDMONSTER_TEMPLATES_DIR' not set: \
                    %r" % e)
        self.read_template()
        self.npars = len(self.templates.shape) - 1
        self.templates_flat = n.reshape(self.templates, (-1, self.fftnaxis1))
        self.t_fft = n.fft.fft(self.templates_flat)
        self.t2_fft = n.fft.fft(self.templates_flat**2)
        self.f_nulls = []
        self.chi2_null = []
        self.sn2_data = []

    def read_template(self):
        self.templates, self.baselines, self.infodict = \
                read_ndArch(join(self.templatesdir, self.fname))
        self.type = self.infodict['class']
        self.origshape = self.templates.shape
        self.ntemps = self.templates[..., 0].size
        self.fftnaxis1 = two_pad(self.origshape[-1])
        self.tempwave = 10**(self.infodict['coeff0'] +
                             n.arange(self.infodict['nwave']) *
                             self.infodict['coeff1'])
        templates_pad = n.zeros(self.origshape[:-1]+(self.fftnaxis1,))
        templates_pad[..., :self.origshape[-1]] = self.templates
        self.templates = templates_pad

    def create_z_baseline(self, loglam0):
        self.zbase = ((10**loglam0)/self.tempwave) - 1

    def conv_zbounds(self):
        zmaxpix = n.where(abs((self.zbase-self.zmin)) ==
                          n.min(abs(self.zbase-self.zmin)))[0][0]
        zminpix = n.where(abs((self.zbase-self.zmax)) ==
                          n.min(abs(self.zbase-self.zmax)))[0][0]
        return zminpix, zmaxpix

    def zchi2(self, specs, specloglam, ivar, npixstep=1, chi2file=False,
              plate=None, mjd=None, fiberid=None, linelist=None):
        self.chi2file = chi2file
        self.npixstep = npixstep
        # Apply mask by making ivar[i-5:i+5] = 0 where wave[i] is the
        # wavelenght of an emission feature in linelist
        if linelist is not None:
            self.linelist = linelist
            mask = create_mask(linelist, specloglam)
            ivar *= mask
        self.zwarning = n.zeros(specs.shape[0])
        flag_val_unplugged = int('0b10000000', 2)
        flag_val_neg_model = int('0b1000', 2)
        self.create_z_baseline(specloglam[0])
        if (self.zmin != None) and (self.zmax != None) and \
                (self.zmax > self.zmin):

            zminpix, zmaxpix = self.conv_zbounds()
            self.pixoffset = zminpix
            num_z = int(n.floor((zmaxpix - zminpix) / npixstep))
            zinds = zminpix + n.arange(num_z)*npixstep
            self.zbase = self.zbase[zinds]
        else:
            zminpix = 0
            # Number of pixels to be fitted in redshift
            num_z = int(n.floor((self.origshape[-1] - specs.shape[-1]) /
                                npixstep))

        # Create arrays for use in routine
        # Create chi2 array of shape (# of fibers, template_parameter_1,
        # ..., template_parameter_N, # of redshifts)
        zchi2arr = n.zeros((specs.shape[0], self.templates_flat.shape[0],
                            num_z))
        temp_zwarning = n.zeros(zchi2arr.shape)

        # Pad data and SSPs to a power of 2 for faster FFTs
        data_pad = n.zeros(specs.shape[:-1] + (self.fftnaxis1,), dtype=float)
        data_pad[..., :specs.shape[-1]] = specs
        ivar_pad = n.zeros(ivar.shape[:-1] + (self.fftnaxis1,), dtype=float)
        ivar_pad[..., :specs.shape[-1]] = ivar

        # Pre-compute FFTs for use in convolutions
        data_fft = n.fft.fft(data_pad * ivar_pad)
        ivar_fft = n.fft.fft(ivar_pad)

        if self.npoly > 0:
            # Compute poly terms, noting that they will stay fixed with
            # the data - assumes data is passed in as shape (nfibers, npix)
            polyarr = poly_array(self.npoly, specs.shape[1])
            pmat = n.zeros((self.npoly+1, self.npoly+1, self.fftnaxis1),
                           dtype=float)
            bvec = n.zeros((self.npoly+1, self.fftnaxis1), dtype=float)

            # Pad to a power of 2 for faster FFTs
            poly_pad = n.zeros((self.npoly, self.fftnaxis1), dtype=float)
            poly_pad[..., :polyarr.shape[-1]] = polyarr

            # Pre-compute FFTs for use in convolutions
            poly_fft = n.zeros((ivar_pad.shape[0], self.npoly, self.fftnaxis1),
                               dtype=complex)
            for i in range(self.npoly):
                poly_fft[:, i, :] = n.fft.fft(poly_pad[i] * ivar_pad)

        # Compute z for all fibers

        for i in range(specs.shape[0]):  # Loop over fibers

            start = time.time()

            # print 'INFO Fitting fiber %s of %s for template %s' % \
            #        (i+1, specs.shape[0], self.fname)

            # If flux is all zeros, flag as unplugged according to BOSS
            # zwarning flags and don't bother with doing fit
            if len(n.where(specs[i] != 0.)[0]) == 0:
                self.zwarning[i] = int(self.zwarning[i]) | flag_val_unplugged
            else:  # Otherwise, go ahead and do fit

                self.sn2_data.append(n.sum((specs[i]**2)*ivar[i]))

                if self.npoly > 0:
                    for ipos in range(self.npoly):
                        bvec[ipos+1] = n.sum(poly_pad[ipos] * data_pad[i] *
                                             ivar_pad[i])
                    for ipos in range(self.npoly):
                        for jpos in range(self.npoly):
                            pmat[ipos+1, jpos+1] = n.sum(poly_pad[ipos] *
                                                         poly_pad[jpos] *
                                                         ivar_pad[i])
                    try:
                        f_null = linalg.solve(pmat[1:, 1:, 0], bvec[1:, 0])
                    except n.linalg.linalg.LinAlgError:
                        f_null = 0
                    self.f_nulls.append(f_null)
                    self.chi2_null.append(self.sn2_data[i] -
                                          n.dot(n.dot(f_null,
                                                      pmat[1:, 1:, 0]),
                                                f_null))
                else:
                    self.chi2_null.append(self.sn2_data[i])
                # print 'INFO Chi^2_null value is %s' % self.chi2_null[i]
                # Loop over templates
                # multiprocessing

                func_args = []
                for j in range(self.templates_flat.shape[0]):
                    if self.npoly > 0:
                        arguments = {"j": j,
                                     "poly_fft": poly_fft[i],
                                     "t_fft": self.t_fft[j],
                                     "t2_fft": self.t2_fft[j],
                                     "data_fft": data_fft[i],
                                     "ivar_fft": ivar_fft[i],
                                     "pmat_pol": pmat,
                                     "bvec_pol": bvec,
                                     "chi2_0": self.sn2_data[i],
                                     "chi2_null": self.chi2_null[i],
                                     "num_z": num_z,
                                     "npixstep": self.npixstep,
                                     "zminpix": zminpix,
                                     "flag_val_neg_model": flag_val_neg_model
                                     }
                    else:
                        arguments = {"j": j,
                                     "t_fft": self.t_fft[j],
                                     "t2_fft": self.t2_fft[j],
                                     "data_fft": data_fft[i],
                                     "ivar_fft": ivar_fft[i],
                                     "chi2_0": self.sn2_data[i],
                                     "num_z": num_z,
                                     "npixstep": self.npixstep,
                                     "zminpix": zminpix,
                                     "flag_val_neg_model": flag_val_neg_model
                                     }
                    func_args.append(arguments)

                results = None
                if self.nproc > 1:
                    pool = multiprocessing.Pool(self.nproc)
                    if self.npoly > 0:
                        results = pool.map(_zchi2, func_args)
                    else:
                        results = pool.map(_zchi2_no_poly, func_args)
                    pool.close()
                    pool.join()
                else:
                    if self.npoly > 0:
                        results = [_zchi2(x) for x in func_args]
                    else:
                        results = [_zchi2_no_poly(x) for x in func_args]

                for result in results:
                    j = result[0]
                    zchi2arr[i, j] = result[1]
                    temp_zwarning[i, j] = result[2]

                stop = time.time()

                print("INFO fitted spectrum %d/%d, chi2_null=%f, %d templates \
                    in %s, npoly=%d, using %d procs in %f \
                    sec" % (i+1, specs.shape[0], self.chi2_null[i],
                            self.templates_flat.shape[0], self.fname,
                            self.npoly, self.nproc, stop-start))

        # Use only neg_model flag from best fit model/redshift and add
        # it to self.zwarning
        for i in range(self.zwarning.shape[0]):
            minpos = (n.where(zchi2arr[i] == n.min(zchi2arr[i]))[0][0],
                      n.where(zchi2arr[i] == n.min(zchi2arr[i]))[1][0])
            self.zwarning[i] = int(self.zwarning[i]) | \
                int(temp_zwarning[i, minpos[0], minpos[1]])
        zchi2arr = n.reshape(zchi2arr, (specs.shape[0],) + self.origshape[:-1]
                             + (num_z,))
        bestl = n.where(zchi2arr == n.min(zchi2arr))[-1][0]
        thisz = ((10**(specloglam[0]))/self.tempwave[bestl+zminpix])-1

        # return zchi2arr
        self.zchi2arr = zchi2arr
        # self.store_models(specs, ivar)
        if self.chi2file is True:
            if (plate is not None) & (mjd is not None) & (fiberid is not None):
                write_chi2arr(plate, mjd, fiberid, self.zchi2arr)
            else:
                print('WARNING Plate/mjd/fiberid not given - unable to write \
                    chi2 file!')
        else:
            print('INFO Not writing chi2')

    def store_models(self, specs, ivar):
        self.models = n.zeros((specs.shape))
        for i in range(self.models.shape[0]):
            minloc = n.unravel_index(self.zchi2arr[i].argmin(),
                                     self.zchi2arr[i].shape)
            pmat = n.zeros((specs.shape[-1], self.npoly+1))
            this_temp = self.templates[minloc[:-1]]
            pmat[:, 0] = this_temp[(minloc[-1]*self.npixstep) +
                                   self.pixoffset:(minloc[-1]*self.npixstep) +
                                   self.pixoffset+specs.shape[-1]]
            polyarr = poly_array(self.npoly, specs.shape[-1])
            pmat[:, 1:] = n.transpose(polyarr)
            ninv = n.diag(ivar[i])
            try:  # Some eBOSS spectra have ivar[i] = 0 for all i
                '''
                f = linalg.solve( n.dot(n.dot(n.transpose(pmat),ninv),pmat),
                                n.dot( n.dot(n.transpose(pmat),ninv),specs[i]))
                '''
                f = nnls(n.dot(n.dot(n.transpose(pmat), ninv), pmat),
                         n.dot(n.dot(n.transpose(pmat), ninv), specs[i]))
                f = n.array(f)[0]
                self.models[i] = n.dot(pmat, f)
            except Exception as e:
                self.models[i] = n.zeros(specs.shape[-1])
                print("Exception: %r" % e)

    def write_chi2arr(self, plate, mjd, fiberid):
        prihdu = fits.PrimaryHDU(self.zchi2arr)
        thdulist = fits.HDUList([prihdu])
        try:
            rsr = environ['REDMONSTER_SPECTRO_REDUX']
            run2d = environ['RUN2D']
            run1d = environ['RUN1D']
            if (rsr is not None) & (run2d is not None) & (run1d is not None):
                testpath = join(rsr, run2d, '%s' % plate, run1d)
                if exists(testpath):
                    dest = testpath
                else:
                    try:
                        makedirs(testpath)
                        dest = testpath
                    except Exception as e:
                        print("Exception: %r" % e)
                        dest = None
        except Exception as e:
            print("Exception: %r" % e)
            dest = None
        if dest is not None:
            try:
                thdulist.writeto(join(dest, '%s' % 'chi2arr-%s-%s-%s-%03d.fits'
                                      % (self.type, plate, mjd, fiberid)),
                                 overwrite=True)
                print('Writing chi2 file to %s' %
                      join(dest, '%s' % 'chi2arr-%s-%s-%s-%03d.fits' %
                           (self.type, plate, mjd, fiberid)))
            except Exception as e:
                print("Exception: %r" % e)
                print('Environment variables not set or path does not exist - \
                        not writing chi2 file!')
        else:
            print('Environment variables not set or path does not exist - \
                    not writing chi2 file!')
