import numpy as np
import math
from humpday.objectives.deapobjectives import schwefel, schaffer, bohachevsky, griewank, rastrigin, shekel, rosenbrock

# Some test objective functions to help guide optimizer choices
# -------------------------------------------------------------
#
# We'll use DEAP's set of groovy benchmarks, and landscapes, swarmpackagepy also
#
# See pretty pictures at https://deap.readthedocs.io/en/master/api/benchmarks.html#deap.benchmarks
# Some hardness assessment is at https://github.com/nathanrooy/landscapes#available-functions-from-single_objective but we'll do our own


## Basis of tricky functions

import datetime

DAY = datetime.datetime.today().day
OFFSET = DAY / 50
POWER = 1 + (DAY % 3) / 3.0
SHIFT = DAY / 100


def smoosh(ui):
    """ Distort the interval to avoid obvious minima and avoid memorization """
    ui_rotate = ui + SHIFT % 1.0
    ui_shift = ui_rotate + SHIFT
    xi = ui_shift ** POWER
    low = SHIFT ** POWER
    high = (1 + SHIFT) ** POWER
    yi = (xi - low) / (high - low)
    return yi ** POWER


def schwefel_on_cube(u: [float]) -> float:
    # https://deap.readthedocs.io/en/master/api/benchmarks.html#deap.benchmarks.schwefel
    u_squished = [1000 * (smoosh(ui) - 0.5) for ui in u]
    try:
        return 0.001 * schwefel(u_squished)[0] / 0.71063
    except Exception as e:
        raise Exception(e)


def griewank_on_cube(u: [float]) -> float:
    # https://deap.readthedocs.io/en/master/api/benchmarks.html#deap.benchmarks.griewank
    u_squished = [1200 * (ui ** 1.1 - 0.5) for ui in u]
    return griewank(u_squished)[0] / 0.532075


def rastrigin_on_cube(u: [float]) -> float:
    # https://deap.readthedocs.io/en/master/api/benchmarks.html#deap.benchmarks.rastrigin
    u_squished = [10.24 * (ui ** 1.1 - 0.5) for ui in u]
    return 0.01 * rastrigin(u_squished)[0] / 0.059697


def bohachevsky_on_cube(u: [float]) -> float:
    # https://deap.readthedocs.io/en/master/api/benchmarks.html#deap.benchmarks.bohachevsky
    u_squished = [10 * (ui ** 1.1 - 0.5) for ui in u]
    return 1.0 + bohachevsky(u_squished)[0]


def rosenbrock_on_cube(u: [float]) -> float:
    # https://deap.readthedocs.io/en/master/api/benchmarks.html#deap.benchmarks.rosenbrock
    u_squished = [200 * (ui ** 1.1 - 0.5) for ui in u]
    return 1 + 0.1 * rosenbrock(u_squished)[0] / 0.008949


def shaffer_on_cube(u: [float]) -> float:
    # https://deap.readthedocs.io/en/master/api/benchmarks.html#deap.benchmarks.schaffer
    u_squished = [200 * (ui ** 1.1 - 0.5) for ui in u]
    return 0.01 * schaffer(u_squished)[0] / (0.1042133 * 0.71809)


def shekel_on_cube(u: [float]) -> float:
    # https://deap.readthedocs.io/en/master/api/benchmarks.html#deap.benchmarks.schaffer

    n_dim = len(u)
    NUMMAX = 15
    A = 10 * np.random.rand(NUMMAX, n_dim)
    C = np.random.rand(NUMMAX)
    u_squished = [800 * (smoosh(ui) - 0.5) for ui in u]
    return 1.2298 - shekel(u_squished, A, C)[0]


## Combinations

def deap_combo1_on_cube(u: [float]) -> float:
    return 0.3 * (schwefel_on_cube(u) + griewank_on_cube(u) + shekel_on_cube(u)) / 1.883


def deap_combo2_on_cube(u: [float]) -> float:
    return 0.5 * (shaffer_on_cube(u) + shekel_on_cube(u)) - 0.1075


def deap_combo3_on_cube(u: [float]) -> float:
    return 0.5 * (rosenbrock_on_cube(u) + bohachevsky_on_cube(u) + shekel_on_cube(u)) / 1.88


DEAP_OBJECTIVES = [schwefel_on_cube, rastrigin_on_cube, griewank_on_cube,
                   bohachevsky_on_cube, rosenbrock_on_cube, shaffer_on_cube, shekel_on_cube,
                   deap_combo1_on_cube, deap_combo2_on_cube, deap_combo3_on_cube]


# By hand...

def rosenbrock_modified_on_cube(u: [float]) -> float:
    """ https://en.wikipedia.org/wiki/Rosenbrock_function """
    u_scaled = [4 * ui - 2 for ui in u]
    if len(u) == 1:
        return (0.25 - u_scaled[0]) ** 2
    else:
        return 5 + 0.001 * np.sum(
            [100 * (ui_plus - ui * ui) + (1 - ui) * (1 - ui) for ui, ui_plus in zip(u_scaled[1:], u_scaled)])


# According to http://infinity77.net/global_optimization/test_functions.html#test-functions-index
# there are some really hard ones
# See https://github.com/andyfaff/ampgo/blob/master/%20ampgo%20--username%20andrea.gavana%40gmail.com/go_benchmark.py
# See also https://arxiv.org/pdf/1308.4008v1.pdf

def damavandi_on_cube(u: [float]) -> float:
    """ A trivial multi-dimensional extension of Damavandi's function """
    return 0.01 * damavandi2(u[0], u[1]) - 0.46


def damavandi2(u1, u2) -> float:
    """ Pretty evil function this one """
    # http://infinity77.net/global_optimization/test_functions_nd_D.html#go_benchmark.Damavandi
    x1 = u1 / 14.
    x2 = u2 / 14.
    numerator = math.sin(math.pi * (x1 - 2.0)) * math.sin(math.pi * (x2 - 2.0))
    denumerator = (math.pi ** 2) * (x1 - 2.0) * (x2 - 2.0)
    factor1 = 1.0 - (abs(numerator / denumerator)) ** 5.0
    factor2 = 2 + (x1 - 7.0) ** 2.0 + 2 * (x2 - 7.0) ** 2.0
    return factor1 * factor2


def paviani_on_cube(u: [float]) -> float:
    # http://infinity77.net/global_optimization/test_functions_nd_P.html#go_benchmark.Paviani
    x = np.array([2.001 + 5.996 * smoosh(ui) for ui in u])

    def safe_np_log(x):
        lb = np.array([1e-6] * len(x))
        xup = np.maximum(x, lb)
        return np.log(xup)

    return float(np.sum(safe_np_log(x - 2) ** 2.0 + safe_np_log(10.0 - x) ** 2.0) - np.prod(x) ** 0.2) / 8.6456


# Landscapes

from landscapes.single_objective import styblinski_tang, zakharov, salomon, rotated_hyper_ellipsoid, qing, michalewicz


def styblinski_tang_on_cube(u: [float]) -> float:
    u_scaled = [10 * (smoosh(ui) - 0.5) for ui in u]
    return 3.3499 + 0.01 * styblinski_tang(u_scaled)


def zakharov_on_cube(u: [float]) -> float:
    u_scaled = [15 * smoosh(ui) - 10 for ui in u]
    return 0.01 * zakharov(u_scaled) / 0.3462


def salomon_on_cube(u: [float]) -> float:
    u_scaled = [200 * smoosh(ui) - 100 for ui in u]
    return salomon(u_scaled) / 3.09999


def rotated_hyper_ellipsoid_on_cube(u: [float]) -> float:
    u_scaled = [2 * 65.536 * smoosh(ui) - 65.536 for ui in u]
    return 0.1 * rotated_hyper_ellipsoid(u_scaled)


def qing_on_cube(u: [float]) -> float:
    u_scaled = [1000 * smoosh(ui) - 500 for ui in u]
    return qing(u_scaled) / 0.01805


def michaelewicz_on_cube(u: [float]) -> float:
    u_scaled = [4 * smoosh(ui) - 2 for ui in u]
    return 1.4439 + 0.1 * michalewicz(u_scaled, m=20)


def landscapes_combo1_on_cube(u: [float]) -> float:
    return (qing_on_cube(u) + michaelewicz_on_cube(u)) / (1.5744 * 1.4688)


def landscapes_combo2_on_cube(u: [float]) -> float:
    return (rotated_hyper_ellipsoid_on_cube(u) + salomon_on_cube(u)) / (6.7555 * 0.82)


def landscapes_combo3_on_cube(u: [float]) -> float:
    return (2 + zakharov_on_cube(u) + styblinski_tang_on_cube(u)) / 4.4329


LANDSCAPES_OBJECTIVES = [styblinski_tang_on_cube, zakharov_on_cube, salomon_on_cube, rotated_hyper_ellipsoid_on_cube,
                         qing_on_cube, michaelewicz_on_cube, landscapes_combo1_on_cube, landscapes_combo2_on_cube,
                         landscapes_combo3_on_cube]


# Some copied from peabox
# https://github.com/stromatolith/peabox/blob/master/peabox/peabox_testfuncs.py
# as that isn't deployed to PyPI as far as I can determine


def ackley_on_cube(u: [float]) -> float:
    # allow parameter range -32.768<=x(i)<=32.768, global minimum at x=(0,0,...,0)
    rescaled_u = [2 * 32.768 * smoosh(ui) - 32.768 for ui in u]
    x = np.asfarray(rescaled_u)
    ndim = len(x)
    a = 20.;
    b = 0.2;
    c = 2. * math.pi
    return (-a * np.exp(-b * np.sqrt(1. / ndim * np.sum(x ** 2))) - np.exp(
        1. / ndim * np.sum(np.cos(c * x))) + a + np.exp(1.)) / 20.0


# Adapted from https://github.com/SISDevelop/SwarmPackagePy/blob/master/SwarmPackagePy/testFunctions.py

def cross_on_cube(u):
    x = [5 * smoosh(ui) - 2.5 for ui in u]
    return round(-0.0001 * (abs(math.sin(x[0]) * math.sin(x[1]) * math.exp(abs(100 -
                                                                               math.sqrt(sum([i ** 2 for i in
                                                                                              x])) / math.pi))) + 1) ** 0.1,
                 7)


def powers_on_cube(u):
    x = [5 * smoosh(ui) - 2.5 for ui in u]
    return sum([abs(x[i]) ** (i + 2) for i in range(len(x))])


def booth_on_cube(u):
    x = [5 * smoosh(ui) - 2.5 for ui in u]
    return sum([abs(x[i]) ** (i + 2) for i in range(len(x))])


def matyas_on_cube(u):
    x = [3 * smoosh(ui) - 1.5 for ui in u]

    def sphere_function(x):
        return sum([i ** 2 for i in x])

    return 0.26 * sphere_function(x) - 0.48 * x[0] * x[1]


def drop_wave_on_cube(u):
    x = [3 * smoosh(ui) - 1.5 for ui in u]

    def sphere_function(x):
        return sum([i ** 2 for i in x])

    return -(1 + math.cos(12 * math.sqrt(sphere_function(x)))) / (0.5 * sphere_function(x) + 2)


SWARM_OBJECTIVES = [cross_on_cube, powers_on_cube,
                    booth_on_cube, matyas_on_cube, drop_wave_on_cube]

A_CLASSIC_OBJECTIVE = rastrigin_on_cube  # Just pick one for testing

MISC_OBJECTIVES = [paviani_on_cube, damavandi_on_cube, rosenbrock_modified_on_cube, ackley_on_cube]

CLASSIC_OBJECTIVES = DEAP_OBJECTIVES + LANDSCAPES_OBJECTIVES + MISC_OBJECTIVES + SWARM_OBJECTIVES

if __name__ == "__main__":
    for objective in CLASSIC_OBJECTIVES:
        objective(u=[0.0, 0.5, 1.0])
        objective(u=[0.0, 0.5, 0.0, 0.0, 1.0])
    print(len(CLASSIC_OBJECTIVES))
