import json
import argparse
import os.path
from glob import glob
from tqdm import tqdm
from multiprocessing import Pool
import torch
import sys
from functools import partial

sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from post_processors.openai_api_callback import majority_voting_predict
from data.deepseek_math_utils import eval_script


def load_rewards(reward_file, re_index):
    if os.path.exists(reward_file):
        rewards = json.load(open(reward_file))
    else:
        rewards = []
        for i, file in enumerate(sorted(glob(reward_file))):
            print(file)
            sub_rewards = json.load(open(file))
            if re_index:
                for item in sub_rewards:
                    item["index"] = f"{item['index']}_{i}"
            rewards.extend(sub_rewards)

    return rewards


def reward_reduction(ending_logits, reduction: str = "min", norm: bool = True):
    if norm:
        ending_logits = torch.tensor(ending_logits)
        ending_probs = torch.softmax(ending_logits, dim=-1).tolist()
        step_rewards = [step[1] for step in ending_probs]
    else:
        step_rewards = [step[1] for step in ending_logits]
    if reduction == "min":
        reward = min(step_rewards)
    elif reduction == "product":
        reward = 1
        for prob in step_rewards:
            reward *= prob
    elif reduction == "sum":
        reward = sum(step_rewards)
    else:
        raise ValueError(f"Invalid reduction method: {reduction}")
    return reward


def weighted_majority_voting_predict(preds, weights):
    pred2weight = {}
    for pred, weight in zip(preds, weights):
        if pred not in pred2weight:
            pred2weight[pred] = 0
        pred2weight[pred] += weight
    return max(pred2weight, key=pred2weight.get)


def _init(id2reward):
    global _id2reward
    _id2reward = id2reward


def _worker(item, reduction, norm, sc_top_k=None):
    if not item["response"] or not item["pred"] or not item["res"]:
        return {
            "missing": 1,
            "reward_missing": 0,
            "pred_missing": 0,
            "seq_too_long": 0,
            "sorted_results": [],
            "sc_res": False
        }
    unsorted_results = []
    reward_missing = 0
    pred_missing = 0
    seq_too_long = 0
    for i, (resp, pred, r) in enumerate(zip(item["response"], item["pred"], item["res"])):
        resp_id = f"{item['id']}_{i}"
        if resp_id not in _id2reward:
            reward_missing += 1
            continue
        if not pred:
            pred_missing += 1
            continue

        process_rewards = _id2reward[resp_id]
        if len(process_rewards["ending_logits"]) == 0:
            seq_too_long += 1
            continue

        assert resp == process_rewards["response"], f"{resp} \n\n {process_rewards['response']} \n\n ========="

        reward = reward_reduction(process_rewards["ending_logits"], reduction, norm)
        unsorted_results.append((resp, pred, r, reward))
    sorted_results = sorted(unsorted_results, key=lambda x: x[-1], reverse=True)

    sc_top_k_res = {}
    if sc_top_k and sorted_results:
        for k in sc_top_k:
            preds = [r[1] for r in sorted_results[:k]]
            sc_pred = majority_voting_predict(preds)
            if sc_pred != "":
                sc_res = eval_script.eval_math({"prediction": sc_pred, "answer": item["label"]})
            else:
                sc_res = False
            sc_top_k_res[k] = sc_res

    sc_k_res = {}
    if sc_top_k and unsorted_results:
        for k in sc_top_k:
            preds = [r[1] for r in unsorted_results[:k]]
            sc_pred = majority_voting_predict(preds)
            if sc_pred != "":
                sc_res = eval_script.eval_math({"prediction": sc_pred, "answer": item["label"]})
            else:
                sc_res = False
            sc_k_res[k] = sc_res

    weighted_best_of_k = {}
    if sc_top_k and unsorted_results:
        for k in sc_top_k:
            preds = [r[1] for r in unsorted_results[:k]]
            weights = torch.softmax(torch.tensor([r[-1] for r in unsorted_results]), dim=0).tolist()
            best_of_k_pred = weighted_majority_voting_predict(preds, weights)
            if best_of_k_pred != "":
                best_of_k_res = eval_script.eval_math({"prediction": best_of_k_pred, "answer": item["label"]})
            else:
                best_of_k_res = False
            weighted_best_of_k[k] = best_of_k_res

    return {
        "missing": 0,
        "reward_missing": reward_missing,
        "pred_missing": pred_missing,
        "seq_too_long": seq_too_long,
        "sorted_results": sorted_results,
        "sc_res": item["sc_res"],
        "sc_top_k_res": sc_top_k_res,
        "weighted_best_of_k": weighted_best_of_k,
        "sc_k_res": sc_k_res,
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--response_file", type=str, required=True)
    parser.add_argument("--reward_file", type=str, required=True)
    parser.add_argument("--reduction", type=str, default="min")
    parser.add_argument("--raw_logits", action="store_true", default=False)
    parser.add_argument("--num_workers", type=int, default=4)
    parser.add_argument("--re_index", action="store_true", default=False)
    args = parser.parse_args()

    if os.path.exists(args.response_file):
        responses = json.load(open(args.response_file))
    else:
        responses = []
        for file in glob(args.response_file):
            print(file)
            responses.extend(json.load(open(file)))

    print(len(responses[0]["response"]), len(responses[0]["pred"]), len(responses[0]["res"]))
    print(responses[0]["id"])

    rewards = load_rewards(args.reward_file, args.re_index)
    print(rewards[0]['index'])
    id2reward = {item["index"]: item for item in rewards}

    # _k = [1, 3, 5]
    _k = [3, 5, 4, 8, 16, 32, 64, 128]
    prm_pass_at_k = {k: 0 for k in _k}
    missing = 0
    missing_reward = 0
    pred_missing = 0
    seq_too_long = 0
    sc_cnt = 0
    ultimate_results = []
    # for item in tqdm(responses):
    #     sorted_results = []
    #     if not item["response"] or not item["pred"] or not item["res"]:
    #         missing += 1
    #         continue
    #     for i, (resp, pred, r) in enumerate(zip(item["response"], item["pred"], item["res"])):
    #         resp_id = f"{item['id']}_{i}"
    #         if resp_id not in id2reward:
    #             missing_reward += 1
    #             continue
    #         process_rewards = id2reward[resp_id]
    #         if len(process_rewards["ending_logits"]) == 0:
    #             continue
    #
    #         assert resp == process_rewards["response"], f"{resp} \n\n {process_rewards['response']} \n\n ========="
    #
    #         reward = reward_reduction(process_rewards["ending_logits"], args.reduction, not args.raw_logits)
    #         sorted_results.append((resp, pred, r, reward))
    #
    #     if not sorted_results:
    #         continue
    #     sorted_results = sorted(sorted_results, key=lambda x: x[-1], reverse=True)
    #     ultimate_results.append(sorted_results)
    #     for k in _k:
    #         if any([r[2] for r in sorted_results[:k]]):
    #             prm_pass_at_k[k] += 1
    #
    #     if item["sc_res"]:
    #         sc_cnt += 1
    with Pool(args.num_workers, initializer=_init, initargs=(id2reward,)) as pool:
        annotate = partial(_worker, reduction=args.reduction, norm=(not args.raw_logits), sc_top_k=_k)
        results = list(tqdm(pool.imap_unordered(annotate, responses), total=len(responses)))

    sc_top_k = {k: 0 for k in _k}
    weighted_best_of_k = {k: 0 for k in _k}
    sc_k = {k: 0 for k in _k}
    for item in results:
        missing += item["missing"]
        missing_reward += item["reward_missing"]
        pred_missing += item["pred_missing"]
        seq_too_long += item["seq_too_long"]
        sorted_results = item["sorted_results"]

        if item["sc_res"]:
            sc_cnt += 1

        if not sorted_results:
            continue

        ultimate_results.append(sorted_results)
        for k in _k:
            if any([r[2] for r in sorted_results[:k]]):
                prm_pass_at_k[k] += 1

        for k, v in item["sc_top_k_res"].items():
            if v:
                sc_top_k[k] += 1

        for k, v in item["weighted_best_of_k"].items():
            if v:
                weighted_best_of_k[k] += 1

        for k, v in item["sc_k_res"].items():
            if v:
                sc_k[k] += 1

    print(f"Total: {len(responses)}")
    print(f"Missing: {missing}")
    print(f"Missing reward: {missing_reward}")
    print(f"Missing pred: {pred_missing}")
    print(f"Seq too long: {seq_too_long}")
    print(f"SC: {sc_cnt}")
    for k, v in prm_pass_at_k.items():
        print(f"PRM pass at {k}: {v}")
        print(f"PRM pass at {k} rate: {v / len(responses) * 100:.2f}%")
    for k, v in sc_top_k.items():
        print(f"SC pass at {k}: {v}")
        print(f"SC pass at {k} rate: {v / len(responses) * 100:.2f}%")
    print(f"SC rate: {sc_cnt / len(responses) * 100:.2f}%")
    for k, v in weighted_best_of_k.items():
        print(f"Weighted best of k pass at {k}: {v}")
        print(f"Weighted best of k pass at {k} rate: {v / len(responses) * 100:.2f}%")
    for k, v in sc_k.items():
        print(f"SC k pass at {k}: {v}")
        print(f"SC k pass at {k} rate {v / len(responses) * 100:.2f}%")

    json.dump(ultimate_results[:100], open("debug.json", "w"), indent=2)


if __name__ == '__main__':
    main()