import sys
sys.path.append(".")
from utils.complify import mutate
from utils.solve import solve 
import time
import utils.gpt as gpt

statement = """
(declare-fun initial_weight () Int)
(declare-fun lost_weight () Int)
(declare-fun gained_weight1 () Int)
(declare-fun lost_weight2 () Int)
(declare-fun gained_weight2 () Int)
(declare-fun final_weight () Int)
(assert (= initial_weight 1))
(assert (= lost_weight 1))
; (gained_weight1 = (lost_weight * 78))
(assert (= gained_weight1 (* lost_weight 78)))
; (lost_weight2 = ((lost_weight * 3) + 198))
(assert (= lost_weight2 (+ (* lost_weight 3) 198)))
; (gained_weight2 = (lost_weight * 124))
(assert (= gained_weight2 (* lost_weight 124)))
; (final_weight = (((gained_weight2 + gained_weight1 + initial_weight) - (lost_weight2 + lost_weight)) * 162))
(assert (= final_weight (* (- (+ gained_weight2 gained_weight1 initial_weight) (+ lost_weight2 lost_weight)) 162)))
(check-sat)
(get-value (final_weight ))
"""

if __name__ == "__main__":
    t0 = time.time()
    # print(statement)
    ok, solution = solve(statement)
    print(ok, solution)
    new_statement = mutate(statement, tactic='mcmc', rm_aux=True, \
                                                allow_real=False, simplify=True, rename=True)
    print(new_statement)
    ok, solution = solve(new_statement)
    t1 = time.time()
    print(t1-t0, solution)
        
    # client = gpt_utils.set_api_key(1)
    # prompt = f"Translate the math problem formulated with SMT-LIB back to a natural language problem. Please ensure the problem to be a math word problem. {statement}"
    # natural_problem = gpt_utils.gpt4_response(client, prompt, [], top_p=0.9)
    # print(natural_problem)

    # client = gpt_utils.set_api_key(1)
    # prompt = f"Solve the following math word problem. Let us think step by step. {natural_problem}"
    # solution = gpt_utils.gpt4_response(client, prompt, [], top_p=0.9)
    # print(solution)