import time

import numpy as np
from numpy import sin, cos, pi, exp, sqrt, abs
from scipy.optimize import rosen


class SimpleQuadratic:

    def fun(self, x):
        return np.dot(x, x)

    def der(self, x):
        return 2. * x

    def hess(self, x):
        return 2. * np.eye(x.size)


class AsymmetricQuadratic:

    def fun(self, x):
        return np.dot(x, x) + x[0]

    def der(self, x):
        d = 2. * x
        d[0] += 1
        return d

    def hess(self, x):
        return 2. * np.eye(x.size)


class SlowRosen:

    def fun(self, x):
        time.sleep(40e-6)
        return rosen(x)


class LJ:
    """
    The Lennard Jones potential

    a mathematically simple model that approximates the interaction between a
    pair of neutral atoms or molecules.
    https://en.wikipedia.org/wiki/Lennard-Jones_potential

    E = sum_ij V(r_ij)

    where r_ij is the cartesian distance between atom i and atom j, and the
    pair potential has the form

    V(r) = 4 * eps * ( (sigma / r)**12 - (sigma / r)**6

    Notes
    -----
    the double loop over many atoms makes this *very* slow in Python.  If it
    were in a compiled language it would be much faster.
    """

    def __init__(self, eps=1.0, sig=1.0):
        self.sig = sig
        self.eps = eps

    def vij(self, r):
        return 4. * self.eps * ((self.sig / r)**12 - (self.sig / r)**6)

    def dvij(self, r):
        p7 = 6. / self.sig * (self.sig / r)**7
        p13 = -12. / self.sig * (self.sig / r)**13
        return 4. * self.eps * (p7 + p13)

    def fun(self, coords):
        natoms = coords.size // 3
        coords = np.reshape(coords, [natoms, 3])
        energy = 0.
        for i in range(natoms):
            for j in range(i + 1, natoms):
                dr = coords[j, :] - coords[i, :]
                r = np.linalg.norm(dr)
                energy += self.vij(r)
        return energy

    def der(self, coords):
        natoms = coords.size // 3
        coords = np.reshape(coords, [natoms, 3])
        energy = 0.
        grad = np.zeros([natoms, 3])
        for i in range(natoms):
            for j in range(i + 1, natoms):
                dr = coords[j, :] - coords[i, :]
                r = np.linalg.norm(dr)
                energy += self.vij(r)
                g = self.dvij(r)
                grad[i, :] += -g * dr/r
                grad[j, :] += g * dr/r
        grad = grad.reshape([natoms * 3])
        return grad

    def get_random_configuration(self):
        rnd = np.random.uniform(-1, 1, [3 * self.natoms])
        return rnd * float(self.natoms)**(1. / 3)


class LJ38(LJ):
    natoms = 38
    target_E = -173.928427


class LJ30(LJ):
    natoms = 30
    target_E = -128.286571


class LJ20(LJ):
    natoms = 20
    target_E = -77.177043


class LJ13(LJ):
    natoms = 13
    target_E = -44.326801


class Booth:
    target_E = 0.
    solution = np.array([1., 3.])
    xmin = np.array([-10., -10.])
    xmax = np.array([10., 10.])

    def fun(self, coords):
        x, y = coords
        return (x + 2. * y - 7.)**2 + (2. * x + y - 5.)**2

    def der(self, coords):
        x, y = coords
        dfdx = 2. * (x + 2. * y - 7.) + 4. * (2. * x + y - 5.)
        dfdy = 4. * (x + 2. * y - 7.) + 2. * (2. * x + y - 5.)
        return np.array([dfdx, dfdy])


class Beale:
    target_E = 0.
    solution = np.array([3., 0.5])
    xmin = np.array([-4.5, -4.5])
    xmax = np.array([4.5, 4.5])

    def fun(self, coords):
        x, y = coords
        p1 = (1.5 - x + x * y)**2
        p2 = (2.25 - x + x * y**2)**2
        p3 = (2.625 - x + x * y**3)**2
        return p1 + p2 + p3

    def der(self, coords):
        x, y = coords
        dfdx = (2. * (1.5 - x + x * y) * (-1. + y) +
                2. * (2.25 - x + x * y**2) * (-1. + y**2) +
                2. * (2.625 - x + x * y**3) * (-1. + y**3))
        dfdy = (2. * (1.5 - x + x * y) * (x) +
                2. * (2.25 - x + x * y**2) * (2. * y * x) +
                2. * (2.625 - x + x * y**3) * (3. * x * y**2))
        return np.array([dfdx, dfdy])


"""
Global Test functions for minimizers.

HolderTable, Ackey and Levi have many competing local minima and are suited
for global minimizers such as basinhopping or differential_evolution.
(https://en.wikipedia.org/wiki/Test_functions_for_optimization)

See also https://mpra.ub.uni-muenchen.de/2718/1/MPRA_paper_2718.pdf

"""


class HolderTable:
    target_E = -19.2085
    solution = [8.05502, 9.66459]
    xmin = np.array([-10, -10])
    xmax = np.array([10, 10])
    stepsize = 2.
    temperature = 2.

    def fun(self, x):
        return - abs(sin(x[0]) * cos(x[1]) * exp(abs(1. - sqrt(x[0]**2 +
                     x[1]**2) / pi)))

    def dabs(self, x):
        """derivative of absolute value"""
        if x < 0: