import numpy as np

from gp import Chromosome, ops
from gp.population import ramped_initialization
from gp.primitive import Var

operations = ops.ALL
terminals = [Var(f"x{i}") for i in range(1, 9)]

chrom = Chromosome.full_init(terminals, operations, max_depth=3)
print("Depth:", chrom.root.get_subtree_depth())
print("Formula:", chrom)
print("Tree:\n", chrom.root.to_str_tree())

values = [
    np.array([1.0]),
    np.array([2.0]),
    np.array([3.0]),
    np.array([4.0]),
    np.array([5.0]),
    np.array([6.0]),
    np.array([7.0]),
    np.array([8.0]),
]
context = {var: value for var, value in zip(terminals, values)}
print("Value for ", values, ":", chrom.root.eval(context))

population = ramped_initialization(
    5,
    [3, 4, 5, 6, 7, 8],
    terminals,
    operations,
)
print("Population size:", len(population))
print("Population:")
[print(str(chrom)) for chrom in population]