import numpy as np


def felli(x):
    N = x.shape[0]
    if N < 2:
        raise ValueError("dimension must be greater than one")
    exponents = np.arange(N) / (N - 1)
    return float(np.sum((1e6 ** exponents) * (x ** 2)))


def escma():
    #Initialization

    # User defined input parameters
    N = 10                       
    xmean = np.random.rand(N)
    sigma = 0.5
    stopfitness = 1e-10
    stopeval = int(1e3 * N**2)

    # Strategy parameter setting: Selection
    lambda_ = 4 + int(np.floor(3 * np.log(N)))
    mu = lambda_ / 2.0

    # muXone recombination weights
    weights = np.log(mu + 0.5) - np.log(np.arange(1, int(mu) + 1))
    mu = int(np.floor(mu))
    weights = weights / np.sum(weights)
    mueff = np.sum(weights)**2 / np.sum(weights**2)

    # Strategy parameter setting: Adaptation
    cc = (4 + mueff / N) / (N + 4 + 2 * mueff / N)
    cs = (mueff + 2) / (N + mueff + 5)
    c1 = 2 / ((N + 1.3)**2 + mueff)
    cmu = min(1 - c1,
              2 * (mueff - 2 + 1 / mueff) / ((N + 2)**2 + 2 * mueff))
    damps = 1 + 2 * max(0, np.sqrt((mueff - 1) / (N + 1)) - 1) + cs

    # Initialize dynamic (internal) strategy parameters and constants
    pc = np.zeros(N)
    ps = np.zeros(N)
    B = np.eye(N)
    D = np.eye(N)
    C = B @ D @ (B @ D).T
    eigeneval = 0
    chiN = np.sqrt(N) * (1 - 1/(4*N) + 1/(21 * N**2))

    #Generation Loop
    counteval = 0
    arx = np.zeros((N, lambda_))
    arz = np.zeros((N, lambda_))
    arfitness = np.zeros(lambda_)

    while counteval < stopeval:

        # Generate and evaluate lambda offspring
        for k in range(lambda_):
            arz[:, k] = np.random.randn(N)
            arx[:, k] = xmean + sigma * (B @ D @ arz[:, k])
            arfitness[k] = felli(arx[:, k])
            counteval += 1

        # Sort by fitness and compute weighted mean into xmean
        idx = np.argsort(arfitness)
        arfitness = arfitness[idx]
        arindex = idx

        xold = xmean.copy()
        xmean = arx[:, arindex[:mu]] @ weights
        zmean = arz[:, arindex[:mu]] @ weights

        # Cumulation: Update evolution paths
        ps = (1 - cs) * ps + np.sqrt(cs * (2 - cs) * mueff) * (B @ zmean)
        norm_ps = np.linalg.norm(ps)
        hsig = norm_ps / np.sqrt(1 - (1 - cs)**(2 * counteval / lambda_)) / chiN < \
               (1.4 + 2 / (N + 1))
        hsig = 1.0 if hsig else 0.0

        pc = (1 - cc) * pc + hsig * np.sqrt(cc * (2 - cc) * mueff) * (B @ D @ zmean)

        # Adapt covariance matrix C
        BDz = B @ D @ arz[:, arindex[:mu]]
        C = (1 - c1 - cmu) * C \
            + c1 * (np.outer(pc, pc) + (1 - hsig) * cc * (2 - cc) * C) \
            + cmu * BDz @ np.diag(weights) @ BDz.T

        # Adapt step-size sigma
        sigma = sigma * np.exp((cs / damps) * (norm_ps / chiN - 1))

        # Update B and D from C (Eigenzerlegung, O(N^2))
        if counteval - eigeneval > lambda_ / ((c1 + cmu) * N * 10):
            eigeneval = counteval
            # enforce symmetry
            C = (C + C.T) / 2.0
            eigvals, B = np.linalg.eigh(C)
            D = np.diag(np.sqrt(eigvals))

        # Break, if fitness is good enough
        if arfitness[0] <= stopfitness:
            break

        # Escape flat fitness, or better terminate?
        if arfitness[0] == arfitness[int(np.ceil(0.7 * lambda_)) - 1]:
            sigma = sigma * np.exp(0.2 + cs / damps)
            print("warning: flat fitness, consider reformulating the objective")

        print(f"{counteval}: {arfitness[0]}")

    #Final Message
    print(f"{counteval}: {arfitness[0]}")
    xmin = arx[:, arindex[0]]
    return xmin


if __name__ == "__main__":
    xmin = escma()
    print("Bestes gefundenes x:", xmin)
    print("f(xmin) =", felli(xmin))