from __future__ import division, absolute_import, print_function
from past.builtins import xrange

import numpy as np
import os
import sys
import esutil
import time

from .fgcmUtilities import _pickle_method
from .fgcmUtilities import objFlagDict
from .fgcmUtilities import retrievalFlagDict

import types
try:
    import copy_reg as copyreg
except ImportError:
    import copyreg

import multiprocessing
from multiprocessing import Pool

from .sharedNumpyMemManager import SharedNumpyMemManager as snmm

copyreg.pickle(types.MethodType, _pickle_method)

## FIXME: derivatives should not be zero when hitting the boundary (check)

class FgcmChisq(object):
    """
    Class which computes the chi-squared for the fit.

    parameters
    ----------
    fgcmConfig: FgcmConfig
       Config object
    fgcmPars: FgcmParameters
       Parameter object
    fgcmStars: FgcmStars
       Stars object
    fgcmLUT: FgcmLUT
       LUT object

    Config variables
    ----------------
    nCore: int
       Number of cores to run in multiprocessing
    nStarPerRun: int
       Number of stars per run.  More can use more memory.
    noChromaticCorrections: bool
       If set to True, then no chromatic corrections are applied.  (bad idea).
    """

    def __init__(self,fgcmConfig,fgcmPars,fgcmStars,fgcmLUT):

        self.fgcmLog = fgcmConfig.fgcmLog

        #self.fgcmLog.log('INFO','Initializing FgcmChisq')
        self.fgcmLog.info('Initializing FgcmChisq')

        # does this need to be shm'd?
        self.fgcmPars = fgcmPars

        # this is shm'd
        self.fgcmLUT = fgcmLUT

        # also shm'd
        self.fgcmStars = fgcmStars

        # need to configure
        self.nCore = fgcmConfig.nCore
        self.ccdStartIndex = fgcmConfig.ccdStartIndex
        self.nStarPerRun = fgcmConfig.nStarPerRun
        self.noChromaticCorrections = fgcmConfig.noChromaticCorrections

        # these are the standard *band* I10s
        self.I10StdBand = fgcmConfig.I10StdBand

        self.illegalValue = fgcmConfig.illegalValue

        if (fgcmConfig.useSedLUT and self.fgcmLUT.hasSedLUT):
            self.useSedLUT = True
        else:
            self.useSedLUT = False

        self.resetFitChisqList()

        # this is the default number of parameters
        self.nActualFitPars = self.fgcmPars.nFitPars
        #self.fgcmLog.log('INFO','Default: fit %d parameters.' % (self.nActualFitPars))
        self.fgcmLog.info('Default: fit %d parameters.' % (self.nActualFitPars))

        self.clearMatchCache()

    def resetFitChisqList(self):
        """
        Reset the recorded list of chi-squared values.
        """

        self.fitChisqs = []

    def clearMatchCache(self):
        """
        Clear the pre-match cache.  Note that this isn't working right.
        """
        self.matchesCached = False
        self.goodObs = None
        self.goodStarsSub = None

    def __call__(self,fitParams,fitterUnits=False,computeDerivatives=False,computeSEDSlopes=False,useMatchCache=False,debug=False,allExposures=False,includeReserve=False,fgcmGray=None):
        """
        Compute the chi-squared for a given set of parameters.

        parameters
        ----------
        fitParams: numpy array of floats
           Array with the numerical values of the parameters (properly formatted).
        fitterUnits: bool, default=False
           Are the units of fitParams normalized for the minimizer?
        computeDerivatives: bool, default=False
           Compute fit derivatives?
        computeSEDSlopes: bool, default=False
           Compute SED slopes from magnitudes?
        useMatchCache: bool, default=False
           Cache observation matches.  Do not use!
        debug: bool, default=False
           Debug mode with no multiprocessing
        allExposures: bool, default=False
           Compute using all exposures, including flagged/non-photometric
        includeReserve: bool, default=False
           Compute using all objects, including those put in reserve.
        fgcmGray: FgcmGray, default=None
           CCD Gray information for computing with "ccd crunch"
        """

        # computeDerivatives: do we want to compute the derivatives?
        # computeSEDSlope: compute SED Slope and recompute mean mags?
        # fitterUnits: units of th fitter or "true" units?

        self.computeDerivatives = computeDerivatives
        self.computeSEDSlopes = computeSEDSlopes
        self.fitterUnits = fitterUnits
        self.allExposures = allExposures
        self.useMatchCache = useMatchCache
        self.includeReserve = includeReserve
        self.fgcmGray = fgcmGray    # may be None

        self.fgcmLog.debug('FgcmChisq: computeDerivatives = %d' %
                         (int(computeDerivatives)))
        self.fgcmLog.debug('FgcmChisq: computeSEDSlopes = %d' %
                         (int(computeSEDSlopes)))
        self.fgcmLog.debug('FgcmChisq: fitterUnits = %d' %
                         (int(fitterUnits)))
        self.fgcmLog.debug('FgcmChisq: allExposures = %d' %
                         (int(allExposures)))
        self.fgcmLog.debug('FgcmChisq: includeReserve = %d' %
                         (int(includeReserve)))

        startTime = time.time()

        if (self.allExposures and (self.computeDerivatives or
                                   self.computeSEDSlopes)):
            raise ValueError("Cannot set allExposures and computeDerivatives or computeSEDSlopes")
        self.fgcmPars.reloadParArray(fitParams,fitterUnits=self.fitterUnits)
        self.fgcmPars.parsToExposures()

        # and reset numbers if necessary
        if (not self.allExposures):
            snmm.getArray(self.fgcmStars.objMagStdMeanHandle)[:] = 99.0
            snmm.getArray(self.fgcmStars.objMagStdMeanNoChromHandle)[:] = 99.0
            snmm.getArray(self.fgcmStars.objMagStdMeanErrHandle)[:] = 99.0

        # do we want to include reserve stars?
        if (self.includeReserve):
            # this mask will filter everything but RESERVED
            resMask = 255 & ~objFlagDict['RESERVED']
            goodStars,=np.where((snmm.getArray(self.fgcmStars.objFlagHandle) & resMask) == 0)
        else:
            goodStars,=np.where(snmm.getArray(self.fgcmStars.objFlagHandle) == 0)

        self.fgcmLog.info('Found %d good stars for chisq' % (goodStars.size))

        if (goodStars.size == 0):
            raise ValueError("No good stars to fit!")

        # do global pre-matching before giving to workers, because
        #  it is faster this way

        obsObjIDIndex = snmm.getArray(self.fgcmStars.obsObjIDIndexHandle)
        obsExpIndex = snmm.getArray(self.fgcmStars.obsExpIndexHandle)
        obsFlag = snmm.getArray(self.fgcmStars.obsFlagHandle)

        if (self.useMatchCache and self.matchesCached) :
            # we have already done the matching
            self.fgcmLog.info('Retrieving cached matches')
            goodObs = self.goodObs
            goodStarsSub = self.goodStarsSub
        else:
            # we need to do matching
            preStartTime=time.time()
            self.fgcmLog.info('Pre-matching stars and observations...')
            goodStarsSub,goodObs = esutil.numpy_util.match(goodStars,
                                                           obsObjIDIndex,
                                                           presorted=True)

            if (goodStarsSub[0] != 0.0):
                raise ValueError("Very strange that the goodStarsSub first element is non-zero.")

            if (not self.allExposures):
                # cut out all bad exposures and bad observations
                gd,=np.where((self.fgcmPars.expFlag[obsExpIndex[goodObs]] == 0) &
                             (obsFlag[goodObs] == 0))
            else:
                # just cut out bad observations
                gd,=np.where(obsFlag[goodObs] == 0)

            # crop out both goodObs and goodStarsSub
            goodObs=goodObs[gd]
            goodStarsSub=goodStarsSub[gd]

            self.fgcmLog.info('Pre-matching done in %.1f sec.' %
                             (time.time() - preStartTime))

            if (self.useMatchCache) :
                self.fgcmLog.info('Caching matches for next iteration')
                self.matchesCached = True
                self.goodObs = goodObs
                self.goodStarsSub = goodStarsSub

        self.nSums = 2  # chisq, nobs
        if (self.computeDerivatives):
            # we have one for each of the derivatives
            # and a duplicate set to track which parameters were "touched"
            self.nSums += 2*self.fgcmPars.nFitPars

        self.debug = debug
        if (self.debug):
            # debug mode: single core
            self.totalHandleDict = {}
            self.totalHandleDict[0] = snmm.createArray(self.nSums,dtype='f8')

            self._worker((goodStars,goodObs))

            partialSums = snmm.getArray(self.totalHandleDict[0])[:]
        else:
            # regular multi-core

            # make a dummy process to discover starting child number
            proc = multiprocessing.Process()
            workerIndex = proc._identity[0]+1
            proc = None

            self.totalHandleDict = {}
            for thisCore in xrange(self.nCore):
                self.totalHandleDict[workerIndex + thisCore] = (
                    snmm.createArray(self.nSums,dtype='f8'))

            # split goodStars into a list of arrays of roughly equal size

            prepStartTime = time.time()
            nSections = goodStars.size // self.nStarPerRun + 1
            goodStarsList = np.array_split(goodStars,nSections)

            # is there a better way of getting all the first elements from the list?
            #  note that we need to skip the first which should be zero (checked above)
            #  see also fgcmBrightObs.py
            # splitValues is the first of the goodStars in each list
            splitValues = np.zeros(nSections-1,dtype='i4')
            for i in xrange(1,nSections):
                splitValues[i-1] = goodStarsList[i][0]

            # get the indices from the goodStarsSub matched list (matched to goodStars)
            splitIndices = np.searchsorted(goodStars[goodStarsSub], splitValues)

            # and split along the indices
            goodObsList = np.split(goodObs,splitIndices)

            workerList = list(zip(goodStarsList,goodObsList))

            # reverse sort so the longest running go first
            workerList.sort(key=lambda elt:elt[1].size, reverse=True)

            self.fgcmLog.info('Using %d sections (%.1f seconds)' %
                             (nSections,time.time()-prepStartTime))

            self.fgcmLog.info('Running chisq on %d cores' % (self.nCore))

            # make a pool
            pool = Pool(processes=self.nCore)
            pool.map(self._worker,workerList,chunksize=1)
            pool.close()
            pool.join()

            # sum up the partial sums from the different jobs
            partialSums = np.zeros(self.nSums,dtype='f8')
            for thisCore in xrange(self.nCore):
                partialSums[:] += snmm.getArray(
                    self.totalHandleDict[workerIndex + thisCore])[:]

        if (not self.allExposures):
            # we get the number of fit parameters by counting which of the parameters
            #  have been touched by the data (number of touches is irrelevant)

            if (self.computeDerivatives):
                nonZero, = np.where(partialSums[self.fgcmPars.nFitPars:
                                                    2*self.fgcmPars.nFitPars] > 0)
                self.nActualFitPars = nonZero.size
                self.fgcmLog.info('Actually fit %d parameters.' % (self.nActualFitPars))

            fitDOF = partialSums[-1] - float(self.nActualFitPars)

            if (fitDOF <= 0):
                raise ValueError("Number of parameters fitted is more than number of constraints! (%d > %d)" % (self.fgcmPars.nFitPars,partialSums[-1]))

            fitChisq = partialSums[-2] / fitDOF
            if (self.computeDerivatives):
                dChisqdP = partialSums[0:self.fgcmPars.nFitPars] / fitDOF

            # want to append this...
            self.fitChisqs.append(fitChisq)

            self.fgcmLog.info('Chisq/dof = %.2f (%d iterations)' %
                             (fitChisq, len(self.fitChisqs)))

        else:
            try:
                fitChisq = self.fitChisqs[-1]
            except:
                fitChisq = 0.0

        # free shared arrays
        for key in self.totalHandleDict.keys():
            snmm.freeArray(self.totalHandleDict[key])

        self.fgcmLog.info('Chisq computation took %.2f seconds.' %
                         (time.time() - startTime))

        self.fgcmStars.magStdComputed = True
        if (self.allExposures):
            self.fgcmStars.allMagStdComputed = True

        if (self.computeDerivatives):
            return fitChisq, dChisqdP
        else:
            return fitChisq

    def _worker(self,goodStarsAndObs):
        """
        Multiprocessing worker for FgcmChisq.  Not to be called on its own.

        parameters
        ----------
        goodStarsAndObs: tuple[2]
           (goodStars, goodObs)
        """

        # NOTE: No logging is allowed in the _worker method

        workerStartTime = time.time()

        goodStars = goodStarsAndObs[0]
        goodObs = goodStarsAndObs[1]

        if self.debug:
            thisCore = 0
        else:
            thisCore = multiprocessing.current_process()._identity[0]

        objMagStdMean = snmm.getArray(self.fgcmStars.objMagStdMeanHandle)
        objMagStdMeanNoChrom = snmm.getArray(self.fgcmStars.objMagStdMeanNoChromHandle)
        objMagStdMeanErr = snmm.getArray(self.fgcmStars.objMagStdMeanErrHandle)
        objSEDSlope = snmm.getArray(self.fgcmStars.objSEDSlopeHandle)
        objNGoodObs = snmm.getArray(self.fgcmStars.objNGoodObsHandle)

        obsObjIDIndex = snmm.getArray(self.fgcmStars.obsObjIDIndexHandle)

        obsExpIndex = snmm.getArray(self.fgcmStars.obsExpIndexHandle)
        obsBandIndex = snmm.getArray(self.fgcmStars.obsBandIndexHandle)
        obsLUTFilterIndex = snmm.getArray(self.fgcmStars.obsLUTFilterIndexHandle)
        obsCCDIndex = snmm.getArray(self.fgcmStars.obsCCDHandle) - self.ccdStartIndex
        obsFlag = snmm.getArray(self.fgcmStars.obsFlagHandle)
        obsSecZenith = snmm.getArray(self.fgcmStars.obsSecZenithHandle)
        obsMagADU = snmm.getArray(self.fgcmStars.obsMagADUHandle)
        # obsMagADUErr = snmm.getArray(self.fgcmStars.obsMagADUErrHandle)
        obsMagADUModelErr = snmm.getArray(self.fgcmStars.obsMagADUModelErrHandle)
        obsMagStd = snmm.getArray(self.fgcmStars.obsMagStdHandle)

        # and fgcmGray stuff (if desired)
        if (self.fgcmGray is not None):
            ccdGray = snmm.getArray(self.fgcmGray.ccdGrayHandle)
            # this is ccdGray[expIndex, ccdIndex]
            # and we only apply when > self.illegalValue
            # same sign as FGCM_DUST (QESys)

        # and the arrays for locking access
        objMagStdMeanLock = snmm.getArrayBase(self.fgcmStars.objMagStdMeanHandle).get_lock()
        obsMagStdLock = snmm.getArrayBase(self.fgcmStars.obsMagStdHandle).get_lock()

        # cut these down now, faster later
        obsObjIDIndexGO = obsObjIDIndex[goodObs]
        obsBandIndexGO = obsBandIndex[goodObs]
        obsLUTFilterIndexGO = obsLUTFilterIndex[goodObs]
        obsExpIndexGO = obsExpIndex[goodObs]
        obsSecZenithGO = obsSecZenith[goodObs]
        obsCCDIndexGO = obsCCDIndex[goodObs]

        # which observations are used in the fit?
        _,obsFitUseGO = esutil.numpy_util.match(self.fgcmPars.fitBandIndex,
                                                obsBandIndexGO)

        # now refer to obsBandIndex[goodObs]
        # add GO to index names that are cut to goodObs
        # add GOF to index names that are cut to goodObs[obsFitUseGO]

        lutIndicesGO = self.fgcmLUT.getIndices(obsLUTFilterIndexGO,
                                               self.fgcmPars.expPWV[obsExpIndexGO],
                                               self.fgcmPars.expO3[obsExpIndexGO],
                                               #np.log(self.fgcmPars.expTau[obsExpIndexGO]),
                                               self.fgcmPars.expLnTau[obsExpIndexGO],
                                               self.fgcmPars.expAlpha[obsExpIndexGO],
                                               obsSecZenithGO,
                                               obsCCDIndexGO,
                                               self.fgcmPars.expPmb[obsExpIndexGO])
        I0GO = self.fgcmLUT.computeI0(self.fgcmPars.expPWV[obsExpIndexGO],
                                      self.fgcmPars.expO3[obsExpIndexGO],
                                      #np.log(self.fgcmPars.expTau[obsExpIndexGO]),
                                      self.fgcmPars.expLnTau[obsExpIndexGO],
                                      self.fgcmPars.expAlpha[obsExpIndexGO],
                                      obsSecZenithGO,
                                      self.fgcmPars.expPmb[obsExpIndexGO],
                                      lutIndicesGO)
        I10GO = self.fgcmLUT.computeI1(self.fgcmPars.expPWV[obsExpIndexGO],
                                       self.fgcmPars.expO3[obsExpIndexGO],
                                       #np.log(self.fgcmPars.expTau[obsExpIndexGO]),
                                       self.fgcmPars.expLnTau[obsExpIndexGO],
                                       self.fgcmPars.expAlpha[obsExpIndexGO],
                                       obsSecZenithGO,
                                       self.fgcmPars.expPmb[obsExpIndexGO],
                                       lutIndicesGO) / I0GO

        qeSysGO = self.fgcmPars.expQESys[obsExpIndexGO]

        obsMagGO = obsMagADU[goodObs] + 2.5*np.log10(I0GO) + qeSysGO

        if (self.fgcmGray is not None):
            # We want to apply the "CCD Gray Crunch"
            # make sure we aren't adding something crazy, but this shouldn't happen
            # because we're filtering good observations (I hope!)
            ok,=np.where(ccdGray[obsExpIndexGO, obsCCDIndexGO] > self.illegalValue)
            obsMagGO[ok] += ccdGray[obsExpIndexGO[ok], obsCCDIndexGO[ok]]

        # Compute the sub-selected error-squared, using model error when available
        obsMagErr2GO = obsMagADUModelErr[goodObs]**2.

        if (self.computeSEDSlopes):
            # first, compute mean mags (code same as below.  FIXME: consolidate, but how?)

            # make temp vars.  With memory overhead

            wtSum = np.zeros_like(objMagStdMean,dtype='f8')
            objMagStdMeanTemp = np.zeros_like(objMagStdMean)

            np.add.at(wtSum,
                      (obsObjIDIndexGO,obsBandIndexGO),
                      1./obsMagErr2GO)
            np.add.at(objMagStdMeanTemp,
                  (obsObjIDIndexGO,obsBandIndexGO),
                  obsMagGO/obsMagErr2GO)

            # these are good object/bands that were observed
            gd=np.where(wtSum > 0.0)

            # and acquire lock to save the values
            objMagStdMeanLock.acquire()

            objMagStdMean[gd] = objMagStdMeanTemp[gd] / wtSum[gd]
            objMagStdMeanErr[gd] = np.sqrt(1./wtSum[gd])

            # and release the lock.
            objMagStdMeanLock.release()

            if (self.useSedLUT):
                self.fgcmStars.computeObjectSEDSlopesLUT(goodStars,self.fgcmLUT)
            else:
                self.fgcmStars.computeObjectSEDSlopes(goodStars)

        # compute linearized chromatic correction
        deltaStdGO = 2.5 * np.log10((1.0 +
                                   objSEDSlope[obsObjIDIndexGO,
                                               obsBandIndexGO] * I10GO) /
                                  (1.0 + objSEDSlope[obsObjIDIndexGO,
                                                     obsBandIndexGO] *
                                   self.I10StdBand[obsBandIndexGO]))

        if self.noChromaticCorrections:
            # NOT RECOMMENDED
            deltaStdGO *= 0.0

        # we can only do this for calibration stars.
        #  must reference the full array to save

        # acquire lock when we write to and retrieve from full array
        obsMagStdLock.acquire()

        obsMagStd[goodObs] = obsMagGO + deltaStdGO
        # this is cut here
        obsMagStdGO = obsMagStd[goodObs]

        # we now have a local cut copy, so release
        obsMagStdLock.release()

        # kick out if we're just computing magstd for all exposures
        if (self.allExposures) :
            # kick out
            return None

        # compute mean mags

        # we make temporary variables.  These are less than ideal because they
        #  take up the full memory footprint.  MAYBE look at making a smaller
        #  array just for the stars under consideration, but this would make the
        #  indexing in the np.add.at() more difficult

        wtSum = np.zeros_like(objMagStdMean,dtype='f8')
        objMagStdMeanTemp = np.zeros_like(objMagStdMean)
        objMagStdMeanNoChromTemp = np.zeros_like(objMagStdMeanNoChrom)

        np.add.at(wtSum,
                  (obsObjIDIndexGO,obsBandIndexGO),
                  1./obsMagErr2GO)

        np.add.at(objMagStdMeanTemp,
                  (obsObjIDIndexGO,obsBandIndexGO),
                  obsMagStdGO/obsMagErr2GO)

        # And the same thing with the non-chromatic corrected values
        np.add.at(objMagStdMeanNoChromTemp,
                  (obsObjIDIndexGO,obsBandIndexGO),
                  obsMagGO/obsMagErr2GO)

        # which objects/bands have observations?
        gd=np.where(wtSum > 0.0)

        # and acquire lock to save the values
        objMagStdMeanLock.acquire()

        objMagStdMean[gd] = objMagStdMeanTemp[gd] / wtSum[gd]
        objMagStdMeanNoChrom[gd] = objMagStdMeanNoChromTemp[gd] / wtSum[gd]
        objMagStdMeanErr[gd] = np.sqrt(1./wtSum[gd])

        # also make local copies for Good Observations
        objMagStdMeanGO = objMagStdMean[obsObjIDIndexGO,obsBandIndexGO]
        objMagStdMeanErr2GO = objMagStdMeanErr[obsObjIDIndexGO,obsBandIndexGO]**2.

        # and release the lock.
        objMagStdMeanLock.release()

        # compute delta-mags

        deltaMagGO = (obsMagStdGO - objMagStdMeanGO)

        # Note that this is the model error when we have it
        obsWeightGO = 1. / obsMagErr2GO

        deltaMagWeightedGOF = deltaMagGO[obsFitUseGO] * obsWeightGO[obsFitUseGO]

        partialChisq = np.sum(deltaMagGO[obsFitUseGO]**2. * obsWeightGO[obsFitUseGO])

        partialArray = np.zeros(self.nSums,dtype='f8')
        partialArray[-2] = partialChisq
        partialArray[-1] = obsFitUseGO.size

        if (self.computeDerivatives):
            unitDict=self.fgcmPars.getUnitDict(fitterUnits=self.fitterUnits)

            # this is going to be ugly.  wow, how many indices and sub-indices?
            #  or does it simplify since we need all the obs on a night?
            #  we shall see!  And speed up!

            (dLdPWVGO,dLdO3GO,dLdTauGO,dLdAlphaGO) = (
                self.fgcmLUT.computeLogDerivatives(lutIndicesGO,
                                                   I0GO))

            if (self.fgcmLUT.hasI1Derivatives):
                (dLdPWVI1GO,dLdO3I1GO,dLdTauI1GO,dLdAlphaI1GO) = (
                    self.fgcmLUT.computeLogDerivativesI1(lutIndicesGO,
                                                         I0GO,
                                                         I10GO,
                                                         objSEDSlope[obsObjIDIndexGO,
                                                                     obsBandIndexGO]))
                dLdPWVGO += dLdPWVI1GO
                dLdO3GO += dLdO3I1GO
                dLdTauGO += dLdTauI1GO
                dLdAlphaGO += dLdAlphaI1GO

            # we have objMagStdMeanErr[objIndex,:] = \Sum_{i"} 1/\sigma^2_{i"j}
            #   note that this is summed over all observations of an object in a band
            #   so that this is already done

            # we need magdLdp = \Sum_{i'} (1/\sigma^2_{i'j}) dL(i',j|p)
            #   note that this is summed over all observations in a filter that
            #   touch a given parameter

            # set up arrays
            magdLdPWVIntercept = np.zeros((self.fgcmPars.nCampaignNights,
                                           self.fgcmPars.nFitBands))
            magdLdPWVPerSlope = np.zeros_like(magdLdPWVIntercept)
            magdLdPWVOffset = np.zeros_like(magdLdPWVIntercept)
            magdLdTauIntercept = np.zeros_like(magdLdPWVIntercept)
            magdLdTauPerSlope = np.zeros_like(magdLdPWVIntercept)
            magdLdTauOffset = np.zeros_like(magdLdPWVIntercept)
            magdLdAlpha = np.zeros_like(magdLdPWVIntercept)
            magdLdO3 = np.zeros_like(magdLdPWVIntercept)

            magdLdPWVScale = np.zeros(self.fgcmPars.nFitBands,dtype='f4')
            magdLdTauScale = np.zeros_like(magdLdPWVScale)

            magdLdPWVRetrievedScale = np.zeros(self.fgcmPars.nFitBands,dtype='f4')
            magdLdPWVRetrievedOffset = np.zeros_like(magdLdPWVRetrievedScale)
            magdLdPWVRetrievedNightlyOffset = np.zeros_like(magdLdPWVIntercept)

            magdLdWashIntercept = np.zeros((self.fgcmPars.nWashIntervals,
                                            self.fgcmPars.nFitBands))
            magdLdWashSlope = np.zeros_like(magdLdWashIntercept)

            # note below that objMagStdMeanErr2GO is the the square of the error,
            #  and already cut to [obsObjIDIndexGO,obsBandIndexGO]

            ##########
            ## O3
            ##########

            expNightIndexGOF = self.fgcmPars.expNightIndex[obsExpIndexGO[obsFitUseGO]]
            uNightIndex = np.unique(expNightIndexGOF)

            np.add.at(magdLdO3,
                      (expNightIndexGOF,obsBandIndexGO[obsFitUseGO]),
                      dLdO3GO[obsFitUseGO] / obsMagErr2GO[obsFitUseGO])
            np.multiply.at(magdLdO3,
                           (expNightIndexGOF,obsBandIndexGO[obsFitUseGO]),
                           objMagStdMeanErr2GO[obsFitUseGO])
            np.add.at(partialArray[self.fgcmPars.parO3Loc:
                                       (self.fgcmPars.parO3Loc+
                                        self.fgcmPars.nCampaignNights)],
                      expNightIndexGOF,
                      deltaMagWeightedGOF * (
                    (dLdO3GO[obsFitUseGO] -
                     magdLdO3[expNightIndexGOF,obsBandIndexGO[obsFitUseGO]])))

            partialArray[self.fgcmPars.parO3Loc +
                         uNightIndex] *= (2.0 / unitDict['o3Unit'])
            partialArray[self.fgcmPars.nFitPars +
                         self.fgcmPars.parO3Loc +
                         uNightIndex] += 1

            ###########
            ## Alpha
            ###########

            np.add.at(magdLdAlpha,
                      (expNightIndexGOF,obsBandIndexGO[obsFitUseGO]),
                      dLdAlphaGO[obsFitUseGO] / obsMagErr2GO[obsFitUseGO])
            np.multiply.at(magdLdAlpha,
                           (expNightIndexGOF,obsBandIndexGO[obsFitUseGO]),
                           objMagStdMeanErr2GO[obsFitUseGO])
            np.add.at(partialArray[self.fgcmPars.parAlphaLoc:
                                       (self.fgcmPars.parAlphaLoc+
                                        self.fgcmPars.nCampaignNights)],
                      expNightIndexGOF,
                      deltaMagWeightedGOF * (
                    (dLdAlphaGO[obsFitUseGO] -
                     magdLdAlpha[expNightIndexGOF,obsBandIndexGO[obsFitUseGO]])))

            partialArray[self.fgcmPars.parAlphaLoc +
                         uNightIndex] *= (2.0 / unitDict['alphaUnit'])
            partialArray[self.fgcmPars.nFitPars +
                         self.fgcmPars.parAlphaLoc +
                         uNightIndex] += 1

            ###########
            ## PWV External
            ###########

            if (self.fgcmPars.hasExternalPWV and not self.fgcmPars.useRetrievedPWV):
                hasExtGOF,=np.where(self.fgcmPars.externalPWVFlag[obsExpIndexGO[obsFitUseGO]])
                uNightIndexHasExt = np.unique(expNightIndexGOF[hasExtGOF])

                # PWV Nightly Offset
                np.add.at(magdLdPWVOffset,
                          (expNightIndexGOF[hasExtGOF],
                           obsBandIndexGO[obsFitUseGO[hasExtGOF]]),
                          dLdPWVGO[obsFitUseGO[hasExtGOF]] /
                          obsMagErr2GO[obsFitUseGO[hasExtGOF]])
                np.multiply.at(magdLdPWVOffset,
                               (expNightIndexGOF[hasExtGOF],
                                obsBandIndexGO[obsFitUseGO[hasExtGOF]]),
                               objMagStdMeanErr2GO[obsFitUseGO[hasExtGOF]])
                np.add.at(partialArray[self.fgcmPars.parExternalPWVOffsetLoc:
                                           (self.fgcmPars.parExternalPWVOffsetLoc+
                                            self.fgcmPars.nCampaignNights)],
                          expNightIndexGOF[hasExtGOF],
                          deltaMagWeightedGOF[hasExtGOF] * (
                        (dLdPWVGO[obsFitUseGO[hasExtGOF]] -
                         magdLdPWVOffset[expNightIndexGOF[hasExtGOF],
                                         obsBandIndexGO[obsFitUseGO[hasExtGOF]]])))
                partialArray[self.fgcmPars.parExternalPWVOffsetLoc +
                             uNightIndexHasExt] *= (2.0 / unitDict['pwvUnit'])
                partialArray[self.fgcmPars.nFitPars +
                             self.fgcmPars.parExternalPWVOffsetLoc +
                             uNightIndexHasExt] += 1

                # PWV Global Scale
                np.add.at(magdLdPWVScale,
                          obsBandIndexGO[obsFitUseGO[hasExtGOF]],
                          self.fgcmPars.expPWV[obsExpIndexGO[obsFitUseGO[hasExtGOF]]] *
                          dLdPWVGO[obsFitUseGO[hasExtGOF]] /
                          obsMagErr2GO[obsFitUseGO[hasExtGOF]])
                np.multiply.at(magdLdPWVScale,
                               obsBandIndexGO[obsFitUseGO[hasExtGOF]],
                               objMagStdMeanErr2GO[obsFitUseGO[hasExtGOF]])
                partialArray[self.fgcmPars.parExternalPWVScaleLoc] = 2.0 * (
                    np.sum(deltaMagWeightedGOF[hasExtGOF] * (
                            self.fgcmPars.expPWV[obsExpIndexGO[obsFitUseGO[hasExtGOF]]] *
                            dLdPWVGO[obsFitUseGO[hasExtGOF]] -
                            magdLdPWVScale[obsBandIndexGO[obsFitUseGO[hasExtGOF]]])) /
                    unitDict['pwvGlobalUnit'])
                partialArray[self.fgcmPars.nFitPars +
                             self.fgcmPars.parExternalPWVScaleLoc] += 1

            ################
            ## PWV Retrieved
            ################

            if (self.fgcmPars.useRetrievedPWV):
                hasRetrievedPWVGOF, = np.where((self.fgcmPars.compRetrievedPWVFlag[obsExpIndexGO[obsFitUseGO]] &
                                                retrievalFlagDict['EXPOSURE_RETRIEVED']) > 0)

                if hasRetrievedPWVGOF.size > 0:
                    # note this might be zero-size on first run

                    # PWV Retrieved Global Scale
                    np.add.at(magdLdPWVRetrievedScale,
                              obsBandIndexGO[obsFitUseGO[hasRetrievedPWVGOF]],
                              self.fgcmPars.expPWV[obsExpIndexGO[obsFitUseGO[hasRetrievedPWVGOF]]] *
                              dLdPWVGO[obsFitUseGO[hasRetrievedPWVGOF]] /
                              obsMagErr2GO[obsFitUseGO[hasRetrievedPWVGOF]])
                    np.multiply.at(magdLdPWVRetrievedScale,
                                   obsBandIndexGO[obsFitUseGO[hasRetrievedPWVGOF]],
                                   objMagStdMeanErr2GO[obsFitUseGO[hasRetrievedPWVGOF]])
                    partialArray[self.fgcmPars.parRetrievedPWVScaleLoc] = 2.0 * (
                        np.sum(deltaMagWeightedGOF[hasRetrievedPWVGOF] * (
                                self.fgcmPars.expPWV[obsExpIndexGO[obsFitUseGO[hasRetrievedPWVGOF]]] *
                                dLdPWVGO[obsFitUseGO[hasRetrievedPWVGOF]] -
                                magdLdPWVRetrievedScale[obsBandIndexGO[obsFitUseGO[hasRetrievedPWVGOF]]])) /
                        unitDict['pwvGlobalUnit'])
                    partialArray[self.fgcmPars.nFitPars +
                                 self.fgcmPars.parRetrievedPWVScaleLoc] += 1

                    if self.fgcmPars.useNightlyRetrievedPWV:
                        # PWV Retrieved Nightly Offset

                        uNightIndexHasRetrievedPWV = np.unique(expNightIndexGOF[hasRetrievedPWVGOF])

                        np.add.at(magdLdPWVRetrievedNightlyOffset,
                                  (expNightIndexGOF[hasRetrievedPWVGOF],
                                   obsBandIndexGO[obsFitUseGO[hasRetrievedPWVGOF]]),
                                  dLdPWVGO[obsFitUseGO[hasRetrievedPWVGOF]] /
                                  obsMagErr2GO[obsFitUseGO[hasRetrievedPWVGOF]])
                        np.multiply.at(magdLdPWVRetrievedNightlyOffset,
                                       (expNightIndexGOF[hasRetrievedPWVGOF],
                                        obsBandIndexGO[obsFitUseGO[hasRetrievedPWVGOF]]),
                                       objMagStdMeanErr2GO[obsFitUseGO[hasRetrievedPWVGOF]])
                        np.add.at(partialArray[self.fgcmPars.parRetrievedPWVNightlyOffsetLoc:
                                                   (self.fgcmPars.parRetrievedPWVNightlyOffsetLoc+
                                                    self.fgcmPars.nCampaignNights)],
                                  expNightIndexGOF[hasRetrievedPWVGOF],
                                  deltaMagWeightedGOF[hasRetrievedPWVGOF] * (
                                (dLdPWVGO[obsFitUseGO[hasRetrievedPWVGOF]] -
                                 magdLdPWVRetrievedNightlyOffset[expNightIndexGOF[hasRetrievedPWVGOF],
                                                                 obsBandIndexGO[obsFitUseGO[hasRetrievedPWVGOF]]])))
                        partialArray[self.fgcmPars.parRetrievedPWVNightlyOffsetLoc +
                                     uNightIndexHasRetrievedPWV] *= (2.0 / unitDict['pwvUnit'])
                        partialArray[self.fgcmPars.nFitPars +
                                     self.fgcmPars.parRetrievedPWVNightlyOffsetLoc +
                                     uNightIndexHasRetrievedPWV] += 1

                    else:
                        # PWV Retrieved Global Offset
                        np.add.at(magdLdPWVRetrievedOffset,
                                  obsBandIndexGO[obsFitUseGO[hasRetrievedPWVGOF]],
                                  dLdPWVGO[obsFitUseGO[hasRetrievedPWVGOF]] /
                                  obsMagErr2GO[obsFitUseGO[hasRetrievedPWVGOF]])
                        np.multiply.at(magdLdPWVRetrievedOffset,
                                       obsBandIndexGO[obsFitUseGO[hasRetrievedPWVGOF]],
                                       objMagStdMeanErr2GO[obsFitUseGO[hasRetrievedPWVGOF]])
                        partialArray[self.fgcmPars.parRetrievedPWVOffsetLoc] = 2.0 * (
                            np.sum(deltaMagWeightedGOF[hasRetrievedPWVGOF] * (
                                    dLdPWVGO[obsFitUseGO[hasRetrievedPWVGOF]] -
                                    magdLdPWVRetrievedOffset[obsBandIndexGO[obsFitUseGO[hasRetrievedPWVGOF]]])) /
                            unitDict['pwvGlobalUnit'])
                        partialArray[self.fgcmPars.nFitPars +
                                     self.fgcmPars.parRetrievedPWVOffsetLoc] += 1

            else:
                ###########
                ## PWV No External
                ###########

                noExtGOF, = np.where(~self.fgcmPars.externalPWVFlag[obsExpIndexGO[obsFitUseGO]])
                uNightIndexNoExt = np.unique(expNightIndexGOF[noExtGOF])

                # PWV Nightly Intercept

                np.add.at(magdLdPWVIntercept,
                          (expNightIndexGOF[noExtGOF],
                           obsBandIndexGO[obsFitUseGO[noExtGOF]]),
                          dLdPWVGO[obsFitUseGO[noExtGOF]] /
                          obsMagErr2GO[obsFitUseGO[noExtGOF]])
                np.multiply.at(magdLdPWVIntercept,
                               (expNightIndexGOF[noExtGOF],
                                obsBandIndexGO[obsFitUseGO[noExtGOF]]),
                               objMagStdMeanErr2GO[obsFitUseGO[noExtGOF]])
                np.add.at(partialArray[self.fgcmPars.parPWVInterceptLoc:
                                           (self.fgcmPars.parPWVInterceptLoc+
                                            self.fgcmPars.nCampaignNights)],
                          expNightIndexGOF[noExtGOF],
                          deltaMagWeightedGOF[noExtGOF] * (
                        (dLdPWVGO[obsFitUseGO[noExtGOF]] -
                         magdLdPWVOffset[expNightIndexGOF[noExtGOF],
                                         obsBandIndexGO[obsFitUseGO[noExtGOF]]])))

                partialArray[self.fgcmPars.parPWVInterceptLoc +
                             uNightIndexNoExt] *= (2.0 / unitDict['pwvUnit'])
                partialArray[self.fgcmPars.nFitPars +
                             self.fgcmPars.parPWVInterceptLoc +
                             uNightIndexNoExt] += 1

                # PWV Nightly Percent Slope
                np.add.at(magdLdPWVPerSlope,
                          (expNightIndexGOF[noExtGOF],
                           obsBandIndexGO[obsFitUseGO[noExtGOF]]),
                          self.fgcmPars.expDeltaUT[obsExpIndexGO[obsFitUseGO[noExtGOF]]] *
                          self.fgcmPars.expPWV[obsExpIndexGO[obsFitUseGO[noExtGOF]]] *
                          dLdPWVGO[obsFitUseGO[noExtGOF]] /
                          obsMagErr2GO[obsFitUseGO[noExtGOF]])
                np.multiply.at(magdLdPWVPerSlope,
                               (expNightIndexGOF[noExtGOF],
                                obsBandIndexGO[obsFitUseGO[noExtGOF]]),
                               objMagStdMeanErr2GO[obsFitUseGO[noExtGOF]])
                np.add.at(partialArray[self.fgcmPars.parPWVPerSlopeLoc:
                                           (self.fgcmPars.parPWVPerSlopeLoc+
                                            self.fgcmPars.nCampaignNights)],
                          expNightIndexGOF[noExtGOF],
                          deltaMagWeightedGOF[noExtGOF] * (
                        (self.fgcmPars.expDeltaUT[obsExpIndexGO[obsFitUseGO[noExtGOF]]] *
                         self.fgcmPars.expPWV[obsExpIndexGO[obsFitUseGO[noExtGOF]]] *
                         dLdPWVGO[obsFitUseGO[noExtGOF]] -
                         magdLdPWVPerSlope[expNightIndexGOF[noExtGOF],
                                           obsBandIndexGO[obsFitUseGO[noExtGOF]]])))

                partialArray[self.fgcmPars.parPWVPerSlopeLoc +
                             uNightIndex] *= (2.0 / unitDict['pwvPerSlopeUnit'])
                partialArray[self.fgcmPars.nFitPars +
                             self.fgcmPars.parPWVPerSlopeLoc] += 1

            #############
            ## Tau External
            #############

            if (self.fgcmPars.hasExternalTau):
                hasExtGOF,=np.where(self.fgcmPars.externalTauFlag[obsExpIndexGO[obsFitUseGO]])
                uNightIndexHasExt = np.unique(expNightIndexGOF[hasExtGOF])

                # Tau Nightly Offset
                np.add.at(magdLdTauOffset,
                          (expNightIndexGOF[hasExtGOF],
                           obsBandIndexGO[obsFitUseGO[hasExtGOF]]),
                          dLdTauGO[obsFitUseGO[hasExtGOF]] /
                          obsMagErr2GO[obsFitUseGO[hasExtGOF]])
                np.multiply.at(magdLdTauOffset,
                               (expNightIndexGOF[hasExtGOF],
                                obsBandIndexGO[obsFitUseGO[hasExtGOF]]),
                               objMagStdMeanErr2GO[obsFitUseGO[hasExtGOF]])
                np.add.at(partialArray[self.fgcmPars.parExternalTauOffsetLoc:
                                           (self.fgcmPars.parExternalTauOffsetLoc+
                                            self.fgcmPars.nCampaignNights)],
                          expNightIndexGOF[hasExtGOF],
                          deltaMagWeightedGOF[hasExtGOF] * (
                        (dLdTauGO[obsFitUseGO[hasExtGOF]] -
                         magdLdTauOffset[expNightIndexGOF[hasExtGOF],
                                         obsBandIndexGO[obsFitUseGO[hasExtGOF]]])))

                partialArray[self.fgcmPars.parExternalTauOffsetLoc +
                             uNightIndexHasExt] *= (2.0 / unitDict['tauUnit'])
                partialArray[self.fgcmPars.nFitPars +
                             self.fgcmPars.parExternalTauOffsetLoc +
                             uNightIndexHasExt] += 1

                # Tau Global Scale
                ## MAYBE: is this correct with the logs?
                np.add.at(magdLdTauScale,
                          obsBandIndexGO[obsFitUseGO[hasExtGOF]],
                          self.fgcmPars.expTau[obsExpIndexGO[obsFitUseGO[hasExtGOF]]] *
                          dLdTauGO[obsFitUseGO[hasExtGOF]] /
                          obsMagErr2GO[obsFitUseGO[hasExtGOF]])
                np.multiply.at(magdLdTauScale,
                               obsBandIndexGO[obsFitUseGO[hasExtGOF]],
                               objMagStdMeanErr2GO[obsFitUseGO[hasExtGOF]])
                partialArray[self.fgcmPars.parExternalTauScaleLoc] = 2.0 * (
                    np.sum(deltaMagWeightedGOF[hasExtGOF] * (
                            self.fgcmPars.expTau[obsExpIndexGO[obsFitUseGO[hasExtGOF]]] *
                            dLdTauGO[obsFitUseGO[hasExtGOF]] -
                            magdLdPWVScale[obsBandIndexGO[obsFitUseGO[hasExtGOF]]])) /
                    unitDict['tauUnit'])
                partialArray[self.fgcmPars.nFitPars +
                             self.fgcmPars.parExternalTauScaleLoc] += 1

            ###########
            ## Tau No External
            ###########

            noExtGOF, = np.where(~self.fgcmPars.externalTauFlag[obsExpIndexGO[obsFitUseGO]])
            uNightIndexNoExt = np.unique(expNightIndexGOF[noExtGOF])

            # lnTau Nightly Intercept
            np.add.at(magdLdTauIntercept,
                      (expNightIndexGOF[noExtGOF],
                       obsBandIndexGO[obsFitUseGO[noExtGOF]]),
                      dLdTauGO[obsFitUseGO[noExtGOF]] /
                      obsMagErr2GO[obsFitUseGO[noExtGOF]])
            np.multiply.at(magdLdTauIntercept,
                           (expNightIndexGOF[noExtGOF],
                            obsBandIndexGO[obsFitUseGO[noExtGOF]]),
                           objMagStdMeanErr2GO[obsFitUseGO[noExtGOF]])
            np.add.at(partialArray[self.fgcmPars.parLnTauInterceptLoc:
                                       (self.fgcmPars.parLnTauInterceptLoc+
                                        self.fgcmPars.nCampaignNights)],
                      expNightIndexGOF[noExtGOF],
                      deltaMagWeightedGOF[noExtGOF] * (
                    (dLdTauGO[obsFitUseGO[noExtGOF]] -
                     magdLdTauOffset[expNightIndexGOF[noExtGOF],
                                     obsBandIndexGO[obsFitUseGO[noExtGOF]]])))

            partialArray[self.fgcmPars.parLnTauInterceptLoc +
                         uNightIndexNoExt] *= (2.0 / unitDict['lnTauUnit'])
            partialArray[self.fgcmPars.nFitPars +
                         self.fgcmPars.parLnTauInterceptLoc +
                         uNightIndexNoExt] += 1

            # lnTau nightly slope
            np.add.at(magdLdTauPerSlope,
                      (expNightIndexGOF[noExtGOF],
                       obsBandIndexGO[obsFitUseGO[noExtGOF]]),
                      self.fgcmPars.expDeltaUT[obsExpIndexGO[obsFitUseGO[noExtGOF]]] *
                      #self.fgcmPars.expTau[obsExpIndexGO[obsFitUseGO[noExtGOF]]] *
                      dLdTauGO[obsFitUseGO[noExtGOF]] /
                      obsMagErr2GO[obsFitUseGO[noExtGOF]])
            np.multiply.at(magdLdTauPerSlope,
                           (expNightIndexGOF[noExtGOF],
                            obsBandIndexGO[obsFitUseGO[noExtGOF]]),
                           objMagStdMeanErr2GO[obsFitUseGO[noExtGOF]])
            np.add.at(partialArray[self.fgcmPars.parLnTauSlopeLoc:
                                       (self.fgcmPars.parLnTauSlopeLoc+
                                        self.fgcmPars.nCampaignNights)],
                      expNightIndexGOF[noExtGOF],
                      deltaMagWeightedGOF[noExtGOF] * (
                    (self.fgcmPars.expDeltaUT[obsExpIndexGO[obsFitUseGO[noExtGOF]]] *
                     #self.fgcmPars.expTau[obsExpIndexGO[obsFitUseGO[noExtGOF]]] *
                     dLdTauGO[obsFitUseGO[noExtGOF]] -
                     magdLdTauPerSlope[expNightIndexGOF[noExtGOF],
                                       obsBandIndexGO[obsFitUseGO[noExtGOF]]])))

            partialArray[self.fgcmPars.parLnTauSlopeLoc +
                         uNightIndexNoExt] *= (2.0 / unitDict['lnTauSlopeUnit'])
            partialArray[self.fgcmPars.nFitPars +
                         self.fgcmPars.parLnTauSlopeLoc +
                         uNightIndexNoExt] += 1

            #############
            ## Washes (QE Sys)
            #############

            expWashIndexGOF = self.fgcmPars.expWashIndex[obsExpIndexGO[obsFitUseGO]]
            uWashIndex = np.unique(expWashIndexGOF)

            # Wash Intercept
            np.add.at(magdLdWashIntercept,
                      (expWashIndexGOF,obsBandIndexGO[obsFitUseGO]),
                      1./obsMagErr2GO[obsFitUseGO])
            np.multiply.at(magdLdWashIntercept,
                           (expWashIndexGOF,obsBandIndexGO[obsFitUseGO]),
                           objMagStdMeanErr2GO[obsFitUseGO])
            np.add.at(partialArray[self.fgcmPars.parQESysInterceptLoc:
                                       (self.fgcmPars.parQESysInterceptLoc +
                                        self.fgcmPars.nWashIntervals)],
                      expWashIndexGOF,
                      deltaMagWeightedGOF * (
                    (1.0 - magdLdWashIntercept[expWashIndexGOF,
                                               obsBandIndexGO[obsFitUseGO]])))

            partialArray[self.fgcmPars.parQESysInterceptLoc +
                         uWashIndex] *= (2.0 / unitDict['qeSysUnit'])
            partialArray[self.fgcmPars.nFitPars +
                         self.fgcmPars.parQESysInterceptLoc +
                         uWashIndex] += 1

            # Wash Slope
            np.add.at(magdLdWashSlope,
                      (expWashIndexGOF,obsBandIndexGO[obsFitUseGO]),
                      (self.fgcmPars.expMJD[obsExpIndexGO[obsFitUseGO]] -
                       self.fgcmPars.washMJDs[expWashIndexGOF]) /
                       obsMagErr2GO[obsFitUseGO])
            np.multiply.at(magdLdWashSlope,
                           (expWashIndexGOF,obsBandIndexGO[obsFitUseGO]),
                           objMagStdMeanErr2GO[obsFitUseGO])
            np.add.at(partialArray[self.fgcmPars.parQESysSlopeLoc:
                                       (self.fgcmPars.parQESysSlopeLoc +
                                        self.fgcmPars.nWashIntervals)],
                      expWashIndexGOF,
                      deltaMagWeightedGOF * (
                    (self.fgcmPars.expMJD[obsExpIndexGO[obsFitUseGO]] -
                     self.fgcmPars.washMJDs[expWashIndexGOF]) -
                    magdLdWashSlope[expWashIndexGOF,
                                    obsBandIndexGO[obsFitUseGO]]))
            partialArray[self.fgcmPars.parQESysSlopeLoc +
                         uWashIndex] *= (2.0 / unitDict['qeSysSlopeUnit'])
            partialArray[self.fgcmPars.nFitPars +
                         self.fgcmPars.parQESysSlopeLoc +
                         uWashIndex] += 1

        # note that this store doesn't need locking because we only access
        #  a given array from a single process

        totalArr = snmm.getArray(self.totalHandleDict[thisCore])
        totalArr[:] = totalArr[:] + partialArray

        # and we're done
        return None

    def __getstate__(self):
        # Don't try to pickle the logger.

        state = self.__dict__.copy()
        del state['fgcmLog']
        return state


#!/usr/bin/env python
# coding: utf-8

# # Face Recognition
#
# Welcome! In this assignment, you're going to build a face recognition system. Many of the ideas presented here are from [FaceNet](https://arxiv.org/pdf/1503.03832.pdf). In the lecture, you also encountered [DeepFace](https://research.fb.com/wp-content/uploads/2016/11/deepface-closing-the-gap-to-human-level-performance-in-face-verification.pdf).
#
# Face recognition problems commonly fall into one of two categories:
#
# **Face Verification** "Is this the claimed person?" For example, at some airports, you can pass through customs by letting a system scan your passport and then verifying that you (the person carrying the passport) are the correct person. A mobile phone that unlocks using your face is also using face verification. This is a 1:1 matching problem.
#
# **Face Recognition** "Who is this person?" For example, the video lecture showed a [face recognition video](https://www.youtube.com/watch?v=wr4rx0Spihs) of Baidu employees entering the office without needing to otherwise identify themselves. This is a 1:K matching problem.
#
# FaceNet learns a neural network that encodes a face image into a vector of 128 numbers. By comparing two such vectors, you can then determine if two pictures are of the same person.
#
# By the end of this assignment, you'll be able to:
#
# * Differentiate between face recognition and face verification
# * Implement one-shot learning to solve a face recognition problem
# * Apply the triplet loss function to learn a network's parameters in the context of face recognition
# * Explain how to pose face recognition as a binary classification problem
# * Map face images into 128-dimensional encodings using a pretrained model
# * Perform face verification and face recognition with these encodings
#
# **Channels-last notation**
#
# For this assignment, you'll be using a pre-trained model which represents ConvNet activations using a "channels last" convention, as used during the lecture and in previous programming assignments.
#
# In other words, a batch of images will be of shape $(m, n_H, n_W, n_C)$.

# ## Table of Contents
#
# - [1 - Packages](#1)
# - [2 - Naive Face Verification](#2)
# - [3 - Encoding Face Images into a 128-Dimensional Vector](#3)
#     - [3.1 - Using a ConvNet to Compute Encodings](#3-1)
#     - [3.2 - The Triplet Loss](#3-2)
#         - [Exercise 1 - triplet_loss](#ex-1)
# - [4 - Loading the Pre-trained Model](#4)
# - [5 - Applying the Model](#5)
#     - [5.1 - Face Verification](#5-1)
#         - [Exercise 2 - verify](#ex-2)
#     - [5.2 - Face Recognition](#5-2)
#         - [Exercise 3 - who_is_it](#ex-3)
# - [6 - References](#6)

#
# ## 1 - Packages
#
# Go ahead and run the cell below to import the packages you'll need.

# In[1]:

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, ZeroPadding2D, Activation, Input, concatenate
from tensorflow.keras.models import Model
from tensorflow.keras.layers import BatchNormalization
from tensorflow.keras.layers import MaxPooling2D, AveragePooling2D
from tensorflow.keras.layers import Concatenate
from tensorflow.keras.layers import Lambda, Flatten, Dense
from tensorflow.keras.initializers import glorot_uniform
from tensorflow.keras.layers import Layer
from tensorflow.keras import backend as K
K.set_image_data_format('channels_last')
import os
import numpy as np
from numpy import genfromtxt
import pandas as pd
import tensorflow as tf
import PIL

get_ipython().run_line_magic('matplotlib', 'inline')
get_ipython().run_line_magic('load_ext', 'autoreload')
get_ipython().run_line_magic('autoreload', '2')

#
# ## 2 - Naive Face Verification
#
# In Face Verification, you're given two images and you have to determine if they are of the same person. The simplest way to do this is to compare the two images pixel-by-pixel. If the distance between the raw images is below a chosen threshold, it may be the same person!
#
#
#
   Figure 1

#
# Of course, this algorithm performs poorly, since the pixel values change dramatically due to variations in lighting, orientation of the person's face, minor changes in head position, and so on.
#
# You'll see that rather than using the raw image, you can learn an encoding, $f(img)$.
#
# By using an encoding for each image, an element-wise comparison produces a more accurate judgement as to whether two pictures are of the same person.

#
# ## 3 - Encoding Face Images into a 128-Dimensional Vector
#
#
# ### 3.1 - Using a ConvNet to Compute Encodings
#
# The FaceNet model takes a lot of data and a long time to train. So following the common practice in applied deep learning, you'll load weights that someone else has already trained. The network architecture follows the Inception model from [Szegedy *et al*..](https://arxiv.org/abs/1409.4842) An Inception network implementation has been provided for you, and you can find it in the file `inception_blocks_v2.py` to get a closer look at how it is implemented.
#
# *Hot tip:* Go to "File->Open..." at the top of this notebook. This opens the file directory that contains the `.py` file).
#
# The key things to be aware of are:
#
# - This network uses 160x160 dimensional RGB images as its input. Specifically, a face image (or batch of $m$ face images) as a tensor of shape $(m, n_H, n_W, n_C) = (m, 160, 160, 3)$
# - The input images are originally of shape 96x96, thus, you need to scale them to 160x160. This is done in the `img_to_encoding()` function.
# - The output is a matrix of shape $(m, 128)$ that encodes each input face image into a 128-dimensional vector
#
# Run the cell below to create the model for face images!

# In[2]:

from tensorflow.keras.models import model_from_json

json_file = open('keras-facenet-h5/model.json', 'r')
loaded_model_json = json_file.read()
json_file.close()
model = model_from_json(loaded_model_json)
model.load_weights('keras-facenet-h5/model.h5')

# Now summarize the input and output shapes:

# In[3]:

print(model.inputs)
print(model.outputs)

# By using a 128-neuron fully connected layer as its last layer, the model ensures that the output is an encoding vector of size 128. You then use the encodings to compare two face images as follows:
#
#
#
   Figure 2:
  By computing the distance between two encodings and thresholding, you can determine if the two pictures represent the same person

#
# So, an encoding is a good one if:
#
# - The encodings of two images of the same person are quite similar to each other.
# - The encodings of two images of different persons are very different.
#
# The triplet loss function formalizes this, and tries to "push" the encodings of two images of the same person (Anchor and Positive) closer together, while "pulling" the encodings of two images of different persons (Anchor, Negative) further apart.
#
#

#
   Figure 3:
   In the next section,  you'll call the pictures from left to right: Anchor (A), Positive (P), Negative (N)

#
# ### 3.2 - The Triplet Loss
#
# **Important Note**: Since you're using a pretrained model, you won't actually need to implement the triplet loss function in this assignment. *However*, the triplet loss is the main ingredient of the face recognition algorithm, and you'll need to know how to use it for training your own FaceNet model, as well as other types of image similarity problems. Therefore, you'll implement it below, for fun and edification. :)
#
# For an image $x$, its encoding is denoted as $f(x)$, where $f$ is the function computed by the neural network.
#
#
#
# Training will use triplets of images $(A, P, N)$:
#
# - A is an "Anchor" image--a picture of a person.
# - P is a "Positive" image--a picture of the same person as the Anchor image.
# - N is a "Negative" image--a picture of a different person than the Anchor image.
#
# These triplets are picked from the training dataset. $(A^{(i)}, P^{(i)}, N^{(i)})$ is used here to denote the $i$-th training example.
#
# You'd like to make sure that an image $A^{(i)}$ of an individual is closer to the Positive $P^{(i)}$ than to the Negative image $N^{(i)}$) by at least a margin $\alpha$:
#
# $$
# || f\left(A^{(i)}\right)-f\left(P^{(i)}\right)||_{2}^{2}+\alpha<|| f\left(A^{(i)}\right)-f\left(N^{(i)}\right)||_{2}^{2}
# $$
#
#
# You would thus like to minimize the following "triplet cost":
#
# $$\mathcal{J} = \sum^{m}_{i=1} \large[ \small \underbrace{\mid \mid f(A^{(i)}) - f(P^{(i)}) \mid \mid_2^2}_\text{(1)} - \underbrace{\mid \mid f(A^{(i)}) - f(N^{(i)}) \mid \mid_2^2}_\text{(2)} + \alpha \large ] \small_+ \tag{3}$$
# Here, the notation "$[z]_+$" is used to denote $max(z,0)$.
#
# **Notes**:
#
# - The term (1) is the squared distance between the anchor "A" and the positive "P" for a given triplet; you want this to be small.
# - The term (2) is the squared distance between the anchor "A" and the negative "N" for a given triplet, you want this to be relatively large. It has a minus sign preceding it because minimizing the negative of the term is the same as maximizing that term.
# - $\alpha$ is called the margin. It's a hyperparameter that you pick manually. You'll use $\alpha = 0.2$.
#
# Most implementations also rescale the encoding vectors to haven L2 norm equal to one (i.e., $\mid \mid f(img)\mid \mid_2$=1); you won't have to worry about that in this assignment.
#
#
# ### Exercise 1 - triplet_loss
#
# Implement the triplet loss as defined by formula (3). These are the 4 steps:
#
# 1. Compute the distance between the encodings of "anchor" and "positive": $\mid \mid f(A^{(i)}) - f(P^{(i)}) \mid \mid_2^2$
# 2. Compute the distance between the encodings of "anchor" and "negative": $\mid \mid f(A^{(i)}) - f(N^{(i)}) \mid \mid_2^2$
# 3. Compute the formula per training example: $ \mid \mid f(A^{(i)}) - f(P^{(i)}) \mid \mid_2^2 - \mid \mid f(A^{(i)}) - f(N^{(i)}) \mid \mid_2^2 + \alpha$
# 4. Compute the full formula by taking the max with zero and summing over the training examples:$$\mathcal{J} = \sum^{m}_{i=1} \large[ \small \mid \mid f(A^{(i)}) - f(P^{(i)}) \mid \mid_2^2 - \mid \mid f(A^{(i)}) - f(N^{(i)}) \mid \mid_2^2+ \alpha \large ] \small_+ \tag{3}$$
#
# *Hints*:
#
# - Useful functions: `tf.reduce_sum()`, `tf.square()`, `tf.subtract()`, `tf.add()`, `tf.maximum()`.
#
# - For steps 1 and 2, sum over the entries of $\mid \mid f(A^{(i)}) - f(P^{(i)}) \mid \mid_2^2$ and $\mid \mid     f(A^{(i)}) - f(N^{(i)}) \mid \mid_2^2$.
#
# - For step 4, you will sum over the training examples.
#
# *Additional Hints*:
#
# - Recall that the square of the L2 norm is the sum of the squared differences: $||x - y||_{2}^{2} = \sum_{i=1}^{N}(x_{i} - y_{i})^{2}$
#
# - Note that the anchor, positive and negative encodings are of shape (*m*,128), where *m* is the number of training examples and 128 is the number of elements used to encode a single example.
#
# - For steps 1 and 2, maintain the number of *m* training examples and sum along the 128 values of each encoding. `tf.reduce_sum` has an axis parameter. This chooses along which axis the sums are applied.
#
# - Note that one way to choose the last axis in a tensor is to use negative indexing (axis=-1).
#
# - In step 4, when summing over training examples, the result will be a single scalar value.
#
# - For `tf.reduce_sum` to sum across all axes, keep the default value axis=None.

# In[18]:

# UNQ_C1(UNIQUE CELL IDENTIFIER, DO NOT EDIT)
# GRADED FUNCTION: triplet_loss

def triplet_loss(y_true, y_pred, alpha = 0.2):
    """
    Implementation of the triplet loss as defined by formula (3)

    Arguments:
    y_true -- true labels, required when you define a loss in Keras, you don't need it in this function.
    y_pred -- python list containing three objects:
            anchor -- the encodings for the anchor images, of shape (None, 128)
            positive -- the encodings for the positive images, of shape (None, 128)
            negative -- the encodings for the negative images, of shape (None, 128)

    Returns:
    loss -- real number, value of the loss
    """

    anchor, positive, negative = y_pred[0], y_pred[1], y_pred[2]

    ### START CODE HERE
    #(≈ 4 lines)
    # Step 1: Compute the (encoding) distance between the anchor and the positive
    pos_dist = tf.reduce_sum(tf.square(tf.subtract(anchor, positive)), axis=-1)
    # Step 2: Compute the (encoding) distance between the anchor and the negative
    neg_dist = tf.reduce_sum(tf.square(tf.subtract(anchor, negative)), axis=-1)
    # Step 3: subtract the two previous distances and add alpha.
    basic_loss = pos_dist - neg_dist + alpha
    # Step 4: Take the maximum of basic_loss and 0.0. Sum over the training examples.
    loss = tf.reduce_sum(tf.maximum(basic_loss, 0))
    ### END CODE HERE

    return loss

# In[19]:

# BEGIN UNIT TEST
tf.random.set_seed(1)
y_true = (None, None, None) # It is not used
y_pred = (tf.keras.backend.random_normal([3, 128], mean=6, stddev=0.1, seed = 1),
          tf.keras.backend.random_normal([3, 128], mean=1, stddev=1, seed = 1),
          tf.keras.backend.random_normal([3, 128], mean=3, stddev=4, seed = 1))
loss = triplet_loss(y_true, y_pred)

assert type(loss) == tf.python.framework.ops.EagerTensor, "Use tensorflow functions"
print("loss = " + str(loss))

y_pred_perfect = ([1., 1.], [1., 1.], [1., 1.,])
loss = triplet_loss(y_true, y_pred_perfect, 5)
assert loss == 5, "Wrong value. Did you add the alpha to basic_loss?"
y_pred_perfect = ([1., 1.],[1., 1.], [0., 0.,])
loss = triplet_loss(y_true, y_pred_perfect, 3)
assert loss == 1., "Wrong value. Check that pos_dist = 0 and neg_dist = 2 in this example"
y_pred_perfect = ([1., 1.],[0., 0.], [1., 1.,])
loss = triplet_loss(y_true, y_pred_perfect, 0)
assert loss == 2., "Wrong value. Check that pos_dist = 2 and neg_dist = 0 in this example"
y_pred_perfect = ([0., 0.],[0., 0.], [0., 0.,])
loss = triplet_loss(y_true, y_pred_perfect, -2)
assert loss == 0, "Wrong value. Are you taking the maximum between basic_loss and 0?"
y_pred_perfect = ([[1., 0.], [1., 0.]],[[1., 0.], [1., 0.]], [[0., 1.], [0., 1.]])
loss = triplet_loss(y_true, y_pred_perfect, 3)
assert loss == 2., "Wrong value. Are you applying tf.reduce_sum to get the loss?"
y_pred_perfect = ([[1., 1.], [2., 0.]], [[0., 3.], [1., 1.]], [[1., 0.], [0., 1.,]])
loss = triplet_loss(y_true, y_pred_perfect, 1)
if (loss == 4.):
    raise Exception('Perhaps you are not using axis=-1 in reduce_sum?')
assert loss == 5, "Wrong value. Check your implementation"
# END UNIT TEST

# **Expected Output**:
#
#

#

#
#             loss
#
#
#            527.2598
#
#

#

#
# ## 4 - Loading the Pre-trained Model
#
# FaceNet is trained by minimizing the triplet loss. But since training requires a lot of data and a lot of computation, you won't train it from scratch here. Instead, you'll load a previously trained model in the following cell; which might take a couple of minutes to run.

# In[20]:

FRmodel = model

# Here are some examples of distances between the encodings between three individuals:
#
#

#
   Figure 4:
   Example of distance outputs between three individuals' encodings

#
# Now use this model to perform face verification and face recognition!

#
# ## 5 - Applying the Model
#
# You're building a system for an office building where the building manager would like to offer facial recognition to allow the employees to enter the building.
#
# You'd like to build a face verification system that gives access to a list of people. To be admitted, each person has to swipe an identification card at the entrance. The face recognition system then verifies that they are who they claim to be.
#
#
# ### 5.1 - Face Verification
#
# Now you'll build a database containing one encoding vector for each person who is allowed to enter the office. To generate the encoding, you'll use `img_to_encoding(image_path, model)`, which runs the forward propagation of the model on the specified image.
#
# Run the following code to build the database (represented as a Python dictionary). This database maps each person's name to a 128-dimensional encoding of their face.

# In[21]:

#tf.keras.backend.set_image_data_format('channels_last')
def img_to_encoding(image_path, model):
    img = tf.keras.preprocessing.image.load_img(image_path, target_size=(160, 160))
    img = np.around(np.array(img) / 255.0, decimals=12)
    x_train = np.expand_dims(img, axis=0)
    embedding = model.predict_on_batch(x_train)
    return embedding / np.linalg.norm(embedding, ord=2)

# In[22]:

database = {}
database["danielle"] = img_to_encoding("images/danielle.png", FRmodel)
database["younes"] = img_to_encoding("images/younes.jpg", FRmodel)
database["tian"] = img_to_encoding("images/tian.jpg", FRmodel)
database["andrew"] = img_to_encoding("images/andrew.jpg", FRmodel)
database["kian"] = img_to_encoding("images/kian.jpg", FRmodel)
database["dan"] = img_to_encoding("images/dan.jpg", FRmodel)
database["sebastiano"] = img_to_encoding("images/sebastiano.jpg", FRmodel)
database["bertrand"] = img_to_encoding("images/bertrand.jpg", FRmodel)
database["kevin"] = img_to_encoding("images/kevin.jpg", FRmodel)
database["felix"] = img_to_encoding("images/felix.jpg", FRmodel)
database["benoit"] = img_to_encoding("images/benoit.jpg", FRmodel)
database["arnaud"] = img_to_encoding("images/arnaud.jpg", FRmodel)

# Load the images of Danielle and Kian:

# In[23]:

danielle = tf.keras.preprocessing.image.load_img("images/danielle.png", target_size=(160, 160))
kian = tf.keras.preprocessing.image.load_img("images/kian.jpg", target_size=(160, 160))

# In[24]:

np.around(np.array(kian) / 255.0, decimals=12).shape

# In[25]:

kian

# In[26]:

np.around(np.array(danielle) / 255.0, decimals=12).shape

# In[27]:

danielle

# Now, when someone shows up at your front door and swipes their ID card (thus giving you their name), you can look up their encoding in the database, and use it to check if the person standing at the front door matches the name on the ID.
#
#
# ### Exercise 2 - verify
#
# Implement the `verify()` function, which checks if the front-door camera picture (`image_path`) is actually the person called "identity". You will have to go through the following steps:
#
# - Compute the encoding of the image from `image_path`.
# - Compute the distance between this encoding and the encoding of the identity image stored in the database.
# - Open the door if the distance is less than 0.7, else do not open it.
#
# As presented above, you should use the L2 distance `np.linalg.norm`.
#
# **Note**: In this implementation, compare the L2 distance, not the square of the L2 distance, to the threshold 0.7.
#
# *Hints*:
#
# - `identity` is a string that is also a key in the database dictionary.
# - `img_to_encoding` has two parameters: the image_path and model.

# In[30]:

# UNQ_C2(UNIQUE CELL IDENTIFIER, DO NOT EDIT)
# GRADED FUNCTION: verify

def verify(image_path, identity, database, model):
    """
    Function that verifies if the person on the "image_path" image is "identity".

    Arguments:
        image_path -- path to an image
        identity -- string, name of the person you'd like to verify the identity. Has to be an employee who works in the office.
        database -- python dictionary mapping names of allowed people's names (strings) to their encodings (vectors).
        model -- your Inception model instance in Keras

    Returns:
        dist -- distance between the image_path and the image of "identity" in the database.
        door_open -- True, if the door should open. False otherwise.
    """
    ### START CODE HERE
    # Step 1: Compute the encoding for the image. Use img_to_encoding() see example above. (≈ 1 line)
    encoding = img_to_encoding(image_path, model)
    # Step 2: Compute distance with identity's image (≈ 1 line)
    dist = np.linalg.norm(database[identity] - encoding)
    # Step 3: Open the door if dist < 0.7, else don't open (≈ 3 lines)
    if dist < 0.7:
        print("It's " + str(identity) + ", welcome in!")
        door_open = True
    else:
        print("It's not " + str(identity) + ", please go away")
        door_open = False
    ### END CODE HERE
    return dist, door_open

# Younes is trying to enter the office and the camera takes a picture of him ("images/camera_0.jpg"). Let's run your verification algorithm on this picture:
#
#

# In[31]:

# BEGIN UNIT TEST
assert(np.allclose(verify("images/camera_1.jpg", "bertrand", database, FRmodel), (0.54364836, True)))
assert(np.allclose(verify("images/camera_3.jpg", "bertrand", database, FRmodel), (0.38616243, True)))
assert(np.allclose(verify("images/camera_1.jpg", "younes", database, FRmodel), (1.3963861, False)))
assert(np.allclose(verify("images/camera_3.jpg", "younes", database, FRmodel), (1.3872949, False)))

verify("images/camera_0.jpg", "younes", database, FRmodel)
# END UNIT TEST

# **Expected Output**:
#
#

#

#
#             It's Younes, welcome in!
#
#
#            (0.5992946, True)
#
#

#

# Benoit, who does not work in the office, stole Kian's ID card and tried to enter the office. Naughty Benoit! The camera took a picture of Benoit ("images/camera_2.jpg).
#
#
#
# Run the verification algorithm to check if Benoit can enter.

# In[32]:

verify("images/camera_2.jpg", "kian", database, FRmodel)

# **Expected Output**:
#
#

#

#
#             It's not Kian, please go away
#
#
#            (1.0259346, False)
#
#

#

#
# ### 5.2 - Face Recognition
#
# Your face verification system is mostly working. But since Kian got his ID card stolen, when he came back to the office the next day he couldn't get in!
#
# To solve this, you'd like to change your face verification system to a face recognition system. This way, no one has to carry an ID card anymore. An authorized person can just walk up to the building, and the door will unlock for them!
#
# You'll implement a face recognition system that takes as input an image, and figures out if it is one of the authorized persons (and if so, who). Unlike the previous face verification system, you will no longer get a person's name as one of the inputs.
#
#
# ### Exercise 3 - who_is_it
#
# Implement `who_is_it()` with the following steps:
#
# - Compute the target encoding of the image from `image_path`
# - Find the encoding from the database that has smallest distance with the target encoding.
# - Initialize the `min_dist` variable to a large enough number (100). This helps you keep track of the closest encoding to the input's encoding.
# - Loop over the database dictionary's names and encodings. To loop use for (name, db_enc) in `database.items()`.
# - Compute the L2 distance between the target "encoding" and the current "encoding" from the database. If this distance is less than the min_dist, then set min_dist to dist, and identity to name.

# In[33]:

# UNQ_C3(UNIQUE CELL IDENTIFIER, DO NOT EDIT)
# GRADED FUNCTION: who_is_it

def who_is_it(image_path, database, model):
    """
    Implements face recognition for the office by finding who is the person on the image_path image.

    Arguments:
        image_path -- path to an image
        database -- database containing image encodings along with the name of the person on the image
        model -- your Inception model instance in Keras

    Returns:
        min_dist -- the minimum distance between image_path encoding and the encodings from the database
        identity -- string, the name prediction for the person on image_path
    """

    ### START CODE HERE

    ## Step 1: Compute the target "encoding" for the image. Use img_to_encoding() see example above. ## (≈ 1 line)
    encoding =  img_to_encoding(image_path, model)

    ## Step 2: Find the closest encoding ##

    # Initialize "min_dist" to a large value, say 100 (≈1 line)
    min_dist = 100

    # Loop over the database dictionary's names and encodings.
    for (name, db_enc) in database.items():

        # Compute L2 distance between the target "encoding" and the current db_enc from the database. (≈ 1 line)
        dist = np.linalg.norm(db_enc - encoding)

        # If this distance is less than the min_dist, then set min_dist to dist, and identity to name. (≈ 3 lines)
        if dist < min_dist:
            min_dist = dist
            identity = name
    ### END CODE HERE

    if min_dist > 0.7:
        print("Not in the database.")
    else:
        print ("it's " + str(identity) + ", the distance is " + str(min_dist))

    return min_dist, identity

# Younes is at the front door and the camera takes a picture of him ("images/camera_0.jpg"). Let's see if your `who_it_is()` algorithm identifies Younes.

# In[34]:

# BEGIN UNIT TEST
# Test 1 with Younes pictures
who_is_it("images/camera_0.jpg", database, FRmodel)

# Test 2 with Younes pictures
test1 = who_is_it("images/camera_0.jpg", database, FRmodel)
assert np.isclose(test1[0], 0.5992946)
assert test1[1] == 'younes'

# Test 3 with Younes pictures
test2 = who_is_it("images/younes.jpg", database, FRmodel)
assert np.isclose(test2[0], 0.0)
assert test2[1] == 'younes'
# END UNIT TEST

# **Expected Output**:
#
#

#

#
#             it's Younes, the distance is 0.5992946
#
#
#            (0.5992946, 'younes')
#
#

#

#
# You can change "camera_0.jpg" (picture of Younes) to "camera_1.jpg" (picture of Bertrand) and see the result.

# **Congratulations**!
# You've completed this assignment, and your face recognition system is working well! It not only lets in authorized persons, but now people don't need to carry an ID card around anymore!
#
# You've now seen how a state-of-the-art face recognition system works, and can describe the difference between face recognition and face verification. Here's a quick recap of what you've accomplished:
#
# - Posed face recognition as a binary classification problem
# - Implemented one-shot learning for a face recognition problem
# - Applied the triplet loss function to learn a network's parameters in the context of face recognition
# - Mapped face images into 128-dimensional encodings using a pretrained model
# - Performed face verification and face recognition with these encodings
#
# Great work!

#
#
# **What you should remember**:
#
# - Face verification solves an easier 1:1 matching problem; face recognition addresses a harder 1:K matching problem.
#
# - Triplet loss is an effective loss function for training a neural network to learn an encoding of a face image.
#
# - The same encoding can be used for verification and recognition. Measuring distances between two images' encodings allows you to determine whether they are pictures of the same person.

# **Ways to improve your facial recognition model**:
#
# Although you won't implement these here, here are some ways to further improve the algorithm:
#
# - Put more images of each person (under different lighting conditions, taken on different days, etc.) into the database. Then, given a new image, compare the new face to multiple pictures of the person. This would increase accuracy.
#
# - Crop the images to contain just the face, and less of the "border" region around the face. This preprocessing removes some of the irrelevant pixels around the face, and also makes the algorithm more robust.

#
# ## 6 - References
# 1. Florian Schroff, Dmitry Kalenichenko, James Philbin (2015). [FaceNet: A Unified Embedding for Face Recognition and Clustering](https://arxiv.org/pdf/1503.03832.pdf)
#
# 2. Yaniv Taigman, Ming Yang, Marc'Aurelio Ranzato, Lior Wolf (2014). [DeepFace: Closing the gap to human-level performance in face verification](https://research.fb.com/wp-content/uploads/2016/11/deepface-closing-the-gap-to-human-level-performance-in-face-verification.pdf)
#
# 3. This implementation also took a lot of inspiration from the official FaceNet github repository: https://github.com/davidsandberg/facenet
#
# 4. Further inspiration was found here: https://machinelearningmastery.com/how-to-develop-a-face-recognition-system-using-facenet-in-keras-and-an-svm-classifier/
#
# 5. And here: https://github.com/nyoki-mtl/keras-facenet/blob/master/notebook/tf_to_keras.ipynb

4 - Convolutional Neural Networks/Face_Recognition.py


"""
Defines comparer functions used by FormulaGrader and its subclasses.

Simple Comparer Functions
=========================

A comparer function must have signature
`comparer_func(comparer_params_eval, student_eval, utils)` and should return
True, False, 'partial', or a dictionary with required key 'grade_decimal' and
optional key 'msg'. When `FormulaGrader` (or its subclasses) call your custom
comparer function, `comparer_func`'s argument values are:

- `comparer_params_eval`: The `comparer_params` list, numerically evaluated
  according to variable and function sampling.
- `student_eval`: The student's input, numerically evaluated according to
  variable and function sampling.
- `utils`: A convenience object that may be helpful when writing custom
  comparer functions. It has attributes:

    - `utils.tolerance`: The tolerance specified in grader configuration,
      `0.01%` by default
    - `utils.within_tolerance(x, y)`: checks that `y` is within specified
      tolerance of `x`. Can handle scalars, vectors, and matrices.
      If tolerance was specified as a percentage, then checks that
      `|x-y| < tolerance * x`.

    Comparer functions used inside `MatrixGrader` have the following additional
    `utils` method:

    - `utils.validate_shape(student_eval, shape)`: Checks that `student_eval`
      has specified `shape`, where `shape` is a Numpy shape tuple.

A comparer function must return either:

  - a boolean, or
  - a dictionary with keys:
      - `'grade_decimal'`: number between 0 and 1 (required)
      - `'ok'`: `True` or `False` or `'partial'` (optional, inferred from
        grade_decimal by default)
      - `'msg'`: a feedback message (optional, defaults to `''`)

NOTE: doctests in this module show how the comparer function would be used
      inside a grader

Correlated Comparers
====================
See ./baseclasses.py and ./linear_comparer.py for examples.
"""
from __future__ import print_function, division, absolute_import, unicode_literals

from numbers import Number
import numpy as np

import six
from voluptuous import Schema, Required, Any, Range, All

from mitxgraders.exceptions import InputTypeError, StudentFacingError
from mitxgraders.helpers.validatorfuncs import is_callable, Nullable, text_string
from mitxgraders.helpers.calc.mathfuncs import is_nearly_zero
from mitxgraders.helpers.calc.math_array import are_same_length_vectors, is_vector
from mitxgraders.comparers.baseclasses import Comparer, CorrelatedComparer

def identity_transform(x):
    """
    Returns the input.

    Note: used instead of lambdas because it prints a nice name.
    """
    return x

class EqualityComparer(Comparer):
    """
    This comparer checks for equality between the student and instructor results,
    up to a desired tolerance. If desired, a transforming function can be applied
    before the comparison is carried out.

    comparer_params: ['expect']

    By default, equality_comparer is used in FormulaGrader, NumericalGrader and MatrixGrader.
    >>> from mitxgraders import *
    >>> equality_comparer == EqualityComparer()
    True
    >>> grader = FormulaGrader(
    ...     answers='2*x',
    ...     variables=['x']
    ... )
    >>> grader(None, 'x')['ok']
    False
    >>> grader(None, 'x*2')['ok']
    True

    The following example applies cosine to the expected answer and student input
    before comparison:
    >>> import numpy as np
    >>> grader = FormulaGrader(
    ...     answers={
    ...         'comparer': EqualityComparer(transform=np.cos),
    ...         'comparer_params': ['x']
    ...     },
    ...     variables=['x']
    ... )
    >>> grader(None, 'x')['ok']
    True
    >>> grader(None, '-x')['ok']
    True
    >>> grader(None, 'x + 2*pi')['ok']
    True
    >>> grader(None, 'x + pi')['ok']
    False

    The following example takes the norm of the expected answer and student input
    before comparison. Note the different method of changing the comparer.
    >>> MatrixGrader.set_default_comparer(EqualityComparer(transform=np.linalg.norm))
    >>> grader = MatrixGrader(
    ...     answers='[1, 0, 0]'
    ... )
    >>> grader(None, '[1, 0, 0]')['ok']
    True
    >>> grader(None, '[0, 1, 0]')['ok']
    True
    >>> grader(None, '[1/sqrt(2), 0, 1/sqrt(2)]')['ok']
    True
    >>> MatrixGrader.reset_default_comparer()

    FormulaGrader and MatrixGrader both have set_default_comparer() and
    reset_default_comparer() methods.

    """
    schema_config = Schema({
        Required('transform', default=None): All(
            Nullable(is_callable),
            # if f is None, coerce to identity function
            lambda f: identity_transform if f is None else f
        )
    })

    @staticmethod
    def validate(expected_eval, student_eval, utils):
        if hasattr(utils, 'validate_shape'):
            # in numpy, scalars have empty tuples as their shapes
            shape = tuple() if isinstance(expected_eval, Number) else expected_eval.shape
            utils.validate_shape(student_eval, shape)

    def __call__(self, comparer_params_eval, student_eval, utils):
        expected_eval = comparer_params_eval[0]
        self.validate(expected_eval, student_eval, utils)

        transform = self.config['transform']
        expected_eval = transform(expected_eval)
        student_eval = transform(student_eval)

        return utils.within_tolerance(expected_eval, student_eval)

equality_comparer = EqualityComparer()

class MatrixEntryComparer(CorrelatedComparer):
    """
    Default comparer for MatrixGrader. Compares student and instructor array
    evaluations entry-by-entry for equality. Note that despite the name, this comparer
    works equally well on vectors/matrices/tensors.

    Configuration
    =============
        transform (None | function): same as EqualityComparer (default None)

        entry_partial_credit ('proportional' | number): Determines how partial credit
            is awarded. If set to 'proportional', then credit is proportional to
            the number of correct array entries. If a numeric value betweem 0 and 1
            is provided, this flat rate of partial credit is provided as long as
            some but not all entries are correct. Default is the numeric value 0
            (no partial credit).

        entry_partial_msg (str): A text string message shown when partial credit
            is awarded. The string may optionally contain the formatting key {error_indices},
            which will be replaced with a diagram showing the correct/incorrect entries.
            To show no message, use the the empty string.
            Default value is:
            "Some array entries are incorrect, marked below:\n{error_locations}"
    """

    default_msg = "Some array entries are incorrect, marked below:\n{error_locations}"
    schema_config = EqualityComparer.schema_config.extend({
        Required('entry_partial_credit', default=0): Any(All(float, Range(0, 1)), 0, 1, 'proportional'),
        Required('entry_partial_msg', default=default_msg): text_string
    })

    @staticmethod
    def format_message_with_locations(format_string, locs):
        """
        Returns format_string with {error_locations} replaced by a diagram showing
        correct/incorrect entries.

        Arguments:
            format_string: a string that may contain {error_locations} formatting key.
            locs: a boolean array with False values indicating incorrect entries
        """
        # Not the most elegant way to do these replacements, but this was what
        # I came up with to minimize the amount of extra u prefixes in Python 2

        # These are the edX colors, at least as of July 2019
        bad_str = '\u2717'
        good_str = '\u2713'
        matrix_as_text = six.text_type(locs).replace("  ", " ").replace("[ ", "[")
        matrix_as_text = matrix_as_text.replace("True", good_str).replace("False", bad_str)
        matrix_as_text = matrix_as_text.replace('\n', '
')
        formatted_locs = '
{mat}
'.format(mat=matrix_as_text)
        return format_string.format(error_locations=formatted_locs)

    @staticmethod
    def validate(expected_evals, student_evals, utils):
        for x, y in zip(expected_evals, student_evals):
            EqualityComparer.validate(x, y, utils)

    def __call__(self, comparer_params_evals, student_evals, utils):
        expected_evals = [params[0] for params in comparer_params_evals]
        self.validate(expected_evals, student_evals, utils)

        transform = self.config['transform']
        expected_evals = [transform(x) for x in expected_evals]
        student_evals = [transform(x) for x in student_evals]
        vec_within_tol = np.vectorize(utils.within_tolerance)
        # comparisons_by_eval is a boolean array of entry-by-entry comparisons,
        # one for each comparison. Its numpy shape is (n_evals, *eval_shape)
        comparisons_by_eval = vec_within_tol(expected_evals, student_evals)
        comparisons_summary = np.all(comparisons_by_eval, axis=0)

        num_entries = comparisons_summary.size
        percent_correct = np.sum(comparisons_summary).item()/num_entries
        msg = self.format_message_with_locations(self.config['entry_partial_msg'], comparisons_summary)
        partial_credit = self.config['entry_partial_credit']

        if percent_correct == 1:
            return True
        elif percent_correct == 0:
            return {'ok': False, 'grade_decimal': 0, 'msg': msg}
        elif partial_credit == 'proportional':
            return {'ok': 'partial', 'grade_decimal': percent_correct, 'msg': msg}
        else:
            return {'ok': 'partial', 'grade_decimal': partial_credit, 'msg': msg}

def between_comparer(comparer_params_eval, student_eval, utils):
    """
    Used to check that input is real and between two parameters.

    comparer_params: ['start', 'stop']

    Example:
    >>> from mitxgraders import NumericalGrader
    >>> grader = NumericalGrader(
    ...     answers={
    ...         'comparer': between_comparer,
    ...         'comparer_params': ['1e6', '1e9']
    ...     }
    ... )
    >>> grader(None, '2.5e8')['ok']
    True
    >>> grader(None, '0.001e8')['ok']
    False
    >>> grader(None, '5e7')['ok']
    True

    Input must be real:
    >>> try:
    ...     grader(None, '5e8+2e6*i')['ok']
    ... except InputTypeError as error:
    ...     print(error)
    Input must be real.
    """
    start, stop = comparer_params_eval

    if not np.isreal(student_eval):
        raise InputTypeError("Input must be real.")

    return start <= student_eval <= stop

def congruence_comparer(comparer_params_eval, student_eval, utils):
    """
    Compares the student input to a target, moduli a given modulus.
    Will often set modulus to 2*pi in order to compare angles.

    comparer_params: [target, modulus]

    Example usage:
    >>> from mitxgraders import FormulaGrader
    >>> grader = FormulaGrader(
    ...     answers={
    ...         'comparer': congruence_comparer,
    ...         'comparer_params': [
    ...             'b^2/a', # target
    ...             'c'      # modulus
    ...         ]
    ...     },
    ...     variables=['a', 'b', 'c']
    ... )
    >>> grader(None, 'b^2/a')['ok']
    True
    >>> grader(None, 'b^2/a + 1.5*c')['ok']
    False
    >>> grader(None, 'b^2/a + 2*c  ')['ok']
    True
    """
    expected, modulus = comparer_params_eval

    expected_reduced = expected % modulus
    input_reduced = student_eval % modulus
    return utils.within_tolerance(expected_reduced, input_reduced)

def eigenvector_comparer(comparer_params_eval, student_eval, utils):
    """
    Used to check that a student's answer is an eigenvector of a matrix
    with a given eigenvalue. Ignores scaling of the eigenvector.

    comparer_params: [matrix, eigenvalue]

    Example Usage:
    >>> from mitxgraders import MatrixGrader
    >>> grader = MatrixGrader(
    ...     answers={
    ...         'comparer_params': [
    ...             '[[1, x], [x, -1]]',    # matrix
    ...             'sqrt(1+x^2)'           # eigenvalue
    ...         ],
    ...         'comparer': eigenvector_comparer
    ...     },
    ...     variables=['x']
    ... )
    >>> grader(None, '[1+sqrt(1+x^2), x]')['ok']
    True
    >>> grader(None, '2*[1+sqrt(1+x^2), x]')['ok']
    True
    >>> grader(None, '[1+sqrt(1+x^2), 1]')['ok']
    False
    >>> grader(None, '[0, 0]') == {
    ...     'ok': False,
    ...     'msg': 'Eigenvectors must be nonzero.',
    ...     'grade_decimal': 0
    ... }
    True

    """

    matrix, eigenvalue = comparer_params_eval

    # matrix is square with shape (n, n); student input should have shape (n, )
    expected_input_shape = (matrix.shape[0], )
    utils.validate_shape(student_eval, expected_input_shape)

    expected = eigenvalue * student_eval
    actual = matrix * student_eval

    if utils.within_tolerance(0, np.linalg.norm(student_eval)):
        return {
            'ok': False,
            'grade_decimal': 0,
            'msg': 'Eigenvectors must be nonzero.'
        }

    return utils.within_tolerance(actual, expected)

def vector_span_comparer(comparer_params_eval, student_eval, utils):
    """
    Check whether student's answer is nonzero and in the span of some given
    vectors.

    comparer_params: A list of vectors

    Usage
    =====

    Use a single vector as comparer_params to test whether student input is
    parallel to a particular vector:
    >>> from mitxgraders import MatrixGrader
    >>> grader = MatrixGrader(
    ...     answers={
    ...         'comparer_params': [
    ...             '[3, x, 1 + i]',
    ...         ],
    ...         'comparer': vector_span_comparer
    ...     },
    ...     variables=['x'],
    ... )
    >>> grader(None, '[3, x, 1 + i]')['ok']
    True
    >>> grader(None, '[9, 3*x, 3 + 3*i]')['ok']
    True
    >>> grader(None, '[9, 3*x, 3 - 3*i]')['ok']
    False

    Complex scale factors work, too:
    >>> grader(None, '(4 + 2*i)*[3, x, 1 + i]')['ok']
    True

    Student input should be nonzero:
    >>> result = grader(None, '[0, 0, 0]')
    >>> expected = {
    ...     'ok': False,
    ...     'grade_decimal': 0.0,
    ...     'msg': 'Input should be a nonzero vector.'
    ... }
    >>> result == expected
    True

    Input shape is validated:
    >>> try:
    ...     grader(None, '5')
    ... except InputTypeError as error:
    ...     print(error)
    Expected answer to be a vector, but input is a scalar

    Multiple vectors can be provided:
    >>> grader = MatrixGrader(
    ...     answers={
    ...         'comparer_params': [
    ...             '[1, 1, 0]',    # v0
    ...             '[0, 1, 2]'     # v1
    ...         ],
    ...         'comparer': vector_span_comparer
    ...     },
    ... )

    The vector 2*v0 + 3i*v1 = [2, 2+3i, 6i] is in the span of v0 and v1:
    >>> grader(None, '[2, 2 + 3*i, 6*i]')['ok']
    True

    The comparer_params should be list of equal-length vectors:
    >>> grader = MatrixGrader(
    ...     answers={
    ...         'comparer_params': [
    ...             '[1, 1, 0]',
    ...             '5'
    ...         ],
    ...         'comparer': vector_span_comparer
    ...     },
    ... )
    >>> try:
    ...     grader(None, '[1, 2, 3]')               # doctest: +ELLIPSIS
    ... except StudentFacingError as error:
    ...     print(error)
    Problem Configuration Error: ...to equal-length vectors
    """

    # Validate the comparer params
    if not are_same_length_vectors(comparer_params_eval):
        raise StudentFacingError('Problem Configuration Error: comparer_params '
            'should be a list of strings that evaluate to equal-length vectors')

    # Validate student input shape
    utils.validate_shape(student_eval, comparer_params_eval[0].shape)

    if utils.within_tolerance(0, np.linalg.norm(student_eval)):
        return {
            'ok': False,
            'grade_decimal': 0,
            'msg': 'Input should be a nonzero vector.'
        }

    # Use ordinary least squares to find an approximation to student_eval
    # that lies within the span of given vectors, then check that the
    # residual-sum is small in comparison to student input.
    column_vectors = np.array(comparer_params_eval).transpose()
    # rcond=-1 uses machine precision for testing singular values
    # In numpy 1.14+, use rcond=None fo this behavior. (we use 1.6)
    ols = np.linalg.lstsq(column_vectors, student_eval, rcond=-1)
    error = np.sqrt(ols[1])

    # Check that error is nearly zero, using student_eval as a reference
    # when tolerance is specified as a percentage
    return is_nearly_zero(error, utils.tolerance, reference=student_eval)

def vector_phase_comparer(comparer_params_eval, student_eval, utils):
    """
    Check that student input equals a given input (to within tolerance), up to
    an overall phase factor.

    comparer_params: [target_vector]

    Usage
    =====

    >>> from mitxgraders import MatrixGrader
    >>> grader = MatrixGrader(
    ...     answers={
    ...         'comparer_params': [
    ...             '[1, exp(-i*phi)]',
    ...         ],
    ...         'comparer': vector_phase_comparer
    ...     },
    ...     variables=['phi'],
    ... )

    >>> grader(None, '[1, exp(-i*phi)]')['ok']
    True
    >>> grader(None, '[exp(i*phi/2), exp(-i*phi/2)]')['ok']
    True
    >>> grader(None, '[i, exp(i*(pi/2 - phi))]')['ok']
    True

    >>> grader(None, '[1, exp(+i*phi)]')['ok']
    False
    >>> grader(None, '[2, 2*exp(-i*phi)]')['ok']
    False

    The comparer_params should be list with a single vector:
    >>> grader = MatrixGrader(
    ...     answers={
    ...         'comparer_params': [
    ...             '[1, 1, 0]',
    ...             '[0, 1, 1]'
    ...         ],
    ...         'comparer': vector_phase_comparer
    ...     },
    ... )
    >>> try:
    ...     grader(None, '[1, 2, 3]')               # doctest: +ELLIPSIS
    ... except StudentFacingError as error:
    ...     print(error)
    Problem Configuration Error: ...to a single vector.
    """
    # Validate that author comparer_params evaluate to a single vector
    if not len(comparer_params_eval) == 1 and is_vector(comparer_params_eval[0]):
        raise StudentFacingError('Problem Configuration Error: comparer_params '
            'should be a list of strings that evaluate to a single vector.')

    # We'll check that student input is in the span as target vector and that
    # it has the same magnitude

    in_span = vector_span_comparer(comparer_params_eval, student_eval, utils)

    expected_mag = np.linalg.norm(comparer_params_eval[0])
    student_mag = np.linalg.norm(student_eval)
    same_magnitude = utils.within_tolerance(expected_mag, student_mag)

    return in_span and same_magnitude