# This software was developed by employees of the US Naval Research Laboratory (NRL), an
# agency of the Federal Government. Pursuant to title 17 section 105 of the United States
# Code, works of NRL employees are not subject to copyright protection, and this software
# is in the public domain. PyEBSDIndex is an experimental system. NRL assumes no
# responsibility whatsoever for its use by other parties, and makes no guarantees,
# expressed or implied, about its quality, reliability, or any other characteristic. We
# would appreciate acknowledgment if the software is used. To the extent that NRL may hold
# copyright in countries other than the United States, you are hereby granted the
# non-exclusive irrevocable and unconditional right to print, publish, prepare derivative
# works and distribute this software, in any medium, or authorize others to do so on your
# behalf, on a royalty-free basis throughout the world. You may improve, modify, and
# create derivative works of the software or any portion of the software, and you may copy
# and distribute such modifications or works. Modified works should carry a notice stating
# that you changed the software and should note the date and nature of any such change.
# Please explicitly acknowledge the US Naval Research Laboratory as the original source.
# This software can be redistributed and/or modified freely provided that any derivative
# works bear some notice that they are derived from it, and any modified versions bear
# some notice that they have been modified.
#
# Author: David Rowenhorst;
# The US Naval Research Laboratory Date: 21 Aug 2020

"""Optimization of the pattern center (PC) of EBSD patterns."""

import numpy as np
import multiprocessing
import functools
import scipy.optimize as opt
from timeit import default_timer as timer


__all__ = [
    "optimize",
    "optimize_pso",
]

RADEG = 180.0 / np.pi


def _optfunction(PC_i, indexer, banddat):
    tic = timer()
    PC = np.atleast_2d(PC_i)
    result = np.zeros(PC.shape[0])
    # this loop is here because pyswarms expects a vectorized function
    #print(PC.shape)
    for q in range(PC.shape[0]):

        bandnorm = indexer.bandDetectPlan.radonPlan.radon2pole(
            banddat, PC=PC[q,:], vendor=indexer.vendor
        )
        #print(timer() - tic)
        npoints = banddat.shape[0]
        #n_averages = 0
        #average_fit = 0
        #nbands_fit = 0
        #phase = indexer.phaseLib[0]
        nbands = indexer.bandDetectPlan.nBands
        indexdata, banddat = indexer._indexbandsphase( banddat, bandnorm)



        fit = indexdata[-1]['fit']
        nmatch = indexdata[-1]['nmatch']
        average_fit = fit + 1.0*(nbands - nmatch)
        #average_fit = -1.0*(3.0-fit)*nmatch
        whgood = np.nonzero(fit < 90.0)

        n_averages = len(whgood[0])


        if n_averages < 0.9:
            average_fit = 1000
        else:
            average_fit = np.sum(average_fit[whgood[0]]) + 4.0*(nbands+1)*(npoints - n_averages)
            average_fit /= npoints
            #average_fit /= n_averages
            #average_fit *=  (n_averages*(nbands+1) - nbands_fit)/(n_averages*nbands)
        result[q] = average_fit
    #print(timer()-tic)
    return result


def optimize(pats, indexer, PC0=None, batch=False):
    """Optimize pattern center (PC) (PCx, PCy, PCz) in the convention
    of the :attr:`indexer.vendor` with Nelder-Mead.

    Parameters
    ----------
    pats : numpy.ndarray
        EBSD pattern(s), of shape
        ``(n detector rows, n detector columns)``,
        or ``(n patterns, n detector rows, n detector columns)``.
    indexer : pyebsdindex.ebsd_index.EBSDIndexer
        EBSD indexer instance storing all relevant parameters for band
        detection.
    PC0 : list, optional
        Initial guess of PC. If not given, :attr:`indexer.PC` is used.
        If :attr:`indexer.vendor` is ``"EMSOFT"``, the PC must be four
        numbers, the final number being the pixel size.
    batch : bool, optional
        Default is ``False`` which indicates the fit for a set of
        patterns should be optimized using the cumulative fit for all
        the patterns, and one PC will be returned. If ``True``, then an
        optimization is run for each individual pattern, and an array of
        PC values is returned.

    Returns
    -------
    numpy.ndarray
        Optimized PC.

    Notes
    -----
    SciPy's Nelder-Mead minimization function is used with a tolerance
    ``fatol=0.00001`` between each iteration, ending the optimization
    when the improvement is below this value.
    """
    banddat = indexer.bandDetectPlan.find_bands(pats)
    npoints, nbands = banddat.shape[:2]

    if PC0 is None:
        PC0 = indexer.PC
    emsoftflag = False
    if indexer.vendor == "EMSOFT":  # Convert to EDAX for optimization
        emsoftflag = True
        indexer.vendor = "EDAX"
        patDim = np.array(indexer.bandDetectPlan.patDim)
        delta = indexer.PC
        PCtemp = PC0[0:3]
        patdimnorm = (np.array([patDim[1], patDim[0], np.max(patDim[0:2])]))
        PCtemp[0] *= -1.0
        PCtemp[0] += 0.5 * indexer.bandDetectPlan.patDim[1]
        PCtemp[1] += 0.5 * indexer.bandDetectPlan.patDim[0]
        #PCtemp /= indexer.bandDetectPlan.patDim[1]
        PCtemp /= patdimnorm
        PCtemp[2] /= delta[3]
        PC0 = PCtemp

    if not batch:
        PCopt = opt.minimize(
            _optfunction,
            PC0,
            args=(indexer, banddat),
            method="Nelder-Mead",
            options={"fatol": 0.00001}
        )
        PCoutRet = PCopt['x']
    else:
        PCoutRet = np.zeros((npoints, 3))
        for i in range(npoints):
            PCopt = opt.minimize(
                _optfunction,
                PC0,
                args=(indexer, banddat[i, :].reshape(1, nbands)),
                method="Nelder-Mead"
            )
            PCoutRet[i, :] = PCopt['x']

    if emsoftflag:  # Return original state for indexer
        indexer.vendor = "EMSOFT"
        indexer.PC = delta
        patDim = np.array(indexer.bandDetectPlan.patDim)
        patdimnorm = (np.array([patDim[1], patDim[0], np.max(patDim[0:2])]))
        if PCoutRet.ndim == 2:
            newout = np.zeros((npoints, 4))
            PCoutRet[:, 0] -= 0.5
            #PCoutRet[:, :3] *= indexer.bandDetectPlan.patDim[1]
            PCoutRet[:, :3] *= np.atleast_2d(patdimnorm)
            PCoutRet[:, 0] *= -1.0
            PCoutRet[:, 1] -= 0.5 * patDim[0]
            PCoutRet[:, 2] *= delta[3]
            newout[:, :3] = PCoutRet
            newout[:, 3] = delta[3]
            PCoutRet = newout
        else:
            newout = np.zeros(4)
            PCoutRet[0] -= 0.5
            PCoutRet[:3] *= patdimnorm
            #PCoutRet[:3] *= indexer.bandDetectPlan.patDim[1]
            PCoutRet[1] -= 0.5 * patDim[0]
            PCoutRet[0] *= -1.0
            PCoutRet[2] *= delta[3]
            newout[:3] = PCoutRet
            newout[3] = delta[3]
            PCoutRet = newout

    return PCoutRet


def optimize_pso(
    pats,
    indexer,
    PC0=None,
    batch=False,
    search_limit=0.2,
    early_exit = 0.0001,
    nswarmparticles=30,
    pswarmpar=None,
    niter=50,
    return_cost=False,
    verbose=1
):
    """Optimize pattern center (PC) (PCx, PCy, PCz) in the convention
    of the :attr:`indexer.vendor` with particle swarms.

    Parameters
    ----------
    pats : numpy.ndarray
        EBSD pattern(s), of shape
        ``(n detector rows, n detector columns)``,
        or ``(n patterns, n detector rows, n detector columns)``.
    indexer : pyebsdindex.ebsd_index.EBSDIndexer
        EBSD indexer instance storing all relevant parameters for band
        detection.
    PC0 : list, optional
        Initial guess of PC. If not given, :attr:`indexer.PC` is used.
        If :attr:`indexer.vendor` is ``"EMSOFT"``, the PC must be four
        numbers, the final number being the pixel size.
    batch : bool, optional
        Default is ``False`` which indicates the fit for a set of
        patterns should be optimized using the cumulative fit for all
        the patterns, and one PC will be returned. If ``True``, then an
        optimization is run for each individual pattern, and an array of
        PC values is returned.
    search_limit : float, optional
        Default is 0.2 for all PC values, and sets the +/- limit for the
        optimization search.
    early_exit: float, optional
        Default is 0.0001 for all PC values, and sets a value for which
        the optimum is considered converged before the number of iterations
        is reached.  The optimiztion will exit early if the velocity and distance
        of all the swarm particles is less than the early_exit value.
    nswarmparticles : int, optional
        Number of particles in a swarm. Default is 30.
    pswarmpar : dict, optional
        Particle swarm parameters "c1", "c2", and "w" with defaults 3.5,
        3.5, and 0.8, respectively.
    niter : int, optional
        Number of iterations. Default is 50.
    return_costs: bool, optional
        Set to True to return the cost value as well as the optimum fit PC.
    verbose : int, optional
        Whether to print the parameters and progress of the
        optimization (>= 1) or not (< 1). Default is to print.

    Returns
    -------
    numpy.ndarray
        Optimized PC.
    """
    banddat = indexer.bandDetectPlan.find_bands(pats)
    npoints, nbands = banddat.shape[:2]
    if pswarmpar is None:
        #pswarmpar = {"c1": 3.05, "c2": 1.05, "w": 0.8}
        pswarmpar = {"c1": 3.5, "c2": 3.5, "w": 0.8}
    if nswarmparticles is None:
        #nswarmpoints = int(np.array(search_limit).max() * (10.0/0.2))
        nswarmparticles = 30

    nswarmparticles = max(5, nswarmparticles)

    if PC0 is None:
        PC0 = np.asarray(indexer.PC)
    else:
        PC0 = np.asarray(PC0)

    emsoftflag = False
    if indexer.vendor == "EMSOFT":  # Convert to EDAX for optimization
        emsoftflag = True
        indexer.vendor = "EDAX"
        delta = indexer.PC
        PCtemp = PC0[0:3]
        PCtemp[0] *= -1.0
        PCtemp[0] += 0.5 * indexer.bandDetectPlan.patDim[1]
        PCtemp[1] += 0.5 * indexer.bandDetectPlan.patDim[0]
        PCtemp /= indexer.bandDetectPlan.patDim[1]
        PCtemp[2] /= delta[3]
        PC0 = np.array(PCtemp)



    # optimizer = pso.single.GlobalBestPSO(
    #     n_particles=nswarmpoints,
    #     dimensions=3,
    #     options=pswarmpar,
    #     bounds=(PC0 - np.array(search_limit), PC0 + np.array(search_limit)),
    # )
    optimizer = PSOOpt(dimensions=3, n_particles=nswarmparticles,
                       c1=pswarmpar['c1'],
                       c2 = pswarmpar['c2'], w = pswarmpar['w'], hyperparammethod='auto',
                       early_exit=early_exit)

    if not batch:
        # cost, PCoutRet = optimizer.optimize(
        #     _optfunction, niter, indexer=indexer, banddat=banddat
        # )
        cost, PCoutRet = optimizer.optimize(_optfunction, indexer=indexer, banddat=banddat,
                                            start=PC0, bounds=(PC0 - np.array(search_limit), PC0 + np.array(search_limit)),
                                            niter=niter, verbose=verbose)
        costout = cost
        #print(cost)
    else:
        PCoutRet = np.zeros((npoints, 3))
        if verbose >= 1:
            print('', end='\n')
        costout = np.zeros(npoints, dtype=np.float32)
        for i in range(npoints):
            # cost, PCoutRet[i, :] = optimizer.optimize(
            #     _optfunction, niter, indexer=indexer, banddat=banddat[i, :, :]
            # )


            cost, newPC = optimizer.optimize(_optfunction, indexer=indexer,
                                    banddat=banddat[i, :].reshape(1, nbands),
               start=PC0, bounds=(PC0 - np.array(search_limit), PC0 + np.array(search_limit)),
               niter=niter, verbose=0)

            PCoutRet[i, :] = newPC
            costout[i] = cost
            progress = int(round(10 * float(i) / npoints))
            if verbose >= 1:
                print('', end='\r')
                print('PC found: [',
                      '*' * progress, ' ' * (10 - progress), '] ', i + 1, '/', npoints,
                      '  global best:', "{0:.3g}".format(cost),
                      '  PC opt:', np.array_str(PCoutRet[i,:], precision=4, suppress_small=True),
                      sep='', end='')
        if verbose >= 1:
            print('', end='\n')

    if emsoftflag:  # Return original state for indexer
        indexer.vendor = "EMSOFT"
        indexer.PC = delta
        if PCoutRet.ndim == 2:
            newout = np.zeros((npoints, 4))
            PCoutRet[:, 0] -= 0.5
            PCoutRet[:, :3] *= indexer.bandDetectPlan.patDim[1]
            PCoutRet[:, 1] -= 0.5 * indexer.bandDetectPlan.patDim[0]
            PCoutRet[:, 0] *= -1.0
            PCoutRet[:, 2] *= delta[3]
            newout[:, :3] = PCoutRet
            newout[:, 3] = delta[3]
            PCoutRet = newout
        else:
            newout = np.zeros(4)
            PCoutRet[0] -= 0.5
            PCoutRet[:3] *= indexer.bandDetectPlan.patDim[1]
            PCoutRet[1] -= 0.5 * indexer.bandDetectPlan.patDim[0]
            PCoutRet[0] *= -1.0
            PCoutRet[2] *= delta[3]
            newout[:3] = PCoutRet
            newout[3] = delta[3]
            PCoutRet = newout
    if return_cost is False:
        return PCoutRet
    else:
        return PCoutRet, costout

def _file_opt(fobj, indexer, stride=200, groupsz = 3):
    nCols = fobj.nCols
    nRows = fobj.nRows
    pcopt = np.zeros((int(nRows / stride), int(nCols / stride), 3), dtype=np.float32)

    for i in range(int(nRows / stride)):
        ii = i * stride
        print(ii)
        for j in range(int(nCols / stride)):
            jj = j * stride

            pats = fobj.read_data(
                returnArrayOnly=True,
                convertToFloat=True,
                patStartCount=[ii*nCols + jj, groupsz]
            )

            pc = optimize(pats, indexer)
            pcopt[i, j, :] = pc

    return pcopt


class PSOOpt():
    def __init__(self,
                dimensions=3,
                n_particles=50,
                c1 = 2.05,
                c2 = 2.05,
                w = 0.8,
                hyperparammethod = 'static',
                boundmethod = 'bounce',
                early_exit=None):
        self.n_particles = int(n_particles)
        self.dimensions = int(dimensions)
        self.c1 = c1
        self.c2 = c2
        self.c1i = None
        self.c2i = None
        self.w = w
        self.wi = None
        self.hyperparammethod = hyperparammethod
        self.boundmethod = boundmethod
        self.vellimit = None
        self.start = None
        self.bounds = None
        self.range = None
        self.niter = None
        self.pos = None
        self.vel = None
        self.early_exit = early_exit


    def initializeswarm(self, start=None, bounds=None):


        if start is None:
            if bounds is not None:
                start = 0.5*(bounds[0]+bounds[1])
            else:
                start = np.zeros(self.dimensions, dtype=np.float32)

        self.start = start

        if bounds is None:
            bounds = (-1*np.ones(self.dimensions, dtype=np.float32),np.ones(self.dimensions, dtype=np.float32) )

        self.bounds = bounds
        self.range = self.bounds[1] - self.bounds[0]

        self.pos = np.random.uniform(low=bounds[0], high=bounds[1], size=(self.n_particles, self.dimensions))
        self.pos[0, :] = start

        self.vel = np.random.normal(size=(self.n_particles, self.dimensions), loc=0.0, scale=1.0)
        meanv = np.mean(np.sqrt(np.sum(self.vel**2, axis=1)))
        self.vel *= np.sqrt(np.sum(self.range**2))/(20. * meanv)
        self.vellimit = 4*np.mean(np.sqrt(np.sum(self.vel**2, axis=1)))


        self.pbest = np.zeros(self.n_particles) + np.infty
        self.pbest_loc = np.copy(self.pos)
        self.gbest = np.infty
        self.gbest_loc = start




    def updateswarmbest(self, fun2opt, pool, **kwargs):

        val = np.zeros(self.n_particles)

        #tic = timer()
        for part_i in range(self.n_particles):
            temp = self.pos[part_i, :]
            val[part_i] = fun2opt(temp, **kwargs)
        #print(timer()-tic)
        #pos = self.pos.copy()
        #tic = timer()
        #results = pool.map(functools.partial(fun2opt, **kwargs),list(pos) )
        #print(timer()-tic)
        #print(len(results[0]), type(results[0]))
        #print(len(results))
        #val = np.concatenate(results)

        wh_newpbest = np.nonzero(val < self.pbest)[0]

        self.pbest[wh_newpbest] = val[wh_newpbest]
        self.pbest_loc[wh_newpbest, :] = self.pos[wh_newpbest, :]

        wh_minpbest = np.argmin(self.pbest)
        if self.pbest[wh_minpbest] < self.gbest:
            self.gbest = self.pbest[wh_minpbest]
            self.gbest_loc = self.pbest_loc[wh_minpbest, :]


    def updateswarmvelpos(self):

        w = self.wi
        c1 = self.c1i
        c2 = self.c2i
        r1 = np.random.random((self.n_particles,1))
        r2 = np.random.random((self.n_particles,1))
        nvel = self.vel.copy()
        nvel = w * nvel + \
               c1 * r1 * (self.pbest_loc - self.pos) + \
               c2 * r2 * (self.gbest_loc - self.pos)

        mag = np.expand_dims(np.sqrt(np.sum(nvel**2, axis=1)), axis=1)
        wh_toofast = np.nonzero(mag > self.vellimit)[0]
        #print(nvel.shape, wh_toofast.shape, mag.shape)
        nvel[wh_toofast, :] *= self.vellimit/mag[wh_toofast]

        self.vel = nvel
        self.pos += nvel

        self.boundarycheck()





    def boundarycheck(self):

        if str.lower(self.boundmethod) == 'bounce':
            self.boundarybounce()


    def boundarybounce(self):

        lb,ub = self.bounds
        for d in range(self.dimensions):
            wh_under = np.nonzero(self.pos[:,d] < lb[d])[0]
            self.pos[wh_under,d] = lb[d]
            self.vel[wh_under,d] = np.abs(self.vel[wh_under,d])

            wh_over = np.nonzero(self.pos[:, d] > ub[d])[0]
            self.pos[wh_over, d] = ub[d]
            self.vel[wh_over, d] = -1*np.abs(self.vel[wh_over, d])

    def updatehyperparam(self, iter):
        if str.lower(self.hyperparammethod) == 'auto':

            N = float(self.niter)-1
            self.c1i = (self.c1 - self.c1 / 7) * (N-iter)/N + self.c1 / 7.0
            #self.c2i = (self.c1 - self.c1 / 7) * (iter) / N + self.c1 / 7.0
            self.c2i = (self.c2 - self.c2 / 7) * (iter) / N + self.c2 / 7.0
            self.wi = self.w/2 * ((N - iter)/N)**2 + self.w/2

        else:
            self.c1i = self.c1
            self.c2i = self.c2
            self.wi = self.w


        pass
    def printprogress(self, iter):

        progress = int(round(10*float(iter)/self.niter))
        print('',end='\r' )
        print('Progress [',
              '*' * progress, ' '*(10-progress),'] ', iter+1 , '/', self.niter,
              '  global best:', "{0:.3g}".format(self.gbest),
              '  best loc:', np.array_str(self.gbest_loc, precision=4, suppress_small=True),
              sep='', end='')
    def optimize(self, function, start=None, bounds=None, niter=50, verbose = 1, **kwargs):

        self.initializeswarm(start, bounds)
        early_exit = self.early_exit
        if early_exit is None:
           early_exit = -1.0

        #with multiprocessing.get_context("spawn").Pool(min(multiprocessing.cpu_count(), self.n_particles)) as pool:
        pool = None
        if verbose >= 1:
            print('n_particles:', self.n_particles, 'c1:', self.c1, 'c2:', self.c2, 'w:', self.w )

        self.niter = niter
        for iter in range(niter):
            self.updatehyperparam(iter)
            self.updateswarmbest(function, pool, **kwargs)
            if verbose >= 1:
                self.printprogress(iter)
                #print(np.abs(self.vel).max())
            self.updateswarmvelpos()

            if np.abs(self.vel).max() < early_exit:
                d = abs(self.gbest_loc - self.pos)
                    #print(d.max())
                if d.max() < early_exit:
                    break




        #pool.close()
        #pool.terminate()
        final_best = self.gbest
        final_loc = self.gbest_loc
        if verbose >= 1:
            print('', end='\n')
            print("Optimization finished | best cost: {}, best pos: {}".format(
                final_best, final_loc))
            print(' ')
        return final_best, final_loc



