import numpy as np
from pymoo.core.problem import ElementwiseProblem
from pymoo.core.crossover import Crossover
from pymoo.core.mutation import Mutation
from pymoo.core.sampling import Sampling
from pymoo.algorithms.soo.nonconvex.ga import GA
from pymoo.optimize import minimize
import math

class MySampling(Sampling):

    def _do(self, problem, n_samples, **kwargs):
        """Generate population
        n_samples: The size of the population
        """
        X = np.full((n_samples, problem.n_var), False, dtype=bool)

        for k in range(n_samples):
            I = np.random.permutation(problem.n_var)[:problem.n_max]
            X[k, I] = True
        return X


class BinaryCrossover(Crossover):
    def __init__(self):
        super().__init__(2, 1)

    def _do(self, problem, X, **kwargs):
        """
        Implement crossover operation
        :param problem:
        :param X: Shape [2, n_samples, n_vars] (2 because there are two parents)
        :param kwargs:
        :return:
        """
        n_parents, n_matings, n_var = X.shape

        _X = np.full((self.n_offsprings, n_matings, problem.n_var), False)

        for k in range(n_matings):
            p1, p2 = X[0, k], X[1, k]

            both_are_true = np.logical_and(p1, p2)
            _X[0, k, both_are_true] = True

            n_remaining = problem.n_max - np.sum(both_are_true)

            I = np.where(np.logical_xor(p1, p2))[0]

            S = I[np.random.permutation(len(I))][:n_remaining]
            _X[0, k, S] = True

        return _X


class MyMutation(Mutation):
    def _do(self, problem, X, **kwargs):
        """Implement mutation operation here
        :param problem:
        :param X: Shape [n_samples*2, n_vars]
        :param kwargs:
        :return:
        """
        for i in range(X.shape[0]):
            is_false = np.where(np.logical_not(X[i, :]))[0]
            is_true = np.where(X[i, :])[0]
            X[i, np.random.choice(is_false)] = True
            X[i, np.random.choice(is_true)] = False

        return X

class SubsetProblem(ElementwiseProblem):
    def __init__(self, L, n_max, target):
        super().__init__(n_var=len(L), n_obj=1, n_constr=1)
        self.L = L
        self.n_max = n_max
        self.target = target

    def _evaluate(self, x, out, *args, **kwargs):
        """
        :param x: Bool array
        :param out:
        :param args:
        :param kwargs:
        :return:
        """
        out["F"] = (np.sum(self.L[x]) - self.target)**2
        out["G"] = (self.n_max - np.sum(x)) ** 2

########################################
# create the actual problem to be solved
np.random.seed(1)
L = np.array(range(100))
np.random.shuffle(L)
n_max = 10
target = 398

print('Number of combinations C(100,10): ', math.comb(100, 10))

problem = SubsetProblem(L, n_max, target)

algorithm = GA(
    pop_size=100,
    sampling=MySampling(),
    crossover=BinaryCrossover(),
    mutation=MyMutation(),
    eliminate_duplicates=True)

res = minimize(problem,
               algorithm,
               ('n_gen', 60),
               seed=1,
               verbose=True)
print('Array L: ')
print(L)
print("Function value: %s" % res.F[0])
print("Subset:", np.where(res.X)[0])
selected_values = L[np.where(res.X)[0]]
print('Sum: ', np.sum(selected_values))