import asyncio
import datetime
import json
import logging
import os
import tempfile
import textwrap
import time
from pathlib import Path
from typing import Any, Tuple

from huggingface_hub import hf_hub_download
from huggingface_hub.utils import HfHubHTTPError
from inspect_ai import Task, task
from inspect_ai.dataset import MemoryDataset, Sample

# For recording token counts...
from inspect_ai.model import ModelUsage  # for token count
from inspect_ai.model import (
    ChatMessageAssistant,
    ChatMessageSystem,
    ChatMessageUser,
    GenerateConfig,
    get_model,
)
from inspect_ai.scorer import Score, Scorer, Target, accuracy, scorer, stderr
from inspect_ai.solver import TaskState, use_tools

from astabench.evals.utils import not_implemented_solver
from astabench.util.model import record_model_usage_with_inspect
from astabench.tools import python_session

logger = logging.getLogger(__name__)

# viewable at https://huggingface.co/datasets/allenai/asta-bench
HF_DATASET_DIR = "allenai/asta-bench"
HF_CACHED_RESULTS_DIR = "allenai/asta-bench-solver-data"
RUBRIC_SCORING_MODEL = "claude-3-7-sonnet-20250219"
NULL_SCORE = {
    "overall_score": {
        "average_required_all3": 0,
        "average_optional_all3": 0,
        "average_required_atleast1": 0,
        "average_optional_atleast1": 0,
        "average_required_overall": 0,
        "average_optional_overall": 0,
    }
}

### ======================================================================
###		TASK DEFINITION
### ======================================================================

DEV_DATASET_JSON_FILE = "tasks/e2e_discovery/dev_may_2025.json"
TEST_DATASET_JSON_FILE = "tasks/e2e_discovery/test_may_2025.json"

DEV_HARD_DATASET_JSON_FILE = "tasks/e2e_discovery/dev_hard_jun_2025.json"
TEST_HARD_DATASET_JSON_FILE = "tasks/e2e_discovery/test_hard_jun_2025.json"


def _create_e2e_discovery_task(
    split: str, dataset_file: str, dataset_name: str
) -> Task:
    """Helper function to create e2e_discovery tasks with common configuration."""
    timestamp = datetime.datetime.now().strftime("%m-%d-%Y_%H.%M")
    results_dir = tempfile.mkdtemp(prefix=f"e2e_{split}_" + timestamp + "_")
    print_and_log(f"(Score details can be viewed in {results_dir})")

    raw_data = get_hf_file_contents(dataset_file, repo_id=HF_DATASET_DIR)
    dataset = MemoryDataset(
        samples=[json_to_sample(x, results_dir, dataset_name) for x in raw_data]
    )

    return Task(
        dataset=dataset,
        # Solver that just tells the user to specify a real solver via cli
        solver=not_implemented_solver(),  # placeholder, will be updated in a real call
        setup=[use_tools(python_session())],
        scorer=score_rubric(),
        sandbox=(
            "docker",
            (
                resources.files("astabench.util.sandbox") / "sandbox_compose.yaml"
            ).as_posix(),
        ),
    )


# ----------


def json_to_sample(idea: dict[str, str], results_dir: str, dataset_name: str) -> Sample:
    sample_id = idea["id"]
    task = idea["problem_description"]
    rubrics = idea["evaluation_rubric"]
    return Sample(
        id=sample_id,
        input=task,
        target=json.dumps(rubrics),
        metadata={
            "structured_question": {"task": task},
            "results_dir": results_dir,
            "dataset_name": dataset_name,
        },
    )


# ======================================================================


@task
def e2e_discovery_validation():
    """Task for the E2E Discovery validation set."""
    return _create_e2e_discovery_task("validation", DEV_DATASET_JSON_FILE, "may-2025")


@task
def e2e_discovery_test():
    """Task for the E2E Discovery test set."""
    return _create_e2e_discovery_task("test", TEST_DATASET_JSON_FILE, "may-2025")


@task
def e2e_discovery_hard_validation():
    """Task for the E2E Discovery hard (harpa) validation set."""
    return _create_e2e_discovery_task(
        "validation", DEV_HARD_DATASET_JSON_FILE, "jun-2025-harpa"
    )


@task
def e2e_discovery_hard_test():
    """Task for the E2E Discovery hard (harpa) test set."""
    return _create_e2e_discovery_task(
        "test", TEST_HARD_DATASET_JSON_FILE, "jun-2025-harpa"
    )


# ======================================================================
# 		SCORER
# ======================================================================


@scorer(metrics=[accuracy(), stderr()])
def score_rubric() -> Scorer:
    """Scorer to check if the solution is correct."""

    async def score(state: TaskState, target: Target) -> Score:
        sample_id = state.sample_id
        task = state.metadata["structured_question"]
        results_dir = state.metadata["results_dir"]
        try:
            answer = json.loads(state.output.completion)
        except:
            print_and_log("Answer not in JSON format - ignoring")
            answer = {}
        #       answer = extract_json_from_response(state.output.completion)	# MikeD's suggestion - but I have problems importing this!

        # cache answer details for user's perusal if wanted
        os.makedirs(
            os.path.join(results_dir, "evaluations_json"), exist_ok=True
        )  # check dir exists
        os.makedirs(
            os.path.join(results_dir, "evaluations_html"), exist_ok=True
        )  # check dir exists
        output_file = os.path.join(
            results_dir, "evaluations_json", str(sample_id) + ".extracted.json"
        )
        html_file = os.path.join(
            results_dir, "evaluations_html", str(sample_id) + ".extracted.html"
        )

        if answer is None:
            score, explanation = (
                0,
                "Solver was unsuccessful for task - Assigning a score of 0",
            )
            with open(output_file, "w", encoding="utf-8") as f:
                json.dump(NULL_SCORE, f, indent=4)
            with open(html_file, "w", encoding="utf-8") as f:
                f.write("(No answer returned from solver)")
        else:
            try:
                rubrics = json.loads(target.text)
                score, explanation = await check_rubric(
                    sample_id, task, answer, rubrics, results_dir
                )
            except Exception as e:
                score, explanation = (
                    0,
                    f"Rubric evaluator failed! Error {e}. Assigning a score of 0",
                )
                with open(output_file, "w", encoding="utf-8") as f:
                    json.dump(NULL_SCORE, f, indent=4)
                    with open(html_file, "w", encoding="utf-8") as f:
                        f.write("(No answer returned from solver)")
        print_and_log(f"Overall score: {score:.2f} ({explanation})")

        return Score(
            value=score,
            answer=state.output.completion,
            #            explanation="Rubric scoring completed successfully.",
            explanation=explanation,
        )

    return score


### ======================================================================
###		DOWNLOAD A HF FILE (for the cached solvers)
### ======================================================================


# data = get_hf_file_contents("e2e_discovery/codescientist/may1-2025/evaluations_json/idea-30-simplified.extracted.json", "allenai/asta-bench-solver-data")
# get_hf_file_contents("e2e_discovery/autonora/may5-2025/reports/idea-100-simplified-results.json", "allenai/asta-bench-solver-data")
# get_hf_file_contents("e2e_discovery/faker/may5-2025/evaluations_json/idea-100-simplified.extracted.json", "allenai/asta-bench-solver-data")
def get_hf_file_contents(filename, repo_id=HF_CACHED_RESULTS_DIR):
    """
    Downloads a file from a Hugging Face dataset repo and returns its contents as a string.

    Args:
        filename (str): The name of the file to download (e.g. "idea-5-simplified.json").
        repo_id (str): The Hugging Face repo ID (e.g. "allenai/asta-bench-solver-data").

    Returns:
        str: The contents of the file.
    """
    try:
        local_path = hf_hub_download(
            repo_id=repo_id, filename=filename, repo_type="dataset"
        )
        with open(local_path, "r", encoding="utf-8") as f:
            return json.load(f)
    except HfHubHTTPError:
        print_and_log(f"(Warning: No HuggingFace results file {filename} found)")
        return None
    except json.JSONDecodeError as e:
        print_and_log(f"Invalid JSON format in {filename}: {e}")
        return None


### ======================================================================
###		RECORD TOKEN COUNTS (after a solver does a task)
### ======================================================================


def record_token_counts(results):
    # Record (cached) token usage (backdoor method)
    if results and "results" in results:
        all_token_counts = results["results"].get(
            "token_counts", []
        )  # e.g.  ["model":"gpt-4.1","prompt_tokens":100,"completion_tokens":310,"total_tokens":410]
        for token_counts in all_token_counts:
            model = token_counts["model"]
            prompt_tokens = token_counts.get("prompt_tokens", 0)
            completion_tokens = token_counts.get("completion_tokens", 0)
            total_tokens = token_counts.get("total_tokens", 0)
            usage = ModelUsage(
                input_tokens=prompt_tokens,
                output_tokens=completion_tokens,
                total_tokens=total_tokens,
            )
            record_model_usage_with_inspect(model, usage)
    else:
        print_and_log("(No token counts reported)")


### ======================================================================
###		CHECK THE RUBRIC
### ======================================================================


# Iterate through rubric in rubrics to score different categories about the actual answer
# RETURNS: a score (0-1)
# PJ: question is now a string, not a structure
# async def check_rubric(sample_id, task, answer, rubrics, model=None):	# model is unused
async def check_rubric(
    sample_id: str,
    task: dict[str, Any],
    answer: dict[str, Any],
    rubrics: list[dict[str, Any]],
    results_dir: str,
) -> Tuple[float, str]:
    # select model to use for applying the rubrics (LLM-as-judge)
    model_str = RUBRIC_SCORING_MODEL
    temperature = 0.0
    max_tokens = 8000

    print(
        "DEBUG: scoring task ", sample_id, " using ", model_str, "...", sep="", end=""
    )  # progress tracking for me

    # get the answer elements
    if answer and "results" in answer:
        results = answer["results"]
    else:
        overall_score = 0
        explanation = (
            f"Warning! No results found for {sample_id}. Assigning a score of 0."
        )
        print_and_log(explanation)
        return overall_score, explanation

    paper = results.get("report", "")
    code = results.get("code", [])
    artifacts = results.get("artifacts", [])
    #   trace = results.get("trace","")		# not used for eval

    rubric_evaluation = RubricEvaluation()
    #   rubrics0 = [rubrics[0]]
    evaluated_rubric_output = await rubric_evaluation.evaluate_rubric(
        rubrics,
        paper,
        code,
        artifacts,
        model_str=model_str,
        temperature=temperature,
        max_tokens=max_tokens,
    )

    # Save the output
    output_file = os.path.join(
        results_dir, "evaluations_json", str(sample_id) + ".extracted.json"
    )
    with open(output_file, "w", encoding="utf-8") as f:
        print_and_log(f"\nWrote score analysis (JSON) to {output_file}")
        json.dump(evaluated_rubric_output, f, indent=4)

    # Also save the HTML output
    html_file = os.path.join(
        results_dir, "evaluations_html", str(sample_id) + ".extracted.html"
    )
    try:
        convert_rubric_evaluation_to_html(
            evaluated_rubric_output, sample_id, paper, html_file
        )
    except Exception as e:
        print_and_log(
            f"Error converting rubric evaluation to HTML: {str(e)}", log_level="warning"
        )

    try:
        overall_score = evaluated_rubric_output["overall_score"][
            "average_required_overall"
        ]
        evaluated_rubric = evaluated_rubric_output["evaluated_rubric"]
        explanation = f"{overall_score:.2f} = average of rubric scores: "
        for rubric_element in evaluated_rubric:
            criteria_name = rubric_element["criteria_name"]
            required_or_optional = rubric_element["required_or_optional"]
            if required_or_optional == "required":
                response = rubric_element["response"]
                criteria_met = response["overall_evaluation"]["criteria_met"]
                explanation += f"{criteria_name}: {criteria_met}; "
    except Exception as e:
        print_and_log(
            f"Scoring error {e} Result of evaluate_rubric has an unexpected structure! Assigning a score of 0...",
            log_level="warning",
        )
        explanation = "Failed evaluation: evaluate_rubric has unexpected structure. Assigning a score of 0"
        overall_score = 0

    print()  # newline, for progress-tracking
    return overall_score, explanation


"""
evaluated_rubric_output = 
[   {   "idx": 0,
        "required_or_optional": "required",
        "response": {...,"overall_evaluation": {"criteria_met": 1, "overall_evaluation_explanation": ...}}   },
    {   "idx": 1, ... }, ... ]
"""
# ======================================================================
# 	APPLY RUBRICS (extracted from RubricEvaluation.py)
# ======================================================================

MAX_ARTIFACT_CHARACTERS = 10000  # size limit for artifacts for the evaluator


class RubricEvaluation:
    # Constructor
    def __init__(self):
        pass

    async def evaluate_rubric(
        self,
        rubric: list,
        paper: str,
        code: list,
        artifacts: list,
        model_str: str,
        temperature: float = 0.0,
        max_tokens: int = 8000,
    ) -> dict[str, Any]:
        # Rubric: A list of dictionaries.  Each must (minimally) contain the following keys: `criteria_name`, `criteria_met_question`, and `required_or_optional`.
        # Paper: A string containing the raw paper text to be evaluated.
        # Code: A list of dictionaries.  Each must contain the following keys: `filename` (str), `code` (str).
        # Artifacts: A list of dictionaries. Each must contain the following keys: `filename` (str), `format` (str), `artifact` (str).  `Format` can be: `txt`, `json`, `csv`, `tsv`, `md`, etc.

        # Returns: A dictionary containing the following keys:
        # `evaluated_rubric` (list of dictionaries).  Each dictionary will have one key added: `scores`, that contains sub-keys.
        # `overall_score`: TODO

        #       def mkPrompt(full_rubric:list, rubric_question:dict, paper:str, code:list, artifacts:list, reflection:str=None):   # PEC
        def mkPrompt(
            full_rubric: list,
            rubric_question: str,
            paper: str,
            code: list,
            artifacts: list,
            reflection: str | None = None,
        ):
            prompt = "You are ScientistGPT, the most advanced automated scientific discovery and inference system in the world. You can use your enormous intellect to solve any scientific problem, and always do so by breaking the problem down into its finest details.  You are accurate, faithful, and act with the highest scientific integrity.\n"
            prompt += "\n"
            prompt += "# Task\n"
            prompt += "Your task is to evaluate a scientific result (consisting of a paper, code, and potentially other artifacts like log files) against a single specific question from a rubric.\n"
            prompt += "The goal is to determine if the scientific result meets the criteria requested of the rubric question, and does so correctly, accurately, and faithfully, with scientific integrity.\n"
            prompt += "IMPORTANT: You should assume that the papers, code, and artifacts are written by one of a number of different artificial intelligence systems. Some of those systems are malfunctioning, and producing incorrect, unfaithful, or completely fabricated results.  Your job is to find them.  You will not hurt anyone's feelings by being critical.  You will be making the world a better place by identifying bad science, before it can have a negative impact.\n"
            prompt += "The information required to evaluate the rubric question must be separately evaluted from the three provided sources: the paper, the code, and any additional artifacts that were supplied.\n"
            prompt += "For example, information to answer some rubric questions might be solely available in one source (e.g. code, or paper, or artifacts), while other questions may have information that is provided in multiple sources (e.g. code, paper, and artifacts).\n"
            prompt += "You must evaluate the content of each source separately (as it relates to the rubric question), and also note whether the information in each source is consistent with the other sources, or if there are inconsistencies (which might indicate errors, omissions, or other issues with the science experiment that need to be addressed).\n"
            prompt += "\n"
            prompt += "To perform this task, you will be supplied with the following (below):\n"
            prompt += "- The rubric question\n"
            prompt += "- The paper\n"
            prompt += "- The code\n"
            prompt += "- Any additional artifacts that are available, like log files, results files, etc.\n"
            prompt += "\n"

            prompt += "# Rubric Question\n"
            prompt += "You will be provided with (1) the full rubric (for context), and (2) the specific rubric question that you must evaluate.\n"
            prompt += "\n"
            prompt += "## Full Rubric\n"
            prompt += "The full rubric is:\n"
            prompt += "```\n"
            prompt += json.dumps(full_rubric, indent=4) + "\n"
            prompt += "```\n"
            prompt += "\n"
            prompt += "## Specific Rubric Question\n"
            prompt += "The specific rubric question you must evaluate is:\n"
            prompt += "```\n"
            prompt += rubric_question + "\n"
            prompt += "```\n"
            prompt += "\n"

            prompt += "# Paper\n"
            prompt += "The raw text of the scientific paper is:\n"
            prompt += "```\n"
            prompt += paper + "\n"
            prompt += "```\n"
            prompt += "\n"

            prompt += "# Code\n"
            if code is None:
                code = []
            num_files_code = len(code)
            if num_files_code == 0:
                prompt += "No code was provided.\n"
                prompt += "\n"

            else:
                prompt += (
                    "The code is provided in the following "
                    + str(num_files_code)
                    + " files:\n"
                )
                prompt += "\n"
                for file_idx, file in enumerate(code):
                    filename = file.get("filename", "<missing filename>")
                    _code = file.get("code", "<missing code>")
                    prompt += "File " + str(file_idx) + " (" + str(filename) + "):\n"
                    prompt += "```\n"
                    prompt += _code + "\n"
                    prompt += "```\n"
                    prompt += "\n"

            prompt += "# Other Artifacts\n"
            if artifacts is None:
                artifacts = []
            num_files_artifacts = len(artifacts)
            if num_files_artifacts == 0:
                prompt += "No additional artifacts were provided.\n"
                prompt += "\n"
            else:
                prompt += (
                    "The following "
                    + str(num_files_artifacts)
                    + " additional artifacts were provided:\n"
                )
                prompt += "\n"
                for file_idx, file in enumerate(artifacts):
                    filename = file.get("filename", "<missing filename>")
                    artifact_format = file.get("format", "not provided")
                    artifact = str(file.get("artifact", "<missing artifact>"))
                    if (
                        len(artifact) > MAX_ARTIFACT_CHARACTERS
                    ):  # NEW - stop Claude blowing up
                        artifact = (
                            artifact[:MAX_ARTIFACT_CHARACTERS]
                            + "...[TRUNCATED FOR LENGTH]"
                        )

                    prompt += (
                        "File "
                        + str(file_idx)
                        + " ("
                        + str(filename)
                        + ", format: "
                        + str(format)
                        + "):\n"
                    )
                    prompt += "```\n"

                    # Add any special parsing for artifacts here
                    if artifact_format.strip().lower() == "json":
                        # If the artifact is JSON, format it nicely
                        try:
                            json_obj = json.loads(artifact)
                            prompt += json.dumps(json_obj, indent=4) + "\n"
                        except json.JSONDecodeError:
                            # If JSON decoding fails, just use the raw artifact
                            prompt += artifact + "\n"
                    else:
                        # If the artifact is not JSON, just use the raw artifact
                        prompt += artifact + "\n"
                    prompt += "```\n"
                    prompt += "\n"

            prompt += "# Evaluation Instructions\n"
            prompt += "You must evaluate whether the criteria used in the rubric question were met, and do so accurately, faithfully, and with scientific integrity.\n"
            prompt += "You must perfor this evaluation independently for each of (1) the paper, (2) the code, and (3) any additional artifacts.\n"
            prompt += "You must also note any inconsistencies between the sources, and whether the information in each source is consistent with the other sources.\n"
            prompt += "\n"

            prompt += "## Examples\n"
            prompt += "The following are illustrative cartoon examples:\n"
            prompt += "\n"

            prompt += "### Example 1\n"
            prompt += "Example 1 Rubric Question: `Was a ReAct baseline agent, consisting of a prompt provided to a language model (with separate 'think' and 'act' steps), evaluated on the XYZ benchmark?`\n"
            prompt += "Example 1 Paper:\n"
            prompt += " - Positive example: Paper text includes a table or text describing the ReAct agent being evaluated on the XYZ benchmark. e.g. `The ReAct agent was evaluated on the XYZ benchmark, and the results are shown in Table 1.`\n"
            prompt += " - Negative example: The paper does not mention the ReAct baseline being run, specifies it was run on a different benchmark than XYZ, does not report any results on the ReAct baseline, or has some other flaw that calls this result into question.\n"
            prompt += "Example 1 Code:\n"
            prompt += " - Positive example: The code includes a script that (1) prompts a language model with separate 'think' and 'act' steps, (2) stores or uses the result in some way, and (3) the input to the ReAct model appears to come from the XYZ baseline.\n"
            prompt += " - Negative example: The code does not include a script that prompts a language model with separate 'think' and 'act' steps, or the input to the ReAct model does not appear to come from the XYZ baseline.\n"
            prompt += " - Negative example: The code includes only 'think' or 'act' steps, but not both.\n"
            prompt += " - Negative example: The code includes additional steps that are not part of the ReAct agent (e.g. 'think', 'act', 'analyze', 'decide', etc.)\n"
            prompt += "Example 1 Artifacts:\n"
            prompt += " - Positive example: The artifacts include a log file or other output that shows the ReAct agent was run on the XYZ benchmark, particularly with the prompt, and responses.\n"
            prompt += " - Positive example: The artifacts include a results file showing the performance of the ReAct agent on the XYZ benchmark, either separately (on individual problems in the XYZ benchmark), or in aggregate/average across all problems in XYZ.\n"
            prompt += " - Negative example: The artifacts do not include any log files or other output that shows the ReAct agent was run on the XYZ benchmark.\n"
            prompt += "\n"

            prompt += "# Additional Scenarios and Guidelines\n"
            prompt += "You should consider yourself similar to a scientific peer reviewer.  Meeting the criteria in the rubric question is important, and requires clear evidence, especially to protect from bad actors or unfaithful science.  If full evidence to support a meeting a criteria is missing where it normally should be there, then it's clearly not fulfilling the criteria, and you need to call that out.  Your job isn't to be agreeable, it's to be accurate.  At the same time, if it could reasonably be expected that something shouldn't be reported in either paper, or that it wouldn't exist in the code, or that it wouldn't be a part of the other artifacts, it's OK to give it a pass if the normal way that information is created, vetted, and conveyed is satisfied.\n"
            prompt += "Let's be very clear here, and workshop some scenarios:\n"
            prompt += "Scenario 1: Requires that a specific LLM model be run.\n"
            prompt += " - This would normally be reported in the paper, exist in the code (because this requires code), and may or may not include some artifacts (e.g. the overall results, or input/output for each piece of data provided to the model).  You can't run a model without code, so code must be there.  The code could analyze the artifacts internally, or save them to show it's work, so these may or may not exist.  Reporting them in the paper is conventional.\n"
            prompt += "\n"
            prompt += "Scenario 2: Statistical analysis between two models.\n"
            prompt += " - This would normally be reported in the paper.  If the analysis was performed in the code, it should be clearly visible in the code as well.  If it was in some external package (like R), either the code, or the results files from the analysis, or both should be included.\n"
            prompt += "\n"
            prompt += "Scenario 3: Data was created with certain properties.\n"
            prompt += " - This would normally be reported in the paper.  If the data was created with code, this should be included.  If the analysis to verify the properties was performed with code, that code should be included.  If, instead of a code-based verification, the verification was performed manually (e.g. by a human, or some external package), those artifacts should be included to verify that it has the required properties.\n"
            prompt += "\n"
            prompt += "## Double Checking Code and Artifacts\n"
            prompt += "The code might SAY it's doing something in the comments, but is it really?  If you perform an examination of the code, does it reasonably appear to be able to perform the function that's described in the criteria, or is it likely incorrect, incomplete, simplified, or a fabrication?"
            prompt += "The same applies for data/arficacts -- double check that they are what they say they are, and that they make sense, and showcase the properties that they're required to show (in a faithful, accurate way).\n"
            prompt += "\n"
            if reflection is not None:
                prompt += "## Overall Evaluation\n"
                prompt += "Base your overall evaluation for each rubric question on your evaluation of the individual elements.  If there are important things missing that are required to verify the criteria, or there are other issues, you *should not take this on trust*, you must fail the overall criterion in these cases!  Evaluating something as a pass that does not meet all the standards or that has issues is considered a CRITICAL error!\n"
                prompt += "\n"

            # TODO: Add additional example?

            if reflection is not None:
                prompt += "# Reflection\n"
                prompt += "This is a reflection step.  Previously, you were asked to evaluate the individual elements of the rubric question. Now, you must reflect on your evaluation of the individual elements, and provide an overall evaluation of whether the standard is met based on the evaluations.\n"
                prompt += "Your evaluation was:\n"
                prompt += "```\n"
                prompt += json.dumps(reflection, indent=4) + "\n"
                prompt += "```\n"
                prompt += "\n"

            prompt += "# Output format\n"
            prompt += "The output format is JSON, between a single codeblock (```). You can think and plan as much as you'd like (in plain text) before your official JSON response, but there must be only one codeblock (```), and the only thing in that codeblock must be valid JSON. The JSON will be automatically parsed, so if the content isn't valid JSON, this will be considered a critical error.\n"
            prompt += "Below is the output format, including required keys:\n"
            prompt += "```\n"
            prompt += "{\n"
            prompt += '    "rubric_question": "<The specific rubric question that was evaluated>",\n'
            prompt += '    "standard": "<What is the minimum requirement for evidence? Some combination of Paper, and/or Code, and/or Artifacts? Be specific on what solid standards should be to convincingly show the work was done, in a verifiable, auditable way.>",\n'
            prompt += '    "paper_evaluation": {\n'
            prompt += '        "paper_evaluation_explanation": "<A short evaluation of how the paper does or does not speak to meeting this specific rubric question>",\n'
            prompt += '        "criteria_met": 1, 0, or -1  # 1 = criteria met, 0 = criteria not met, -1 = the paper did not contain any information about this to evaluate this question.\n'
            prompt += "    },\n"
            prompt += '    "code_evaluation": {\n'
            prompt += '        "code_evaluation_explanation": "<A short evaluation of how the code does or does not speak to meeting this specific rubric question>",\n'
            prompt += '        "criteria_met": 1, 0, or -1 # 1 = criteria met, 0 = criteria not met, -1 = the code did not contain any information about this to evaluate this question.\n'
            prompt += "    },\n"
            prompt += '    "artifacts_evaluation": {\n'
            prompt += '        "artifacts_evaluation_explanation": "<A short evaluation of how the artifacts do or do not speak to meeting this specific rubric question>",\n'
            prompt += '        "criteria_met": 1, 0, or -1 # 1 = criteria met, 0 = criteria not met, -1 = the artifacts did not contain any information about this to evaluate this question.\n'
            prompt += "    },\n"
            prompt += '    "consistency_evaluation": {\n'
            prompt += '        "consistency_evaluation_explanation": "<Did the information from the paper, code, and/or artifacts, when available, support each other, or contain contradictions? Be specific in details>",\n'
            prompt += '        "consistency_score": 1 or 0 # 1 = consistent, 0 = some consistency issues\n'
            prompt += "    },\n"
            if reflection is not None:
                prompt += '    "overall_evaluation": {\n'
                prompt += '        "overall_evaluation_explanation": "<A short but CRITICAL evaluation of whether the standard is met based on the evaluations above. Is everything normally required to verify this criterion available, accurate, faithful, and correct, without issues?  If not, this should not pass.>",\n'
                prompt += '        "criteria_met": 1, 0, or -1 # 1 = criteria met, 0 = criteria not met, -1 = the paper did not contain any information about this to evaluate this question.\n'
                prompt += "    }\n"
            prompt += "}\n"
            prompt += "```\n"
            prompt += "\n"
            prompt += "## Notes\n"
            prompt += "The following are important notes:\n"
            prompt += " - The JSON output must be valid JSON, and the only thing in the codeblock (```) must be valid JSON.  If it is not valid JSON, this will be considered a critical error.\n"
            prompt += " - The JSON output must contain all the keys specified in the example above.\n"
            prompt += " - Your answers must be correct, accurate, faithful to the rubric question and the sources, and with the highest scientific integrity.\n"
            prompt += (
                " - You must not hallucinations, or make errors in your evaluation.\n"
            )
            prompt += " - You can think and plan as much as you'd like (in plain text) before your official JSON response, but there must be only one codeblock (```) in the output, and the only thing in that codeblock must be valid JSON.\n"

            return prompt

        start_time = time.time()
        errors = []

        async def evaluate_rubric_question(
            idx,
            rubric,
            rubric_question,
            paper,
            code,
            artifacts,
            model_str,
            temperature,
            max_tokens,
        ):
            rubric_elem = rubric[idx]

            # Make the first prompt
            prompt = mkPrompt(rubric, rubric_question, paper, code, artifacts)
            responseJSON, responseText, _ = await getLLMResponseJSON(
                promptStr=prompt,
                model=model_str,
                maxTokens=max_tokens,
                temperature=temperature,
                jsonOut=True,
            )

            # Make the reflection prompt
            prompt1 = mkPrompt(
                rubric, rubric_question, paper, code, artifacts, reflection=responseJSON
            )
            responseJSON1, responseText1, _ = await getLLMResponseJSON(
                promptStr=prompt1,
                model=model_str,
                maxTokens=max_tokens,
                temperature=temperature,
                jsonOut=True,
            )

            # Pack the response
            packed = {
                "idx": idx,
                "response": responseJSON1,
                "responseText": responseText1,
                "rubric_elem": rubric_elem,
                "cost": 0,
            }

            return packed

        # Evaluate each rubric question
        evaluated_rubric_output = []

        # Create all tasks
        tasks = [
            asyncio.create_task(
                evaluate_rubric_question(
                    idx,
                    rubric,
                    rubric_elem["criteria_met_question"],
                    paper,
                    code,
                    artifacts,
                    model_str=model_str,
                    temperature=temperature,
                    max_tokens=max_tokens,
                )
            )
            for idx, rubric_elem in enumerate(rubric)
        ]

        # Process tasks as they complete
        for coro in asyncio.as_completed(tasks):
            try:
                packed = await coro
                if packed is None:
                    print_and_log(
                        "Error evaluating rubric question.", log_level="warning"
                    )
                    continue
                evaluated_rubric_output.append(packed)
            except Exception:
                logger.exception("Error evaluating rubric question:")
                raise

        # Create the final output
        evaluated_rubric = []
        # Sort the output by the index of the rubric question
        evaluated_rubric_output.sort(key=lambda x: x["idx"])

        for packed in evaluated_rubric_output:
            idx = packed["idx"]
            response = packed["response"]
            rubric_elem = packed["rubric_elem"]
            criteria_name = rubric_elem["criteria_name"]
            required_or_optional = rubric_elem["required_or_optional"]

            pack_step = {
                "idx": idx,
                "rubric_elem": rubric_elem,
                "criteria_name": criteria_name,
                "required_or_optional": required_or_optional,
                "response": response,
            }

            # Add the response to the output
            evaluated_rubric.append(pack_step)

        # Perform overall analysis to get summary scores
        num_required = 0
        num_optional = 0
        count_required_all3 = 0
        count_optional_all3 = 0
        count_required_atleast1 = 0
        count_optional_atleast1 = 0

        count_required_overall = 0
        count_optional_overall = 0

        consistency_issues_list = []
        # Iterate through the evaluated rubric

        for idx, rubric_eval_record in enumerate(evaluated_rubric):
            try:
                # Get the rubric element
                eval_response = rubric_eval_record["response"]

                criteria_name = rubric_eval_record["criteria_name"]
                required_or_optional = (
                    rubric_eval_record["required_or_optional"].strip().lower()
                )

                consistency_issues = 0

                scores = []
                scores.append(eval_response["paper_evaluation"]["criteria_met"])
                scores.append(eval_response["code_evaluation"]["criteria_met"])
                scores.append(eval_response["artifacts_evaluation"]["criteria_met"])
                consistency_issues = eval_response["consistency_evaluation"][
                    "consistency_score"
                ]
                if consistency_issues < 1:
                    consistency_str = (
                        "Rubric question "
                        + str(idx)
                        + " has consistency issues between the sources.\n"
                    )
                    consistency_str += (
                        "Question text: " + eval_response["rubric_question"] + "\n"
                    )
                    consistency_str += (
                        "Consistency evaluation: "
                        + eval_response["consistency_evaluation"][
                            "consistency_evaluation_explanation"
                        ]
                        + "\n"
                    )
                    consistency_issues_list.append(consistency_str)

                # Get the 'overall' evaluation
                overall_evaluation = eval_response["overall_evaluation"]["criteria_met"]

                # Check if there's at least one 1 in the scores
                score_sum = 0
                for score in scores:
                    if score >= 1:
                        score_sum += 1

                if required_or_optional == "required":
                    num_required += 1
                    if score_sum >= 1:
                        count_required_atleast1 += 1
                    if score_sum == 3:
                        count_required_all3 += 1
                    if overall_evaluation >= 1:
                        count_required_overall += 1
                elif required_or_optional == "optional":
                    num_optional += 1
                    if score_sum >= 1:
                        count_optional_atleast1 += 1
                    if score_sum == 3:
                        count_optional_all3 += 1
                    if overall_evaluation >= 1:
                        count_optional_overall += 1
                else:
                    errors.append(
                        "Rubric element "
                        + str(idx)
                        + " does not contain a valid required_or_optional value.  Must be either 'required' or 'optional'. Found: `"
                        + str(required_or_optional)
                        + "`."
                    )

            except Exception as e:
                errors.append(
                    "Error consolodating rubric evaluation for question "
                    + str(idx)
                    + ": "
                    + str(e)
                )
                continue

        # Create average scores
        average_required_all3 = 0.0
        average_optional_all3 = 0.0
        average_required_atleast1 = 0.0
        average_optional_atleast1 = 0.0
        average_required_overall = 0.0
        average_optional_overall = 0.0
        if num_required > 0:
            average_required_all3 = count_required_all3 / num_required
            average_required_atleast1 = count_required_atleast1 / num_required
            average_required_overall = count_required_overall / num_required
        if num_optional > 0:
            average_optional_all3 = count_optional_all3 / num_optional
            average_optional_atleast1 = count_optional_atleast1 / num_optional
            average_optional_overall = count_optional_overall / num_optional

        total_runtime_seconds = time.time() - start_time

        # Create the final output
        evaluated_rubric_output = {
            "evaluated_rubric": evaluated_rubric,
            "overall_score": {
                "average_required_all3": average_required_all3,
                "average_optional_all3": average_optional_all3,
                "average_required_atleast1": average_required_atleast1,
                "average_optional_atleast1": average_optional_atleast1,
                "average_required_overall": average_required_overall,
                "average_optional_overall": average_optional_overall,
            },
            "consistency_issues_list": consistency_issues_list,
            "errors": errors,
            "total_runtime_seconds": round(total_runtime_seconds, 2),
            "total_cost": 0,  # placeholder
        }

        return evaluated_rubric_output


# ======================================================================
# 	POSTPROCESSING OF THE RESULT
# ======================================================================


# Convert the evaluated rubric to HTML, for easy viewing
def convert_rubric_evaluation_to_html(
    evaluated_rubric_struct: dict[str, Any],
    name: str,
    paper: str,
    filenameOut: str,
    code: list | None = None,
    artifacts: list | None = None,
) -> None:
    """Convert the evaluated rubric to HTML format and write to `filenameOut`."""
    TEXT_WRAP_WIDTH = 1000

    htmlOut = "<html>\n"
    htmlOut += "<head><title>Rubric Evaluation (" + str(name) + ")</title></head>\n"
    htmlOut += "<body>\n"
    htmlOut += "<h1>Rubric Evaluation (" + str(name) + ")</h1>\n"

    # Show the overall average score
    htmlOut += "<h2>Overall Average Score</h2>\n"
    # average_required_overall
    htmlOut += (
        "<p>Average Required Overall: <br>"
        + str(
            round(
                evaluated_rubric_struct["overall_score"]["average_required_overall"], 2
            )
        )
        + "</b></p>\n"
    )
    # average_optional_overall
    htmlOut += (
        "<p>Average Optional Overall: "
        + str(
            round(
                evaluated_rubric_struct["overall_score"]["average_optional_overall"], 2
            )
        )
        + "</p>\n"
    )

    htmlOut += "<h2>Rubric Elements</h2>\n"

    # First, make a table showing the valuation according to the rubric.
    # 4 columns (paper, code, artifacts, overall).  First line in each is the score, next lines are the explanation text.  Max 50 characters per line (use textwrap to split).  Centered within each cell.
    # Cells should be green if 1, red if 0, yellow if -1.
    htmlOut += "<table border='1'>\n"
    htmlOut += "<tr>\n"
    htmlOut += "<th>Criteria</th>\n"
    htmlOut += "<th>Standard</th>\n"
    htmlOut += "<th>Paper</th>\n"
    htmlOut += "<th>Code</th>\n"
    htmlOut += "<th>Artifacts</th>\n"
    htmlOut += "<th>Consistency</th>\n"
    htmlOut += "<th>Overall</th>\n"
    htmlOut += "</tr>\n"

    for idx, rubric_eval_record in enumerate(
        evaluated_rubric_struct["evaluated_rubric"]
    ):
        #        try:
        # Get the rubric element
        eval_response = rubric_eval_record["response"]

        criteria_name = rubric_eval_record["criteria_name"]
        required_or_optional = (
            rubric_eval_record["required_or_optional"].strip().lower()
        )
        criteria_met_question = rubric_eval_record["rubric_elem"][
            "criteria_met_question"
        ]

        standard = eval_response["standard"]

        # Get the scores
        paper_score = eval_response["paper_evaluation"]["criteria_met"]
        code_score = eval_response["code_evaluation"]["criteria_met"]
        artifacts_score = eval_response["artifacts_evaluation"]["criteria_met"]
        consistency_score = eval_response["consistency_evaluation"]["consistency_score"]
        overall_score = eval_response["overall_evaluation"]["criteria_met"]

        # Get the explanations
        paper_explanation = eval_response["paper_evaluation"][
            "paper_evaluation_explanation"
        ]
        code_explanation = eval_response["code_evaluation"][
            "code_evaluation_explanation"
        ]
        artifacts_explanation = eval_response["artifacts_evaluation"][
            "artifacts_evaluation_explanation"
        ]
        consistency_explanation = eval_response["consistency_evaluation"][
            "consistency_evaluation_explanation"
        ]
        overall_explanation = eval_response["overall_evaluation"][
            "overall_evaluation_explanation"
        ]

        # Create a new row for the table
        htmlOut += "<tr>\n"
        # Criteria
        htmlOut += (
            "<td><b>"
            + str(criteria_name)
            + "</b><br>"
            + str(required_or_optional)
            + "<br>"
            + "<br>".join(
                textwrap.fill(criteria_met_question, TEXT_WRAP_WIDTH).split("\n")
            )
            + "</td>\n"
        )

        # Standard
        htmlOut += (
            "<td>"
            + "<br>".join(textwrap.fill(standard, TEXT_WRAP_WIDTH).split("\n"))
            + "</td>\n"
        )

        # Add the paper score and explanation
        paper_color = "yellow"
        if paper_score == 1:
            paper_color = "green"
        elif paper_score == 0:
            paper_color = "red"

        # htmlOut += "<td style='background-color: " + str(paper_color) + "; text-align: center;'>" + str(paper_score) + "</td>\n"
        # htmlOut += "<td style='text-align: left;'>" + textwrap.fill(paper_explanation, width=50) + "</td>\n"
        htmlOut += (
            "<td style='background-color: "
            + str(paper_color)
            + "; text-align: center;'>"
            + str(paper_score)
            + "<br>"
            + "<br>".join(
                textwrap.fill(paper_explanation, width=TEXT_WRAP_WIDTH).split("\n")
            )
            + "</td>\n"
        )

        # Code score
        code_color = "yellow"
        if code_score == 1:
            code_color = "green"
        elif code_score == 0:
            code_color = "red"
        # htmlOut += "<td style='background-color: " + str(code_color) + "; text-align: center;'>" + str(code_score) + "<br>" + textwrap.fill(code_explanation, width=TEXT_WRAP_WIDTH) + "</td>\n"
        htmlOut += (
            "<td style='background-color: "
            + str(code_color)
            + "; text-align: center;'>"
            + str(code_score)
            + "<br>"
            + "<br>".join(
                textwrap.fill(code_explanation, width=TEXT_WRAP_WIDTH).split("\n")
            )
            + "</td>\n"
        )

        # Artifacts score
        artifacts_color = "yellow"
        if artifacts_score == 1:
            artifacts_color = "green"
        elif artifacts_score == 0:
            artifacts_color = "red"
        # htmlOut += "<td style='background-color: " + str(artifacts_color) + "; text-align: center;'>" + str(artifacts_score) + "</td>\n"
        # htmlOut += "<td style='text-align: left;'>" + textwrap.fill(artifacts_explanation, width=50) + "</td>\n"
        htmlOut += (
            "<td style='background-color: "
            + str(artifacts_color)
            + "; text-align: center;'>"
            + str(artifacts_score)
            + "<br>"
            + "<br>".join(
                textwrap.fill(artifacts_explanation, width=TEXT_WRAP_WIDTH).split("\n")
            )
            + "</td>\n"
        )

        # Consistency
        consistency_color = "yellow"
        if consistency_score == 1:
            consistency_color = "green"
        elif consistency_score == 0:
            consistency_color = "red"
        # htmlOut += "<td style='background-color: " + str(consistency_color) + "; text-align: center;'>" + str(consistency_score) + "</td>\n"
        # htmlOut += "<td style='text-align: left;'>" + textwrap.fill(consistency_explanation, width=50) + "</td>\n"
        htmlOut += (
            "<td style='background-color: "
            + str(consistency_color)
            + "; text-align: center;'>"
            + str(consistency_score)
            + "<br>"
            + "<br>".join(
                textwrap.fill(consistency_explanation, width=TEXT_WRAP_WIDTH).split(
                    "\n"
                )
            )
            + "</td>\n"
        )

        # Overall score
        overall_color = "yellow"
        if overall_score == 1:
            overall_color = "green"
        elif overall_score == 0:
            overall_color = "red"
        # htmlOut += "<td style='background-color: " + str(overall_color) + "; text-align: center;'>" + str(overall_score) + "</td>\n"
        # htmlOut += "<td style='text-align: left;'>" + textwrap.fill(overall_explanation, width=50) + "</td>\n"
        htmlOut += (
            "<td style='background-color: "
            + str(overall_color)
            + "; text-align: center;'>"
            + str(overall_score)
            + "<br>"
            + "<br>".join(
                textwrap.fill(overall_explanation, width=TEXT_WRAP_WIDTH).split("\n")
            )
            + "</td>\n"
        )

        htmlOut += "</tr>\n"

        # except Exception as e:
        #     print("Error consolodating rubric evaluation for question " + str(idx) + ": " + str(e))
        #     continue
    htmlOut += "</table>\n"

    # Show consistency issues
    htmlOut += "<h2>Consistency Issues</h2>\n"
    if len(evaluated_rubric_struct["consistency_issues_list"]) > 0:
        htmlOut += "<ul>\n"
        for issue in evaluated_rubric_struct["consistency_issues_list"]:
            htmlOut += "<li>" + str(issue) + "</li>\n"
        htmlOut += "</ul>\n"
    else:
        htmlOut += "<p>No consistency issues found.</p>\n"

    # Show errors
    htmlOut += "<h2>Errors</h2>\n"
    if len(evaluated_rubric_struct["errors"]) > 0:
        htmlOut += "<ul>\n"
        for error in evaluated_rubric_struct["errors"]:
            htmlOut += "<li>" + str(error) + "</li>\n"
        htmlOut += "</ul>\n"
    else:
        htmlOut += "<p>No errors found.</p>\n"

    # Show the paper (as preformatted text)
    htmlOut += "<h2>Paper</h2>\n"
    htmlOut += "<pre>" + str(paper) + "</pre>\n"

    if code is not None:
        # Show the code (as preformatted text)
        htmlOut += "<h2>Code</h2>\n"
        for code_elem in code:
            htmlOut += "<h3>Code: " + str(code_elem["filename"]) + "</h3>\n"
            htmlOut += "<pre>" + str(code_elem["code"]) + "</pre>\n"

    if artifacts is not None:
        # Show the artifacts (as preformatted text)
        htmlOut += "<h2>Artifacts</h2>\n"
        for artifact_elem in artifacts:
            htmlOut += "<h3>Artifact: " + str(artifact_elem["filename"]) + "</h3>\n"
            htmlOut += "<pre>" + str(artifact_elem["artifact"]) + "</pre>\n"

    # Complete HTML
    htmlOut += "</body>\n"
    htmlOut += "</html>\n"

    # Save the HTML output to a file
    with open(filenameOut, "w", encoding="utf-8") as f:
        f.write(htmlOut)
    print_and_log(f"Wrote score analysis (HTML) to {str(filenameOut)}")


# ======================================================================
# 		BASIC LLM CALL - OLD
# ======================================================================


DEFAULT_MAX_TOKENS = 8000

#
#   Helper: Get a response from an LLM model
#


# Get a response from an LLM model using litellm
async def getLLMResponseJSON(
    promptStr: str,
    model: str,
    temperature: float = 0,
    maxTokens: int = DEFAULT_MAX_TOKENS,
    jsonOut: bool = True,
):
    MAX_RETRIES = 10
    MAX_GENERATION_TIME_SECONDS = 60 * 5

    try:
        responseJSON, responseText, cost = await _getLLMResponseJSON(
            promptStr,
            model,
            temperature,
            maxTokens,
            jsonOut,
            max_retries=MAX_RETRIES,
            timeout=MAX_GENERATION_TIME_SECONDS,
        )
        return responseJSON, responseText, cost
    except Exception as e:
        # logger.exception prints the traceback
        logger.exception("ERROR: Could not get LLM response.")

        # Check for some known kinds of errors
        errorStr = str(e)
        if "prompt is too long" in errorStr:
            print_and_log("ERROR: Prompt is too long. Giving up.", log_level="warning")
            return None, "", 0

    # If we reach here, something terrible happened
    error_message = "ERROR: Could not get LLM response. Exiting."
    print_and_log(error_message, log_level="warning")
    raise ValueError(error_message)


# ----------


async def _getLLMResponseJSON(
    promptStr: str,
    model: str,
    temperature: float = 0,
    maxTokens: int = DEFAULT_MAX_TOKENS,
    jsonOut: bool = True,
    max_retries: int | None = None,
    timeout: int | None = None,
):
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": promptStr},
            ],
        }
    ]

    if (model == "claude-3-5-sonnet-20240620") and (maxTokens > 4096):
        # TODO
        raise ValueError(
            f'Requested {maxTokens} tokens with {model} but we haven\'t set up {{"anthropic-beta": "max-tokens-3-5-sonnet-2024-07-15"}}'
        )

    response = None

    if "claude" in model:
        inspect_messages = []
        for msg in messages:
            cls = (
                ChatMessageAssistant
                if msg["role"] == "assistant"
                else ChatMessageSystem if msg["role"] == "system" else ChatMessageUser
            )
            inspect_messages.append(cls(content=msg["content"]))
        # Not using max_tokens or temperature; they weren't used in the
        # original when validating the scoring, so they are now vestigial
        # max_tokens=maxTokens, temperature=temperature,
        response = await get_model("anthropic/" + model).generate(
            inspect_messages,
            config=GenerateConfig(
                max_retries=max_retries,
                timeout=timeout,
            ),
        )

    else:
        raise ValueError(f"Don't know how to score with model {model}!")

    print(".", end="")  # progress monitoring in local window

    responseTextObj = response.choices[
        0
    ].message.content  # e.g., responseText = [ContentText(type='text', text="Hi I'm Claude)]
    responseText = responseTextObj[0].text

    responseOutJSON = None
    # Convert the response text to JSON
    try:
        responseOutJSON = json.loads(responseText)
    except Exception:
        pass
    #       print("ERROR: Could not convert response to JSON on first pass. ")

    # Second pass for converting to JSON
    if responseOutJSON is None:
        # Try to find the last JSON block in the response, which starts with "```"
        lines = responseText.split("\n")
        startIdx = 0
        endIdx = 0
        for idx, line in enumerate(lines):
            if line.startswith("```"):
                startIdx = endIdx
                endIdx = idx

        if (startIdx >= 0) and (endIdx > 0):
            # Exclude the start and end line
            linesJSON = lines[startIdx + 1 : endIdx]
            # Join the lines
            linesJSONStr = "\n".join(linesJSON)
            # Try to convert to JSON
            try:
                responseOutJSON = json.loads(linesJSONStr)
            except Exception:
                print_and_log(
                    "ERROR: Could not convert response to JSON on second pass. ",
                    log_level="warning",
                )
        else:
            print_and_log(
                "ERROR: Could not find JSON block in response. ", log_level="warning"
            )

    # Return
    dummy_cost = 0
    return responseOutJSON, responseText, dummy_cost


def print_and_log(message, log_level="info"):
    print(message)
    if log_level == "debug":
        logger.debug(message)
    elif log_level == "warning":
        logger.warning(message)
    elif log_level == "error":
        logger.error(message)
    elif log_level == "info":
        logger.info(message)
    else:
        logger.error(f"ERROR! bad log_level {log_level} passed to print_and_log!")
