import shutil
import subprocess
import json
import os
import anthropic
import re
import time
import itertools
import threading
from queue import Queue
from openai import OpenAI

# TODO: Change the API key(s)
api_key_ant = "XXXXXX"
anthropic_client = anthropic.Anthropic(api_key=api_key_ant)
open_ai_client = OpenAI()

# TODO: Set location of MiniZinc distribution
minizinc_distr = "path/to/minizinc/binary"

# TODO: Change path to original encoding
original_enc = "./Social_Golfers/sg_problem.mzn"

# TODO: Change path to folder with all training instances
instance_folder = "./Social_Golfers/sg_training_instances"


global_timeout = 600
init_temperature= 0.1
temperature_inc = 0.1


def run_enc_on_instance(enc, instance, timeout):
    """
    Runs an encoding on an instance, keeping the timeout

    Parameters:
    enc (String): Path to the encoding
    instance (String): Path to the instance
    timeout (float): The timeout in seconds

    Returns:
    The needed time or one of the Strings "timeout", "error" or "unsatisfiable"
    """

    command = "time timeout " + str(timeout) + "s " + minizinc_distr + " --solver Chuffed --random-seed 10 " + str(enc) + " " + str(instance)

    result = subprocess.run(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    output = result.stdout
    errors = result.stderr

    # Test if streamliner returns error, timeout or unsatisfiable result
    if not output:
        errors = errors.split("user")[0]
        if "Command exited with non-zero status 124" in errors:
            return "timeout"
        else:
            return "ERROR: " + errors
    elif "UNSATISFIABLE" in output:
        return "unsatisfiable"
    elif "UNKNOWN" in output:
        return "timeout"
    elif "ERROR" in output:
        return "error"
    else:
        time = errors.split("user")[0]
        if "Warning" in time or "non-zero status" in time:
            time = time.split("\n")[-1]
        return time
    
def run_sl_on_instance(original_time, sl, instance, thread_id):
    """
    Creates a new encoding with the given streamliner and then calls the function run_enc_on_instance to run the new encoding on the instance

    Parameters:
    original_time (float): Time that the original encoding needed to solve the given instance
    sl (String): new streamlining constraint
    instance (String): Path to the instance
    thread_id (int): thread id

    Returns:
    The needed time or one of the Strings "timeout", "error" or "unsatisfiable"
    """

    new_enc = "./streamliner_temp_" + str(thread_id) + ".mzn"

    # Make copy of original file
    shutil.copy(original_enc, new_enc)

    # Open the file in read mode and read the content
    with open(new_enc, 'r') as file:
        lines = file.readlines()

    # Write the updated content to the output file
    for index, line in enumerate(lines):
        if 'constraint' in line:
            break
    else:
        index = None

    if index is not None:
        lines.insert(index, '\n\n' + sl + '\n\n')

        with open(new_enc, 'w') as file:
            file.writelines(lines)

        return run_enc_on_instance(new_enc, instance, original_time)
    else:
        print('"constraint" not found in the file.')

def evaluate_sl(original_times, potential_steamliners, sl, thread_id, result_queue):
    """
    Evaluates the given streamliner on all instances in the instance folder

    Parameters:
    original_times (list): List of times that the original encoding needed to solve the training instances in the given instance folder
    potential_steamliners (list): List to store all potential streamliners
    sl (String): The streamliner constraint
    thread_id (int): thread id
    result_queue (queue): Queue to write the results of the parallel threads

    Returns:
    One of the Strings "improvement", "all instances are unsatisfiable", "worse than original encoding" or the error message
    """
        
    improves_at_least_one = False
    unsat_counter = 0
    curr_times = {}
    for inst in os.listdir(instance_folder):
        original_time = original_times[str(inst)]
        result = run_sl_on_instance(original_time, sl, instance_folder + "/" + inst, thread_id)
        if not ("ERROR" in result or result in ["unsatisfiable", "timeout", "error"]):
            curr_times[str(inst)] = result
            improves_at_least_one = True
        elif "ERROR" in result:
            result_queue.put((thread_id, result))
            return
        else:
            curr_times[str(inst)] = original_time

        if result == "unsatisfiable":
            unsat_counter += 1

    if improves_at_least_one:
        potential_steamliners[sl] = curr_times
        result_queue.put((thread_id, "improvement"))
    elif unsat_counter == len(os.listdir(instance_folder)):
        result_queue.put((thread_id, "all instances are unsatisfiable"))
    else:
        result_queue.put((thread_id, "worse than original encoding"))

def read_and_evaluate_results(original_times, potential_steamliners, message_content):
    """
    Evaluate the results of the LLM and formulates feedback for the LLM

    Parameters:
    original_times (list): List of times that the original encoding needed to solve the training instances in the given instance folder
    potential_steamliners (list): List to store all potential streamliners
    message_content (String): Content of the response of the LLM

    Returns:
    The previous message of the LLM in a nice format and the feedback string for the next LLM call
    """
    # Read the results and put them into the right format
    json_pattern = re.search(r'\{.*\}', message_content, re.DOTALL)
    json_str = json_pattern.group(0)
    json_str = json_str.replace('\\', '\\\\').replace('\\\\"', '\\"')
    try:
        structured_response = json.loads(json_str)
    except:
        print("An exception occurred while parsing the json format")
        print(json_str)
        results = "An exception occurred while parsing your answer to json format. Please change your answer format."
        return message_content, results

    for key, value in structured_response.items():
        structured_response[key] = value.replace('\\\\', '\\').replace('\n', '')

    # Evaluate given streamliners in parallel
    threads = []
    result_queue = Queue()

    for i in range(1, 6):
        streamliner = "streamliner_" + str(i)
        thread = threading.Thread(target=evaluate_sl, args = (original_times, potential_steamliners, structured_response.get(streamliner), i, result_queue))
        threads.append(thread)
        thread.start()

    for thread in threads:
        thread.join()

    sl_results = [0, 0, 0, 0, 0, 0]
    while not result_queue.empty():
        i, result = result_queue.get()
        sl_results[i] = result

    # Formulate the results as answer for the LLM
    results = f"""Here are the results for the evaluation of the streamliners you provided:
        - streamliner_1: {sl_results[1]}.
        - streamliner_2: {sl_results[2]}.
        - streamliner_3: {sl_results[3]}.
        - streamliner_4: {sl_results[4]}.
        - streamliner_5: {sl_results[5]}.
        Based on this feedback, please give 5 new streamliners.
        """
    
    return json.dumps(structured_response), results

def print_results(potential_steamliners, original_times):
    """
    Evaluate the best combination of three streamliners and prints the results

    Parameters:
    potential_steamliners (list): List to store all potential streamliners
    original_times (list): List of times that the original encoding needed to solve the training instances in the given instance folder
    
    Returns:
    List of best combination of three streamliners
    """
    streamliners = []

    for key, value in potential_steamliners.items():
        streamliners.append(key)

    # Create lists of all possible triple combinations
    if len(streamliners) < 3:
        all_triples = list(itertools.combinations(streamliners, len(streamliners)))
    else:
        all_triples = list(itertools.combinations(streamliners, 3))

    # Evaluate the best combination of three streamliners for the picked instances
    max_time_saved = [0]
    best_triple = []
    for curr_triple in all_triples:
        time_saved_by_curr_triple = []
        for inst in os.listdir(instance_folder):
            orig_time = float(original_times[str(inst)])
            curr_time_saved = 0
            for sl in curr_triple:
                curr_value = float(potential_steamliners[sl][inst])
                curr_time_saved = max(curr_time_saved, orig_time - curr_value)
            time_saved_by_curr_triple.append(curr_time_saved)
        if sum(time_saved_by_curr_triple) > sum(max_time_saved):
            max_time_saved = time_saved_by_curr_triple
            best_triple = curr_triple

    print("Best triple: " + str(best_triple))
    return(best_triple)

def find_original_values():
    """
    Check running times of original encoding

    Returns:
    Dictionary with instances names and corresponding running times
    """
    original_times = {}
    for inst in os.listdir(instance_folder):
        original_times[str(inst)] = run_enc_on_instance(original_enc, instance_folder + "/" + inst, 15)

    return original_times


##### Claude3 Opus #####

def run_variant_on_claude(system_request, adaptive):
    """
    Runs the Stream(LLM) approach on Claude

    Parameters:
    system_request (String): The prompt to use
    multi_shot (boolean): whether adaptive approach should be used

    Returns:
    int, float: The sum of a and b.
    """

    original_times = find_original_values()
    
    # Add 10 percent to the threshholds in order to also consider streamliners that are close to the original times 
    # and have the chance to beat them for some large instances
    original_times = {key: str(float(value) * 1.1) for key, value in original_times.items()}

    potential_steamliners = {}
    curr_temperature = init_temperature

    # Start timer
    start_time = time.time()

    # Global timeout of 10 minutes
    while time.time() - start_time < global_timeout:

        # Read original encoding
        enc_content = ""
        with open(original_enc, 'r') as file:
            enc_content = file.read()

        message_list = []
        message_list.append({"role": "user", "content": [{"type": "text", "text": enc_content}]})

        # Send request to LLM
        message = anthropic_client.messages.create(
            model="claude-3-5-sonnet-20240620",
            max_tokens=4000,
            temperature=curr_temperature,
            system=str(system_request),
            messages=message_list
        )
        
        prev_response, results = read_and_evaluate_results(original_times, potential_steamliners, message.content[0].text)

        # MULTI SHOT
        if adaptive:
            for i in range(2):
                if time.time() - start_time < global_timeout:

                    message_list.append({"role": "assistant", "content": [{"type": "text", "text": prev_response}]})
                    message_list.append({"role": "user", "content": [{"type": "text", "text": str(results)}]})

                    message = anthropic_client.messages.create(
                        model="claude-3-5-sonnet-20240620",
                        max_tokens=4000,
                        temperature=curr_temperature,
                        system=str(system_request),
                        messages=message_list
                    )

                    prev_response, results = read_and_evaluate_results(original_times, potential_steamliners, message.content[0].text)

        if curr_temperature + temperature_inc <= 1:
            curr_temperature += temperature_inc

    print_results(potential_steamliners, original_times)


##### GPT-4o #####

def run_variant_on_gpt(system_request, adaptive):
    """
    Runs the Stream(LLM) approach on GPT

    Parameters:
    system_request (String): The prompt to use
    multi_shot (boolean): whether adaptive approach should be used

    Returns:
    int, float: The sum of a and b.
    """

    original_times = find_original_values()
    
    # Add 10 percent to the threshholds in order to also consider streamliners that are close to the original times 
    # and have the chance to beat them for some large instances
    original_times = {key: str(float(value) * 1.1) for key, value in original_times.items()}

    potential_steamliners = {}
    curr_temperature = init_temperature

    # Start timer
    start_time = time.time()

    # Global timeout of 10 minutes
    while time.time() - start_time < global_timeout:

        # Read original encoding
        enc_content = ""
        with open(original_enc, 'r') as file:
            enc_content = file.read()

        message_list = []
        message_list.append({"role": "system", "content": system_request})
        message_list.append({"role": "user", "content": enc_content})

        # Send request to LLM
        message = open_ai_client.chat.completions.create(
            model="gpt-4o-2024-05-13",
            response_format={ "type": "json_object" },
            messages=message_list,
            temperature=curr_temperature
        )

        prev_response, results = read_and_evaluate_results(original_times, potential_steamliners, message.choices[0].message.content)

        # MULTI SHOT
        if adaptive:
            for i in range(2):
                if time.time() - start_time < global_timeout:

                    message_list.append({"role": "assistant", "content": prev_response})
                    message_list.append({"role": "user", "content": str(results)})

                    message = open_ai_client.chat.completions.create(
                        model="gpt-4o-2024-05-13",
                        response_format={ "type": "json_object" },
                        messages=message_list,
                        temperature=curr_temperature
                    )

                    prev_response, results = read_and_evaluate_results(original_times, potential_steamliners, message.choices[0].message.content)

        if curr_temperature + temperature_inc <= 1:
            curr_temperature += temperature_inc

    print_results(potential_steamliners, original_times)

base_prompt =  "Objective: Analyze the given MiniZinc code and suggest five additional constraints to enhance the problem-solving process. These constraints can include streamlining, implied, symmetry-breaking, or dominance-breaking constraints.\n\nSteps:\n\n1. Analyze Content: Read the provided MiniZinc code. Understand the problem being addressed, including its variables, constraints, and optimization goals. \n2. Generate additional Constraints: Based on your analysis, create five unique constraints. These should offer targeted modifications or restrictions designed to reduce the search space effectively. \n3. Always return your constraints as a JSON object, adhering to the structure: {\"streamliner_1\": \"constraint <MiniZinc constraint>\", ..., \"streamliner_5\": \"constraint <MiniZinc constraint>\"}. Your final output should exclusively be the JSON object containing the five constraints. \n4. As a response, you will get feedback for each constraint. Some constraints might lead to errors, timeouts, or unsatisfiable instances. \n5. Use the information provided in the previous step to generate five new, and hopefully better constraints. \n6. Repeat steps 3 to 5 multiple times. \n\nCompliance Rules: \n\n1. Response Format: Your final output should exclusively be the JSON object containing the five constraints, adhering to the structure: {\"streamliner_1\": \" constraint <MiniZinc constraint>\", ..., \"streamliner_5\": \"constraint <MiniZinc constraint>\"}. Do not forget the semicolon at the end of the constraint! \n2. Code Quality: All MiniZinc code provided for the constraints must be syntactically correct and functional. For some functions, you may need to include an additional library. \n3. Creativity: You're encouraged to be innovative in proposing constraints, keeping in mind their purpose: to narrow down the search space efficiently without oversimplifying the problem."
prompt_combinations = "Objective: Analyze the given MiniZinc code and suggest five additional combinations of constraints to enhance the problem-solving process. These constraints can include streamlining, implied, symmetry-breaking, or dominance-breaking constraints.\n\nSteps:\n\n1. Analyze Content: Read the provided MiniZinc code. Understand the problem being addressed, including its variables, constraints, and optimization goals. \n2. Generate additional Combinations of Constraints: Based on your analysis, create five unique combinations of constraints. These should offer targeted modifications or restrictions designed to reduce the search space effectively. \n3. Always return your constraints as a JSON object, adhering to the structure: {\"streamliner_1\": \"constraint <MiniZinc constraint>; ...; constraint <MiniZinc constraint>;\", ..., \"streamliner_5\": \"constraint <MiniZinc constraint>; ...; constraint <MiniZinc constraint>;\"}. Each combination should consist of at least two constraints. Your final output should exclusively be the JSON object containing the five combinations of constraints. \n4. As a response, you will get feedback for each combination of constraints. Some combinations might lead to errors, timeouts, or unsatisfiable instances. \n5. Use the information provided in the previous step to generate five new, and hopefully better combinations of constraints. \n6. Repeat steps 3 to 5 multiple times. \n\nCompliance Rules: \n\n1. Response Format: Your final output should exclusively be the JSON object containing the five combinations of constraints, adhering to the structure: {\"streamliner_1\": \"constraint <MiniZinc constraint>; ...; constraint <MiniZinc constraint>;\", ..., \"streamliner_5\": \"constraint <MiniZinc constraint>; ...; constraint <MiniZinc constraint>;\"}. Do not forget the semicolon at the end of each constraint! \n2. Code Quality: All MiniZinc code provided for the constraints must be syntactically correct and functional. For some functions, you may need to include an additional library. \n3. Creativity: You're encouraged to be innovative in proposing constraints, keeping in mind their purpose: to narrow down the search space efficiently without oversimplifying the problem."

# TODO: Use the functions run_variant_on_gpt(system_request, adaptive) or run_variant_on_claude(system_request, adaptive) depending on wether you want to use GPT-4o or Claude 3.5 Sonnet
# Parameters:
# system_request (String): The prompt to use
# multi_shot (boolean): whether adaptive approach should be used

# run_variant_on_gpt(system_request, adaptive)
# run_variant_on_claude(system_request, adaptive)

# Example
run_variant_on_gpt(base_prompt, False)
