from time import time

from gbso.cpu.state import CPUState
from gbso.program.test_case import TestCase
from gbso.cpu.insn import *
from gbso.cpu.regs import R8

from gbso.optimize import cost, optimize
from gbso.program.program import Program

prgm = Program(
    insns=[
        INC_R(R8.A),
        INC_R(R8.A),
        INC_R(R8.A),
        INC_R(R8.A),
    ],
)

outputs = [R8.A]
test_cases = [TestCase()]
max_size = 4
num_iters = 1_000_000  # 10_000_000

initial_cost = cost(prgm, test_cases, outputs, prgm)
initial_cycles = prgm.perf()

prgm.display()

print(f"Cost: {initial_cost}")
print(f"Cycles: {initial_cycles}")

start_time = time()

optimized_prgm = optimize(
    prgm,
    max_size=max_size,
    test_cases=test_cases,
    outputs=outputs,
    num_iters=num_iters,
)

end_time = time()

optimized_prgm.display()

optimized_cost = cost(prgm, test_cases, outputs, optimized_prgm)
optimized_cycles = optimized_prgm.perf()

print(f"Cost: {optimized_cost}")
print(f"Cycles: {optimized_cycles}")


optimized_cost = cost(prgm, test_cases, outputs, optimized_prgm)
optimized_cycles = optimized_prgm.perf()

print(f"Cost: {optimized_cost}")
print(f"Cycles: {optimized_cycles}")


optimized_cost = cost(prgm, test_cases, outputs, optimized_prgm)
optimized_cycles = optimized_prgm.perf()

print(f"Cost: {optimized_cost}")
print(f"Cycles: {optimized_cycles}")

print(f"Took {round(end_time - start_time, 3)} seconds")