import jax
import jax.numpy as jnp
from jax import jit, vmap
import functools


def init_cma_es(mean_init, sigma_init, population_size, mu):
    ''' Initialize evolutionary strategy & learning rates. '''
    n_dim = mean_init.shape[0]
    weights_prime = jnp.array(
        [jnp.log((population_size + 1) / 2) - jnp.log(i + 1)
         for i in range(population_size)])
    mu_eff = ((jnp.sum(weights_prime[:mu]) ** 2) /
               jnp.sum(weights_prime[:mu] ** 2))
    mu_eff_minus = ((jnp.sum(weights_prime[mu:]) ** 2) /
                     jnp.sum(weights_prime[mu:] ** 2))

    # lrates for rank-one and rank-μ C updates
    alpha_cov = 2
    c_1 = alpha_cov / ((n_dim + 1.3) ** 2 + mu_eff)
    c_mu = jnp.minimum(1 - c_1 - 1e-8, alpha_cov * (mu_eff - 2 + 1 / mu_eff)
              / ((n_dim + 2) ** 2 + alpha_cov * mu_eff / 2))
    min_alpha = min(1 + c_1 / c_mu,
                    1 + (2 * mu_eff_minus) / (mu_eff + 2),
                    (1 - c_1 - c_mu) / (n_dim * c_mu))
    positive_sum = jnp.sum(weights_prime[weights_prime > 0])
    negative_sum = jnp.sum(jnp.abs(weights_prime[weights_prime < 0]))
    weights = jnp.where(weights_prime >= 0,
                       1 / positive_sum * weights_prime,
                       min_alpha / negative_sum * weights_prime,)
    weights_truncated = jax.ops.index_update(weights, jax.ops.index[mu:], 0)
    c_m = 1

    # lrate for cumulation of step-size control and rank-one update
    c_sigma = (mu_eff + 2) / (n_dim + mu_eff + 5)
    d_sigma = 1 + 2 * jnp.maximum(0, jnp.sqrt((mu_eff - 1) / (n_dim + 1)) - 1) + c_sigma
    c_c = (4 + mu_eff / n_dim) / (n_dim + 4 + 2 * mu_eff / n_dim)
    chi_d = jnp.sqrt(n_dim) * (
        1.0 - (1.0 / (4.0 * n_dim)) + 1.0 / (21.0 * (n_dim ** 2)))

    # Initialize evolution paths & covariance matrix
    p_sigma = jnp.zeros(n_dim)
    p_c = jnp.zeros(n_dim)
    C, D, B = jnp.eye(n_dim), None, None

    memory = {"p_sigma": p_sigma, "p_c": p_c, "sigma": sigma_init,
              "mean": mean_init, "C": C, "D": D, "B": B,
              "generation": 0}

    params = {"mu_eff": mu_eff,
              "c_1": c_1, "c_mu": c_mu, "c_m": c_m,
              "c_sigma": c_sigma, "d_sigma": d_sigma,
              "c_c": c_c, "chi_d": chi_d,
              "weights": weights,
              "weights_truncated": weights_truncated,
              "pop_size": population_size,
              "n_dim": n_dim,
              "tol_x": 1e-12 * sigma_init,
              "tol_x_up": 1e4,
              "tol_fun": 1e-12,
              "tol_condition_C": 1e14,
              "min_generations": 10}
    return params, memory

@functools.partial(jax.jit, static_argnums=(4, 5))
def sample(rng, memory, B, D, n_dim, pop_size):
    """ Jittable Gaussian Sample Helper. """
    z = jax.random.normal(rng, (n_dim, pop_size)) # ~ N(0, I)
    y = B.dot(jnp.diag(D)).dot(z)               # ~ N(0, C)
    y = jnp.swapaxes(y, 1, 0)
    x = memory["mean"] + memory["sigma"] * y    # ~ N(m, σ^2 C)
    return x


@jax.jit
def eigen_decomposition(C, B, D):
    """ Perform eigendecomposition of covariance matrix. """
    if B is not None and D is not None:
        return C, B, D
    C = (C + C.T) / 2
    D2, B = jnp.linalg.eigh(C)
    D = jnp.sqrt(jnp.where(D2 < 0, 1e-20, D2))
    C = jnp.dot(jnp.dot(B, jnp.diag(D ** 2)), B.T)
    return C, B, D


def check_termination(values, params, memory):
    """ Check whether to terminate CMA-ES loop. """
    dC = jnp.diag(memory["C"])
    C, B, D = eigen_decomposition(memory["C"], memory["B"], memory["D"])

    # Stop if generation fct values of recent generation is below thresh.
    if (memory["generation"] > params["min_generations"]
        and jnp.max(values) - jnp.min(values) < params["tol_fun"]):
        print("TERMINATE ----> Convergence/No progress in objective")
        return True

    # Stop if std of normal distrib is smaller than tolx in all coordinates
    # and pc is smaller than tolx in all components.
    if jnp.all(memory["sigma"] * dC < params["tol_x"]) and np.all(
        memory["sigma"] * memory["p_c"] < params["tol_x"]):
        print("TERMINATE ----> Convergence/Search variance too small")
        return True

    # Stop if detecting divergent behavior.
    if memory["sigma"] * jnp.max(D) > params["tol_x_up"]:
        print("TERMINATE ----> Stepsize sigma exploded")
        return True

    # No effect coordinates: stop if adding 0.2-standard deviations
    # in any single coordinate does not change m.
    if jnp.any(memory["mean"] == memory["mean"] + (0.2 * memory["sigma"] * jnp.sqrt(dC))):
        print("TERMINATE ----> No effect when adding std to mean")
        return True

    # No effect axis: stop if adding 0.1-standard deviation vector in
    # any principal axis direction of C does not change m.
    if jnp.all(memory["mean"] == memory["mean"] + (0.1 * memory["sigma"]
                                * D[0] * B[:, 0])):
        print("TERMINATE ----> No effect when adding std to mean")
        return True

    # Stop if the condition number of the covariance matrix exceeds 1e14.
    condition_cov = jnp.max(D) / jnp.min(D)
    if condition_cov > params["tol_condition_C"]:
        print("TERMINATE ----> C condition number exploded")
        return True
    return False


import jax.numpy as jnp


def init_logger(top_k, num_params):
    evo_logger = {"top_values": jnp.zeros(top_k) + 1e10,
                  "top_params": jnp.zeros((top_k, num_params)),
                  "log_top_1": [],
                  "log_top_mean": [],
                  "log_top_std": [],
                  "log_gen_1": [],
                  "log_gen_mean": [],
                  "log_gen_std": [],
                  "log_sigma": [],
                  "log_gen": []}
    return evo_logger


def update_logger(evo_logger, x, fitness, memory, top_k, verbose=False):
    """ Helper function to keep track of top solutions. """
    # Check if there are solutions better than current archive
    vals = jnp.hstack([evo_logger["top_values"], fitness])
    params = jnp.vstack([evo_logger["top_params"], x])
    concat_top = jnp.hstack([jnp.expand_dims(vals, 1), params])
    sorted_top = concat_top[concat_top[:, 0].argsort()]

    # Importantly: Params are stored as flat vectors
    evo_logger["top_values"] = sorted_top[:top_k, 0]
    evo_logger["top_params"] = sorted_top[:top_k, 1:]
    evo_logger["log_top_1"].append(evo_logger["top_values"][0])
    evo_logger["log_top_mean"].append(jnp.mean(evo_logger["top_values"]))
    evo_logger["log_top_std"].append(jnp.std(evo_logger["top_values"]))
    evo_logger["log_gen_1"].append(jnp.min(fitness))
    evo_logger["log_gen_mean"].append(jnp.mean(fitness))
    evo_logger["log_gen_std"].append(jnp.std(fitness))
    evo_logger["log_sigma"].append(memory["sigma"])
    evo_logger["log_gen"].append(memory["generation"])

    if verbose:
        print(evo_logger["log_gen"][-1], evo_logger["top_values"])
    return evo_logger


def rosenbrock_d_dim(x):
    ''' D-Dim. Rosenbrock Fct. Evaluation. '''
    x_i, x_sq, x_p = x[:-1], x[:-1]**2, x[1:]
    return jnp.sum((1 - x_i)**2 + 100*(x_p-x_sq)**2)


def himmelblau_fct(x):
    ''' 2-dim. Himmelblau function. '''
    return (x[0]**2 + x[1] - 11)**2 + (x[0] + x[1]**2 - 7)**2


def six_hump_camel_fct(x):
    ''' 2-dim. 6-Hump Camel function. '''
    p1 = (4 - 2.1*x[0]**2 + x[0]**4/3)*x[0]**2
    p2 = x[0] * x[1]
    p3 = (-4 + 4*x[1]**2)*x[1]**2
    return p1 + p2 + p3


# Toy Problem Evaluation Batch-Jitted Versions
batch_rosenbrock = jit(vmap(rosenbrock_d_dim, in_axes=(0,),
                            out_axes=0))

batch_himmelblau = jit(vmap(himmelblau_fct, in_axes=(0,),
                            out_axes=0))

batch_hump_camel = jit(vmap(six_hump_camel_fct, in_axes=(0,),
                            out_axes=0))
