import json
import os
import pprint
import asyncio
from datetime import datetime
from time import sleep
from tqdm import tqdm
import argparse
from collections import Counter
from distutils.util import strtobool
import multiprocessing as mp


from src.evol.data_utils import load_data
from src.evol.openai_backend import call_chatgpt, LLM
from src.utils.data_utils import extract_answer_math, extract_answer_number
from src.utils.code_utils import execute_tora
from src.utils.math_utils import compare_ans, vote


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--verbose", action="store_true")
    parser.add_argument("--prompt_path", default=None, type=str)
    parser.add_argument("--strategy_path", default=None, type=str)
    parser.add_argument("--dataset", default="gsm", type=str)
    parser.add_argument("--data_path", default=None, type=str)
    parser.add_argument("--model", default="gpt-3.5-turbo", type=str)
    parser.add_argument("--temperature", default=0.0, type=float)
    parser.add_argument("--top_p", default=1.0, type=float)
    parser.add_argument("--max_tokens", default=1024, type=int)
    parser.add_argument("--num_seqs", default=1, type=int)
    parser.add_argument("--num_skips", default=0, type=int)
    parser.add_argument("--input_col", default="question", type=str)
    parser.add_argument("--output_col", default="answer", type=str)
    parser.add_argument("--max_iter", default=3, type=int)
    parser.add_argument("--num_process", default=1, type=int)
    parser.add_argument("--output_path", default=None, type=str)
    args = parser.parse_args()
    return args

def load_prompt(prompt_path):
    with open(prompt_path, 'r', encoding='utf-8') as fp:
        prompt = fp.read().strip()
    return prompt

def stop_tora(result):
    if "\\boxed" in result:
        return True
    return False

def clean_strategy(strategy):
    strategy = strategy.split("## Strategy")[-1].strip()
    strategy = strategy.split("##")[0].strip()
    return strategy

def main(args, samples, idx):
    # load prompt
    prompt = load_prompt(args.prompt_path)
    if idx <= 0:
        print(prompt)
    os.makedirs(f"result/{args.model}/{args.dataset}", exist_ok=True)
    if args.output_path is None:
        output_path = f"result/{args.model}/{args.dataset}/t{args.temperature}_n{args.num_seqs}-train.jsonl"
    else:
        output_path = args.output_path
    print("%" * 30, "Tora", "%" * 30)
    print("Start PID %d and save to %s" % (os.getpid(), output_path))

    if idx != -1:
        output_path = output_path.replace(".jsonl", f"_{idx}.jsonl")
    save_samples, scores = [], []

    llm = LLM()

    with open(output_path, "w" if args.num_skips == 0 else "a") as f:
        pbar = tqdm(
            samples[args.num_skips:], initial=args.num_skips, total=len(samples)
        )
        for i, x in enumerate(pbar):
            question = x[args.input_col]
            strategies = llm.chat(messages=[{"role": "system", "content": "You are a helpful expert for math problem solving."}, {"role": "user", "content": prompt}], model=args.model, stop=["## Question"], max_tokens=args.max_tokens, temperature=args.temperature, num_beams=args.num_seqs)
            strategies = [clean_strategy(s) for s in strategies]
            print(strategies)
            strategies = list(set(strategies))
            outputs = []
            batch_messages = []
            for s in strategies:
                tora_prompt = f"{prompt}\n{s.strip()}\n\n## Question\n{question}\n## Solution"
                batch_messages.append([{"role": "system", "content": "You are a helpful expert for math problem solving."}, {"role": "user", "content": tora_prompt}])
            outputs = asyncio.run(llm.achat(batch_messages, model=args.model, stop=["## Strategy"], max_tokens=args.max_tokens, temperature=args.temperature, num_beams=1))
            outputs = [o[0] for o in outputs]

            pred_anss = [extract_answer_math(o) for o in outputs]
            label_ans = x["answer"]
            score = 0
            for p in pred_anss:
                if compare_ans(p, label_ans):
                    score = 1
                    break
            scores.append(score)
            save_sample = x
            save_sample["generation"] = outputs
            save_sample["strategy"] = strategies
            save_sample["pred_answers"] = pred_anss
            save_sample["label_answer"] = label_ans
            save_sample["score"] = score
            save_samples.append(save_sample)
            f.write(json.dumps(save_sample, ensure_ascii=False, indent=4) + "\n")
            f.flush()
    print(f"Accuracy - {sum(scores) / len(scores)}")

if __name__ == "__main__":
    args = parse_args()
    samples = load_data(args.dataset, args.data_path)
    if args.num_process == 1:
        main(args, samples, idx=-1)
    else:
        num_each_split = int(len(samples) / args.num_process)
        p = mp.Pool(args.num_process)
        for idx in range(args.num_process):
            start = idx * num_each_split
            if idx == args.num_process - 1:
                end = max((idx + 1) * num_each_split, len(samples))
            else:
                end = (idx + 1) * num_each_split
            split_data = samples[start:end]
            p.apply_async(main, args=(args, split_data, idx,))
        p.close()
        p.join()
        print("All of the child processes over!")