import multiprocessing
import random
from collections.abc import Callable
from pathlib import Path

import lmwrapper.caching
import pandas as pd
from lmwrapper.huggingface_wrapper import get_huggingface_lm
from lmwrapper.openai_wrapper import OpenAiModelNames, get_open_ai_lm

from synthegrator.code_problems import CodeProblem
from synthegrator.code_solver import (
    LmBasedSolver,
    LmCodeSolverAutoRegressive,
    ResponseParserSingleEdit,
)
from synthegrator.df_converters import solution_evals_to_df
from synthegrator.execution_threading import solve_and_evaluate_problems
from synthegrator.problem_rendering import (
    PromptRendererSingleEditGeneric,
)
from synthegrator.problem_rendering_insertion_tags import (
    TaggedEditRenderer,
    TaggedEditResponseParser,
)
from synthegrator.synthdatasets.human_eval import yield_human_eval

random.seed(0)

try:
    import torch
except ImportError:
    torch = None

cur_file = Path(__file__).parent.absolute()


def get_experiment_solver_constructors(
    open_ai_models: bool = True,
    hf_models: bool = torch is not None and torch.cuda.is_available(),
    use_edit_renders: bool = False,
) -> dict[str, Callable[[], LmBasedSolver]]:
    """
    Returns a dictionary of solver constructors for all the models
    we want to run in this experiment. The dictionary keys are
    model names. The values are lambda functions that construct
    the solvers. We don't set the values as the actual solvers
    as they might load large models that we don't/can't keep
    all in memory together. So instead by using the lambda
    we can lazily construct the solvers only when they
    are needed.
    """
    solvers = {}
    if use_edit_renders:
        renderer = TaggedEditRenderer()
        response_parser = TaggedEditResponseParser()
    else:
        renderer = PromptRendererSingleEditGeneric()
        response_parser = ResponseParserSingleEdit()

    def construct(lm):
        return LmCodeSolverAutoRegressive(
            lm,
            prompt_renderer=renderer,
            response_parser=response_parser,
            # snap_indent=4,
            include_lm_response=True,
            cache_lm=True,
        )

    if open_ai_models:
        model_names = [
            OpenAiModelNames.gpt_3_5_turbo_instruct,
            OpenAiModelNames.gpt_3_5_turbo,
            #OpenAiModelNames.gpt_4_turbo,
            OpenAiModelNames.gpt_4_1_nano,
            OpenAiModelNames.gpt_4_1,
        ]
        solvers |= {
            # Note: Closure needs an arg so it gets captured
            #   without changing later in the loop.
            str(model_name): (lambda model_name=model_name: construct(
                get_open_ai_lm(
                    model_name,
                    retry_on_rate_limit=True,
                )
            ))
            for model_name in model_names
        }
    if hf_models:
        model_names = [
            "Salesforce/codegen-350M-multi",
            "mistralai/Mistral-7B-v0.1",
            "Qwen/Qwen2.5-Coder-0.5B",
        ]
        solvers |= {
            model_name: (lambda model_name=model_name: construct(
                get_huggingface_lm(
                    model_name,
                    trust_remote_code=True,
                    precision=torch.float16,
                )
            ))
            for model_name in model_names
        }
    return solvers


def run_for_problems_and_serialize(
    problems: list[CodeProblem],
    task_name: str,
    output_directory: Path = cur_file / "results",
    max_limit: int = 500,
    pickle_gzip_whole_solution_eval: bool = False,
    read_file_if_exists: bool = False,
    solver_constructors: dict[
        str,
        Callable[[], LmBasedSolver],
    ] = get_experiment_solver_constructors(),
) -> pd.DataFrame:
    if problems is None:
        raise ValueError("problems parameter cannot be None")

    problems_sample = None

    def get_problems_sample():
        nonlocal problems_sample
        if problems_sample is not None:
            return problems_sample
        if len(problems) > max_limit:
            problems_sample = random.sample(problems, max_limit)
        else:
            problems_sample = problems
        return problems_sample

    thread_goal = max(multiprocessing.cpu_count() // 3, 1)
    print("thread_goal: ", thread_goal)
    lmwrapper.caching.set_cache_dir(
        Path(__file__).parent.absolute() / "lm_wrapper_cache",
    )
    dfs = []
    for solver_name, solver_constructor in solver_constructors.items():
        dest_file = output_directory / f"results_{task_name}_{solver_name.replace('/', '_')}.parquet"
        if read_file_if_exists and dest_file.exists():
            dfs.append(pd.read_parquet(dest_file))
            continue
        print(f"Constructing new solver {solver_name}...")
        solver = solver_constructor()
        print(f"Running {solver_name}...")
        evals = solve_and_evaluate_problems(
            solver=solver,
            problems=get_problems_sample(),
            max_threads_solve=max(min(thread_goal, 3), 1),
            max_threads_eval=thread_goal,
        )
        print(f"Converted {solver_name} to df...")
        new_df = solution_evals_to_df(
            evals,
            solver_key=solver_name,
            pickle_gzip_whole_solution_eval=pickle_gzip_whole_solution_eval,
        )
        print("Converted to df")
        new_df["task_name"] = task_name
        new_df["task_problem_count"] = len(problems)
        for eval in evals:
            assert eval.solution.problem is not None
        new_df.to_parquet(dest_file)
        print("Saved to parquet.", dest_file)
        dfs.append(new_df)
        del solver
    df = pd.concat(dfs)
    scores = df.groupby("solver_key").main_metric__is_success.mean()
    print(scores)
    return df


def get_human_eval_results(
    output_directory: Path = cur_file / "results",
    read_file_if_exists: bool = False,
    max_limit: int = 500,
) -> pd.DataFrame:
    print("Gathering all problems...")
    problems = list(yield_human_eval())
    return run_for_problems_and_serialize(
        problems,
        "human_eval",
        output_directory=output_directory,
        read_file_if_exists=read_file_if_exists,
        max_limit=max_limit,
    )


def main():
    get_human_eval_results()


if __name__ == "__main__":
    main()
