generate_test_data.py

"""Classic (nested-loop) MCQ test data generator using a 2PL IRT model.

Readable and easy to extend. Best for small/medium datasets and learning the model.
Outputs are written under: output/classic/<timestamp>/
"""

from pathlib import Path
from typing import Optional, Union

import numpy as np
import pandas as pd

from output_utils import resolve_run_dir

# Script identity used for the output folder name
SCRIPT_TAG = "classic"
PROJECT_ROOT = Path(__file__).resolve().parent


def generate_item_analysis_data(
    num_questions: int = 20,
    num_candidates: int = 100,
    seed: int = 42,
    output_dir: Optional[Union[str, Path]] = None,
    timestamped: bool = True,
    responses_filename: str = "candidate_responses.csv",
    answer_key_filename: str = "answer_key.csv",
):
    """Generates realistic MCQ test data for Item Analysis using a 2PL IRT model.

    Saves candidate responses and answer key as two separate CSV files under
    ``output/classic/<timestamp>/`` by default (each run gets its own folder).
    """
    np.random.seed(seed)

    out_dir = resolve_run_dir(
        script_tag=SCRIPT_TAG,
        project_root=PROJECT_ROOT,
        output_dir=output_dir,
        timestamped=timestamped,
    )
    responses_path = out_dir / responses_filename
    answer_key_path = out_dir / answer_key_filename

    options = ["A", "B", "C", "D"]
    question_ids = [f"Q{i + 1}" for i in range(num_questions)]
    candidate_ids = [f"CAND_{i + 1:04d}" for i in range(num_candidates)]

    # 1. Define Answer Key
    correct_answers = np.random.choice(options, size=num_questions).tolist()
    answer_key = dict(zip(question_ids, correct_answers))

    # 2. Assign Latent Traits (Item Response Theory parameters)
    abilities = np.random.normal(loc=0.0, scale=1.0, size=num_candidates)
    difficulties = np.random.uniform(low=-2.0, high=2.0, size=num_questions)
    discriminations = np.random.uniform(low=0.5, high=2.0, size=num_questions)

    # 3. Generate Candidate Responses (nested loops — clear, not vectorized)
    response_matrix = []

    for theta in abilities:
        candidate_responses = []
        for q_idx in range(num_questions):
            a = discriminations[q_idx]
            b = difficulties[q_idx]
            correct_opt = correct_answers[q_idx]

            # Logistic 2PL equation for probability of a correct response
            prob_correct = 1.0 / (1.0 + np.exp(-a * (theta - b)))

            if np.random.rand() < prob_correct:
                chosen_option = correct_opt
            else:
                distractors = [opt for opt in options if opt != correct_opt]
                chosen_option = str(np.random.choice(distractors))

            candidate_responses.append(chosen_option)

        response_matrix.append(candidate_responses)

    # 4. Save Candidate Responses to CSV
    df_responses = pd.DataFrame(
        response_matrix, index=candidate_ids, columns=question_ids
    )
    df_responses.index.name = "Candidate_ID"
    df_responses.to_csv(responses_path)

    # 5. Save Answer Key to a separate CSV
    df_key = pd.DataFrame(
        list(answer_key.items()), columns=["Question_ID", "Correct_Answer"]
    )
    df_key.to_csv(answer_key_path, index=False)

    return df_responses, df_key, responses_path, answer_key_path


if __name__ == "__main__":
    QUESTIONS = 100
    CANDIDATES = 300

    df_resp, df_key, resp_path, key_path = generate_item_analysis_data(
        num_questions=QUESTIONS,
        num_candidates=CANDIDATES,
    )

    print(f"✓ Run folder: {resp_path.parent}")
    print(
        f"✓ Saved candidate responses ({CANDIDATES} candidates x {QUESTIONS} questions)"
        f" -> '{resp_path}'"
    )
    print(f"✓ Saved official answer key ({QUESTIONS} questions) -> '{key_path}'\n")

    print("--- Sample Candidate Responses ---")
    print(df_resp.head())

    print("\n--- Answer Key Data ---")
    print(df_key.head())