import numpy as np
import operator
import random
from deap import base, creator, gp, tools, algorithms

# Real pyramid dimensions
base_length = 230.4
height = 146.6

# Calculating other constants
pyramid_dimensions={
    diagonal:           np.sqrt(2 * base_length**2),
    apothem:            np.sqrt((base_length / 2)**2 + height**2),
    slant_height:       np.sqrt(height**2 + (base_length / 2)**2),
    base_perimeter:     4 * base_length,
    base_area:          base_length**2,
    base_angle:         np.arctan(height / (base_length / 2)),
    side_area:          2 * base_length * slant_height,
    side_angle:         np.arctan((base_length / 2) / height),
    volume:             (1 / 3) * base_area * height,
    corner_edge_length: np.sqrt(height**2 + diagonal**2 / 4),
}

# Target value (e - 1)
target = np.e - 1

# For cutsom gen_tree
MAX_TREE_DEPTH = 2


# Create types
creator.create("FitnessMin", base.Fitness, weights=(-1.0,))
creator.create("Individual", gp.PrimitiveTree, fitness=creator.FitnessMin)

# Initialize the primitive set for pyramid constants
pset = gp.PrimitiveSetTyped("MAIN", [float] * 12, float)
pset.addPrimitive(operator.add, [float, float], float)
pset.addPrimitive(operator.sub, [float, float], float)
pset.addPrimitive(operator.mul, [float, float], float)
pset.addPrimitive(operator.truediv, [float, float], float)

# Primitive set for integers
for i in range(1, 11):
    pset.addTerminal(i, name=f"N{i}", ret_type=float)

# Rename arguments for easier understanding
args = {}
for i in range(len(pyramid_dimensions)):


pset.renameArguments(ARG0='base_length', ARG1='height', ARG2='diagonal', ARG3='apothem', ARG4='slant_height', ARG5='base_perimeter', ARG6='base_area', ARG7='base_angle', ARG8='side_area', ARG9='side_angle', ARG10='volume', ARG11='corner_edge_length')

# Custom tree generation that use pyramid constants and integers
def gen_tree(pset, max_depth, current_depth):
    if random.random() > 0.5:
        prim = random.choice(pset.primitives[pset.ret])
    else:
        prim = random.choice(pset.terminals[pset.ret])

    if isinstance(prim, gp.Primitive):
        node = gp.PrimitiveTree([prim])
        for _ in range(prim.arity):
            child_tree = gen_tree(pset, max_depth, current_depth + 1)
            node.extend(child_tree)
    else:
        node = gp.PrimitiveTree([prim])

    return node


# Define the fitness function
def eval_fit(individual, *consts):
    func = gp.compile(individual, pset)
    error = 0
    try:
        result = func(*consts)
        error = abs(result - np.pi) # (np.e - 1)
    except (ZeroDivisionError, OverflowError, ValueError, TypeError):
        error = float('inf')
    return [error]

# Create the toolbox
toolbox = base.Toolbox()
toolbox.register("expr", gp.genHalfAndHalf, pset=pset, min_=1, max_=2)
toolbox.register("individual", tools.initIterate, creator.Individual, lambda: gen_tree(pset, MAX_TREE_DEPTH, 0))
# toolbox.register("individual", tools.initIterate, creator.Individual, toolbox.expr)
toolbox.register("population", tools.initRepeat, list, toolbox.individual)
toolbox.register("compile", gp.compile, pset=pset)
toolbox.register("evaluate", eval_fit, base_length, height, diagonal, apothem, slant_height, base_perimeter, base_area, base_angle, side_area, side_angle, volume, corner_edge_length)
toolbox.register("select", tools.selTournament, tournsize=3)
toolbox.register("mate", gp.cxOnePoint)
toolbox.register("expr_mut", gp.genFull, min_=0, max_=2)
toolbox.register("mutate", gp.mutUniform, expr=toolbox.expr_mut, pset=pset)

# Genetic algorithm parameters
POPULATION_SIZE = 1000
N_GENERATIONS = 100
CROSSOVER_PROB = 0.8
MUTATION_PROB = 0.2

# Main function
def main():
    random.seed(42)
    pop = toolbox.population(n=POPULATION_SIZE)
    hof = tools.HallOfFame(1)

    stats_fit = tools.Statistics(lambda ind: ind.fitness.values)
    stats_size = tools.Statistics(len)
    mstats = tools.MultiStatistics(fitness=stats_fit, size=stats_size)
    mstats.register("avg", np.mean)
    mstats.register("std", np.std)
    mstats.register("min", np.min)
    mstats.register("max", np.max)

    pop, log = algorithms.eaSimple(pop, toolbox, cxpb=CROSSOVER_PROB, mutpb=MUTATION_PROB, ngen=N_GENERATIONS, stats=mstats, halloffame=hof, verbose=True)

    best_ind = hof[0]
    best_func = gp.compile(expr=best_ind, pset=pset)
    best_result = best_func(base_length, height, diagonal, apothem, slant_height, base_perimeter, base_area, base_angle, side_area, side_angle, volume, corner_edge_length)

    print("Best individual: ", best_ind)
    print("Best result: ", best_result)
    print("Error: ", abs(best_result - target))

    return pop, log, hof

if __name__ == "__main__":
    main()