from llama_cpp import Llama  # loads llama-cpp-python so we can run GGUF models
import os  # lets us work with file paths
import random  # lets us pick random items from a list
import requests  # lets us call web APIs over HTTP
import pandas as pd  # helps with working with tables and spreadsheets
import re  # lets us search text with patterns
import json  # lets us parse structured JSON output from the model
from IPython.display import display  # lets us show interactive widgets inside the notebook
import ipywidgets as widgets  # adds dropdown menus and buttons for simple UI

possible_directories = [
    "/home/jovyan/shared/",
    "/home/jovyan/shared_readwrite/",
    "/Users/ericvandusen/SmallLM/Models"
]

existing_directories = []
for directory_path in possible_directories:
    if os.path.exists(directory_path):
        print("Found possible directory:", directory_path)
        existing_directories.append(directory_path)
    else:
        print("Did not find:", directory_path)

if len(existing_directories) > 0:
    model_directory = existing_directories[0]
    print("Using this directory by default:", model_directory)
else:
    model_directory = input("Type a directory path that contains your .gguf files: ")

print("Current model directory:", model_directory)

available_models = []
for filename in os.listdir(model_directory):
    if filename.endswith(".gguf"):
        available_models.append(filename)

if len(available_models) == 0:
    print("No .gguf files found in", model_directory)
else:
    print("Models found in", model_directory)
    dropdown_default = available_models[0]
    for candidate_name in available_models:
        lowercase_name = candidate_name.lower()
        if "qwen" in lowercase_name:
            dropdown_default = candidate_name
            break
    model_dropdown = widgets.Dropdown(
        options=available_models,
        description="Model:",
        value=dropdown_default
    )
    display(model_dropdown)
    print("Use the dropdown to pick a model, then run the next cell to load it.")

selected_model_name = model_dropdown.value
model_path = os.path.join(model_directory, selected_model_name)

model = Llama(
    model_path=model_path,
    n_ctx=2048,
    n_threads=4,
    chat_format="chatml",
    verbose=False
)

# Build the question from parts
warmup_question = "If 3x + 5 = 14, what is the value of x?"
warmup_choice_map = {
    "A": "2",
    "B": "3",
    "C": "4",
    "D": "5"
}

warmup_choices_text = "\n".join(
    f"{label}: {value}" for label, value in warmup_choice_map.items()
)

warmup_prompt = f"""Here is an SAT-style multiple-choice question:

Question: {warmup_question}

Choices:
{warmup_choices_text}

Respond ONLY with valid JSON in this exact format: {{"choice": "B"}}
Choose one letter from A, B, C, or D."""

messages = [
    {
        "role": "system",
        "content": (
            "You are an SAT math solver. "
            "Always respond with valid JSON containing a single field 'choice' "
            "set to one of: A, B, C, or D. No explanation, no extra text."
        )
    },
    {"role": "user", "content": warmup_prompt}
]

warmup_response = model.create_chat_completion(
    messages=messages,
    response_format={
        "type": "json_object",
        "schema": {
            "type": "object",
            "properties": {
                "choice": {
                    "type": "string",
                    "enum": ["A", "B", "C", "D"]
                }
            },
            "required": ["choice"]
        }
    },
    temperature=0.0
)

warmup_raw_text = warmup_response["choices"][0]["message"]["content"].strip()
print("Raw model output:", warmup_raw_text)

# Parse with fallback
warmup_choice = ""
try:
    warmup_json = json.loads(warmup_raw_text)
    warmup_choice = str(warmup_json.get("choice", "")).strip().upper()
except json.JSONDecodeError as parse_error:
    print("JSON parse failed:", parse_error)
    # Fallback: scan for a lone A/B/C/D
    fallback_match = re.search(r'\b([A-D])\b', warmup_raw_text.upper())
    if fallback_match:
        warmup_choice = fallback_match.group(1)

if warmup_choice not in ["A", "B", "C", "D"]:
    print("WARNING: Could not extract a valid choice from model output.")
    warmup_choice = "UNKNOWN"

correct = "B"
print(f"Parsed choice:  {warmup_choice}")
print(f"Correct choice: {correct}")
print(f"Result: {'✓ Correct' if warmup_choice == correct else '✗ Wrong'}")

base_url = "https://pinesat.com/api/questions"

english_questions = requests.get(base_url, params={"section": "english"}).json()
math_questions = requests.get(base_url, params={"section": "math"}).json()

english_nested = pd.DataFrame(english_questions)
math_nested = pd.DataFrame(math_questions)

print("English questions:", len(english_nested))
print("Math questions:", len(math_nested))

english_questions[0]

english_questions[0]["question"]

english_questions[0]["question"]['question']

english_questions[0]["question"]['choices']

english_questions[0]["question"]['paragraph']

english_questions[0]["question"]['correct_answer']

# Pick a math question from the question bank that we can test on
math_questions[0]

# Pick a random question from the English question bank
random_index = random.randint(0, len(english_questions) - 1)
random_entry = english_questions[random_index]

# Pull out the question text, passage, and choices
random_question_text = random_entry["question"]["question"]
random_paragraph_text = random_entry["question"].get("paragraph", "")
random_choices_raw = random_entry["question"]["choices"]

choice_labels = ["A", "B", "C", "D", "E", "F"]
random_choice_map = {}

if isinstance(random_choices_raw, dict):
    for raw_label in random_choices_raw:
        clean_label = str(raw_label).strip().upper()
        random_choice_map[clean_label] = str(random_choices_raw[raw_label])
elif isinstance(random_choices_raw, list):
    for item_index in range(len(random_choices_raw)):
        if item_index < len(choice_labels):
            mapped_label = choice_labels[item_index]
            random_choice_map[mapped_label] = str(random_choices_raw[item_index])
else:
    random_choice_map["A"] = str(random_choices_raw)

valid_labels = []
for label in choice_labels:
    if label in random_choice_map:
        valid_labels.append(label)

random_option_lines = []
for label in valid_labels:
    random_option_lines.append(label + ": " + random_choice_map[label])
random_options_text = "\n".join(random_option_lines)

if len(random_paragraph_text) > 0:
    english_prompt_question_text = "Passage: " + random_paragraph_text + "\n\nQuestion: " + random_question_text
else:
    english_prompt_question_text = "Question: " + random_question_text

english_structured_prompt = """
{question_text}

Options:
{options_text}

Return your answer as JSON with this exact shape:
{{"choice": "A"}}
Use only one label from the provided options.
"""

print("Question:", random_question_text)
if len(random_paragraph_text) > 0:
    print("Paragraph:", random_paragraph_text)
print("Options:")
for option_line in random_option_lines:
    print(option_line)

random_messages = []
random_messages.append({"role": "system", "content": "You are a helpful tutor. Return only valid JSON with one field named choice."})
random_messages.append({"role": "user", "content": english_structured_prompt.format(question_text=english_prompt_question_text, options_text=random_options_text)})

random_response = model.create_chat_completion(
    messages=random_messages,
    response_format={
        "type": "json_object",
        "schema": {
            "type": "object",
            "properties": {
                "choice": {
                    "type": "string",
                    "enum": valid_labels
                }
            },
            "required": ["choice"]
        }
    },
    temperature=0.0
)

random_raw_text = random_response["choices"][0]["message"]["content"]
print("Raw model output:", random_raw_text)

random_choice = ""
try:
    random_json = json.loads(random_raw_text)
    random_choice = str(random_json.get("choice", "")).strip().upper()
except Exception as parse_error:
    print("Could not parse JSON directly:", parse_error)

if random_choice not in valid_labels:
    fallback_pattern = r"\b(" + "|".join(valid_labels) + r")\b"
    fallback_match = re.search(fallback_pattern, random_raw_text.upper())
    if fallback_match is not None:
        random_choice = fallback_match.group(1)

print("Parsed choice:", random_choice)

# Pick a random question from the Math question bank
random_index = random.randint(0, len(math_questions) - 1)
random_entry = math_questions[random_index]

# Pull out the question text and choices
random_question_text = random_entry["question"]["question"]
random_choices_raw = random_entry["question"]["choices"]

choice_labels = ["A", "B", "C", "D", "E", "F"]
random_choice_map = {}

if isinstance(random_choices_raw, dict):
    for raw_label in random_choices_raw:
        clean_label = str(raw_label).strip().upper()
        random_choice_map[clean_label] = str(random_choices_raw[raw_label])
elif isinstance(random_choices_raw, list):
    for item_index in range(len(random_choices_raw)):
        if item_index < len(choice_labels):
            mapped_label = choice_labels[item_index]
            random_choice_map[mapped_label] = str(random_choices_raw[item_index])
else:
    random_choice_map["A"] = str(random_choices_raw)

valid_labels = []
for label in choice_labels:
    if label in random_choice_map:
        valid_labels.append(label)

random_option_lines = []
for label in valid_labels:
    random_option_lines.append(label + ": " + random_choice_map[label])
random_options_text = "\n".join(random_option_lines)

math_structured_prompt = """
Question: {question_text}

Options:
{options_text}

Return your answer as JSON with this exact shape:
{{"choice": "A"}}
Use only one label from the provided options.
"""

print("Question:", random_question_text)
print("Options:")
for option_line in random_option_lines:
    print(option_line)

random_messages = []
random_messages.append({"role": "system", "content": "You are a helpful tutor. Return only valid JSON with one field named choice."})
random_messages.append({"role": "user", "content": math_structured_prompt.format(question_text=random_question_text, options_text=random_options_text)})

random_response = model.create_chat_completion(
    messages=random_messages,
    response_format={
        "type": "json_object",
        "schema": {
            "type": "object",
            "properties": {
                "choice": {
                    "type": "string",
                    "enum": valid_labels
                }
            },
            "required": ["choice"]
        }
    },
    temperature=0.0
)

random_raw_text = random_response["choices"][0]["message"]["content"]
print("Raw model output:", random_raw_text)

random_choice = ""
try:
    random_json = json.loads(random_raw_text)
    random_choice = str(random_json.get("choice", "")).strip().upper()
except Exception as parse_error:
    print("Could not parse JSON directly:", parse_error)

if random_choice not in valid_labels:
    fallback_pattern = r"\b(" + "|".join(valid_labels) + r")\b"
    fallback_match = re.search(fallback_pattern, random_raw_text.upper())
    if fallback_match is not None:
        random_choice = fallback_match.group(1)

print("Parsed choice:", random_choice)

difficulty_widget = widgets.Dropdown(
    options=["Easy", "Medium", "Hard"],
    value="Easy",
    description="Difficulty:"
)

section_widget = widgets.Dropdown(
    options=["English", "Math"],
    value="English",
    description="Section:"
)

num_questions_widget = widgets.Dropdown(
    options=[ 2, 4, 8, 10, 12],
    value=4,
    description="# Questions:"
)

display(difficulty_widget)
display(section_widget)
display(num_questions_widget)
print("Choose your settings above, then run the next cell to test the model.")

# Read the widget values
chosen_difficulty = difficulty_widget.value
chosen_section = section_widget.value
chosen_count = num_questions_widget.value

# Pick the right question list based on section
if chosen_section == "English":
    question_pool = english_questions
else:
    question_pool = math_questions

# Filter by difficulty and shuffle so we get a fresh sample each run
filtered_questions = []
for question_entry in question_pool:
    if question_entry.get("difficulty", "").lower() == chosen_difficulty.lower():
        filtered_questions.append(question_entry)

random.shuffle(filtered_questions)
practice_set = filtered_questions[:chosen_count]

# Pull out question text, choices, paragraph, and correct answer for each item
practice_questions = []
practice_choices = []
practice_paragraphs = []
practice_answers = []
for question_entry in practice_set:
    practice_questions.append(question_entry["question"]["question"])
    practice_choices.append(question_entry["question"]["choices"])
    practice_paragraphs.append(question_entry["question"].get("paragraph", ""))
    practice_answers.append(question_entry["question"].get("correct_answer", ""))

print("Built a practice set of", len(practice_questions), chosen_difficulty, chosen_section, "questions.")
for item_index in range(len(practice_questions)):
    print("\nQ" + str(item_index + 1) + ":", practice_questions[item_index])
    if len(practice_paragraphs[item_index]) > 0:
        print("Passage:", practice_paragraphs[item_index])
    print("Choices:", practice_choices[item_index])
    print("Correct answer:", practice_answers[item_index])

batch_results = []
choice_labels = ["A", "B", "C", "D", "E", "F"]

for item_index in range(len(practice_questions)):
    passage_text = practice_paragraphs[item_index]
    if len(passage_text) > 0:
        full_question_text = "Passage: " + passage_text + "\n\nQuestion: " + practice_questions[item_index]
    else:
        full_question_text = "Question: " + practice_questions[item_index]

    question_choices_raw = practice_choices[item_index]
    question_choice_map = {}

    if isinstance(question_choices_raw, dict):
        for raw_label in question_choices_raw:
            clean_label = str(raw_label).strip().upper()
            question_choice_map[clean_label] = str(question_choices_raw[raw_label])
    elif isinstance(question_choices_raw, list):
        for choice_index in range(len(question_choices_raw)):
            if choice_index < len(choice_labels):
                mapped_label = choice_labels[choice_index]
                question_choice_map[mapped_label] = str(question_choices_raw[choice_index])
    else:
        question_choice_map["A"] = str(question_choices_raw)

    valid_labels = []
    for label in choice_labels:
        if label in question_choice_map:
            valid_labels.append(label)

    option_lines = []
    for label in valid_labels:
        option_lines.append(label + ": " + question_choice_map[label])
    options_text = "\n".join(option_lines)

    batch_structured_prompt = """
{question_text}

Options:
{options_text}

Return your answer as JSON with this exact shape:
{{"choice": "A"}}
Use only one label from the provided options.
"""

    batch_messages = []
    batch_messages.append({"role": "system", "content": "You are a helpful tutor. Return only valid JSON with one field named choice."})
    batch_messages.append({"role": "user", "content": batch_structured_prompt.format(question_text=full_question_text, options_text=options_text)})

    batch_response = model.create_chat_completion(
        messages=batch_messages,
        response_format={
            "type": "json_object",
            "schema": {
                "type": "object",
                "properties": {
                    "choice": {
                        "type": "string",
                        "enum": valid_labels
                    }
                },
                "required": ["choice"]
            }
        },
        temperature=0.0
    )

    batch_raw_text = batch_response["choices"][0]["message"]["content"]

    extracted_answer = ""
    try:
        batch_json = json.loads(batch_raw_text)
        extracted_answer = str(batch_json.get("choice", "")).strip().upper()
    except Exception:
        extracted_answer = ""

    if extracted_answer not in valid_labels:
        fallback_pattern = r"\b(" + "|".join(valid_labels) + r")\b"
        fallback_match = re.search(fallback_pattern, batch_raw_text.upper())
        if fallback_match is not None:
            extracted_answer = fallback_match.group(1)

    is_correct = False
    correct_answer_text = str(practice_answers[item_index]).strip().upper()
    if len(correct_answer_text) > 0 and len(extracted_answer) > 0:
        if extracted_answer == correct_answer_text[0]:
            is_correct = True

    batch_results.append({
        "Section": chosen_section,
        "Difficulty": chosen_difficulty,
        "Question": practice_questions[item_index],
        "Correct Answer": practice_answers[item_index],
        "Model Guess": extracted_answer,
        "Correct?": is_correct
    })

batch_results_table = pd.DataFrame(batch_results)
batch_results_table

# Let's build a way to retrieve the question and model answer pairs from the batch results table, so we can review them more easily.
batch_results_table[1:3][["Question", "Model Guess", "Correct Answer"]]

student_reply_three = input("Describe what you tried and whether the model's answer matched the key.\n")
with open('answers.txt', 'a') as answer_file:
    answer_file.write(student_reply_three)
    answer_file.write('\n')

correct_count = 0
for result_row in batch_results:
    if result_row["Correct?"] == True:
        correct_count = correct_count + 1

print("Score:", correct_count, "out of", len(batch_results))
print(batch_results_table)