"""Genetic operators: selection, crossover, mutation."""

from __future__ import annotations

import random

from factor_engine.gp.operators import (
    AbsOp,
    Neg,
    RankCS,
    TsDecayLinear,
    TsEMA,
    TsMean,
    TsRank,
    TsSlope,
    TsStd,
    TsZScore,
    ZScoreCS,
    choose_binary_operator,
    generate_random_tree,
    make_binary,
    make_unary,
)


def tree_too_large(tree, max_depth, max_nodes):
    return tree.get_depth() > max_depth or tree.get_size() > max_nodes


def tournament_selection(population, fitnesses, k=5):
    idxs = random.sample(range(len(population)), k)
    best_idx = max(idxs, key=lambda i: fitnesses[i])
    return population[best_idx].clone()


def replace_random_subtree(tree, new_subtree):
    tree = tree.clone()
    nodes = tree.get_nodes()
    if len(nodes) <= 1:
        return new_subtree.clone()
    target = random.choice(nodes)
    if target is tree:
        return new_subtree.clone()
    for node in nodes:
        if hasattr(node, "children"):
            for i, child in enumerate(node.children):
                if child is target:
                    node.children[i] = new_subtree.clone()
                    return tree
    return tree


def mutate(tree, max_init_depth):
    if random.random() < 0.55:
        new_subtree = generate_random_tree(1, max_depth=random.randint(2, max(2, max_init_depth - 1)))
        return replace_random_subtree(tree, new_subtree)
    if random.random() < 0.85:
        child = tree.clone()
        op = random.choice([AbsOp, Neg, RankCS, ZScoreCS, TsMean, TsStd, TsZScore, TsRank, TsDecayLinear, TsEMA, TsSlope])
        return make_unary(op, child)
    return make_binary(choose_binary_operator(), tree.clone(), generate_random_tree(1, max_depth=3))


def crossover(parent1, parent2):
    p1 = parent1.clone()
    donor = random.choice(parent2.clone().get_nodes()).clone()
    return replace_random_subtree(p1, donor)