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)
    parser.add_argument("--batch_size", default=10, type=int)
    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("## Parsed Question")[-1].strip()
    strategy = strategy.split("##")[0].strip()
    return strategy


def parse_output(output):
    num_sols = output.count("## Solution")
    if num_sols != 1:
        return None, None
    strategy, solution = output.split("## Solution")
    strategy = clean_strategy(strategy)
    solution = solution.strip()
    return strategy, solution


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}-faith.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 = [], []
    samples = samples[args.num_skips :]
    llm = LLM()
    batch_size = args.batch_size
    print(samples[0]["question"])
    with open(output_path, "w" if args.num_skips == 0 else "a") as f:
        for i in tqdm(range(0, len(samples), batch_size)):
            start = i
            batch_samples = samples[start : start + batch_size]
            batch_messages = [
                [
                    {
                        "role": "system",
                        "content": "You are a helpful expert for math problem solving.",
                    },
                    {
                        "role": "user",
                        "content": prompt.replace("{question}", s["question"]),
                    },
                ]
                for s in batch_samples
            ]
            batch_outputs = asyncio.run(
                llm.achat(
                    batch_messages,
                    model=args.model,
                    stop=["## Question"],
                    max_tokens=args.max_tokens,
                    temperature=args.temperature,
                    num_beams=args.num_seqs,
                )
            )
            print(batch_outputs)
            for s, outputs in zip(batch_samples, batch_outputs):
                parsed_questions, solutions, pred_anss = [], [], []
                for o in outputs:
                    parsed_question, solution = parse_output(o)
                    if parsed_question is None:
                        continue
                    parsed_questions.append(parsed_question)
                    solutions.append(solution)
                    pred_anss.append(extract_answer_math(solution))
                label_ans = s["answer"]
                score = 0
                for p in pred_anss:
                    if compare_ans(p, label_ans):
                        score = 1
                        break
                scores.append(score)
                save_sample = s
                save_sample["generation"] = solutions
                save_sample["parsed_question"] = parsed_questions
                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!")
