adaptive_testing.py

from __future__ import annotations

import os
import csv
import argparse
from pathlib import Path
import numpy as np
import pandas as pd

from maqua_paths import FORMATTED_DATA, POLYTOMIZED_DATA, ADAPTIVE_OUTPUTS, ensure_dir

# Ordered question codes used by the cmodel specification (1..48)
# Derived from formatted_data/questions.csv, excluding binary items and PTSD3.
DEFAULT_QUESTION_ORDER = [
    "A1",
    "A3",
    "A4",
    "ADHD1",
    "ADHD2",
    "ASD2",
    "ASD3",
    "ASD4",
    "ASD5",
    "ASD6",
    "BD2",
    "BD3",
    "ED1",
    "ED2",
    "ED3",
    "ED4",
    "ED5",
    "ED6",
    "G1",
    "G10",
    "G12",
    "G2",
    "G3",
    "G4",
    "G5",
    "G6",
    "G7",
    "G8",
    "G9",
    "G91",
    "OCD1",
    "OCD2",
    "OCD3",
    "OMD1",
    "OMD2",
    "OMD3",
    "OMD4",
    "OMD5",
    "OMD6",
    "PTSD1",
    "PTSD2",
    "SUB1",
    "SUB2",
    "SUB3",
    "SUB4",
    "SUB5",
    "SUB6",
    "nse",
]

EXCLUDED_CODES = {"A2", "ASD1", "BD1", "G11", "SUB7", "PTSD3"}

FACTOR_INDEX_MAP = {
    "F1": [
        34, 35, 36, 37, 38, 39, 1, 2, 3, 11, 12, 6, 7, 8, 9, 10,
        31, 32, 33, 4, 5, 40, 41, 48, 13, 14, 15, 16, 17, 18,
        19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30,
    ],
    "F2": [
        11, 12, 42, 43, 44, 45, 46, 47, 31, 32, 33, 13, 14, 15, 16, 17,
        18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30,
    ],
}


def _import_rpy2():
    """Import rpy2 lazily so --help works even when rpy2 is unavailable."""
    try:
        from rpy2.robjects import pandas2ri  # type: ignore
        from rpy2.robjects import default_converter, conversion  # type: ignore
        from rpy2.robjects.packages import importr  # type: ignore
        import rpy2.robjects as robjects  # type: ignore
        import rpy2.robjects.packages as rpackages  # type: ignore
        from rpy2.robjects.vectors import StrVector  # type: ignore
    except Exception as exc:
        raise RuntimeError(
            "rpy2 is required for adaptive testing. Install rpy2 and ensure R packages mirt/mirtCAT are available."
        ) from exc
    return pandas2ri, default_converter, conversion, importr, robjects, rpackages, StrVector


def _get_question_code_order() -> list[str]:
    """Load question order from formatted_data/questions.csv with fallback."""
    questions_path = FORMATTED_DATA / "questions.csv"
    if questions_path.exists():
        try:
            df = pd.read_csv(questions_path)
            code_col = "code" if "code" in df.columns else df.columns[0]
            codes = df[code_col].astype(str).tolist()
            ordered = [code for code in codes if code not in EXCLUDED_CODES]
            if len(ordered) == len(DEFAULT_QUESTION_ORDER):
                return ordered
            print(
                f"WARNING: Unexpected question count ({len(ordered)}). "
                f"Falling back to default order ({len(DEFAULT_QUESTION_ORDER)})."
            )
        except Exception as exc:
            print(f"WARNING: Failed to load question order from questions.csv: {exc}")
    return DEFAULT_QUESTION_ORDER


def _prepare_irt_matrix(df: pd.DataFrame, label: str, question_order: list[str]) -> pd.DataFrame:
    """Convert long-form discretized data to wide numeric matrix for MIRT."""
    if "user_id" not in df.columns:
        raise ValueError(f"{label} data missing required column: user_id")

    if {"question_code", "discrete_score"}.issubset(df.columns):
        work = df.copy()
        work["question_code"] = work["question_code"].astype(str)
        dup_mask = work.duplicated(subset=["user_id", "question_code"], keep=False)
        if dup_mask.any():
            dup_count = int(dup_mask.sum())
            print(
                f"WARNING: {label} has {dup_count} duplicate user/question rows; "
                "averaging discrete_score."
            )
        wide = (
            work.pivot_table(
                index="user_id",
                columns="question_code",
                values="discrete_score",
                aggfunc="mean",
            )
            .reset_index()
        )
    else:
        wide = df.copy()

    response_cols = [col for col in wide.columns if col != "user_id"]
    if question_order:
        missing = [code for code in question_order if code not in wide.columns]
        if missing:
            preview = ", ".join(missing[:5])
            print(
                f"WARNING: {label} missing {len(missing)} question codes "
                f"(e.g., {preview}). Filling with 0."
            )
            for code in missing:
                wide[code] = np.nan

        extra = [col for col in response_cols if col not in question_order]
        if extra:
            preview = ", ".join(extra[:5])
            print(
                f"WARNING: {label} has {len(extra)} extra columns not in question list "
                f"(e.g., {preview}). Dropping."
            )

        wide = wide[["user_id"] + question_order]
        response_cols = question_order

    # Coerce to numeric
    for col in response_cols:
        wide[col] = pd.to_numeric(wide[col], errors="coerce")

    missing_values = int(wide[response_cols].isna().sum().sum())
    if missing_values > 0:
        print(
            f"WARNING: {label} contains {missing_values} non-numeric/NaN entries; "
            "filling with 0."
        )
        wide[response_cols] = wide[response_cols].fillna(0)

    # If discretized scores, round to integer levels
    if "discrete_score" in df.columns:
        values = wide[response_cols].to_numpy()
        if not np.all(np.isclose(values, np.rint(values))):
            print(f"WARNING: {label} has non-integer discrete scores; rounding.")
        wide[response_cols] = np.rint(values).astype(int)

    return wide


def _build_cmodel(question_order: list[str], active_codes: set[str]) -> str:
    """Build mirt cmodel string using item names after filtering."""
    lines = []
    for factor, indices in FACTOR_INDEX_MAP.items():
        codes = []
        for idx in indices:
            if 1 <= idx <= len(question_order):
                code = question_order[idx - 1]
                if code in active_codes:
                    codes.append(code)
        if codes:
            lines.append(f"{factor}=" + ",".join(codes))
        else:
            print(f"WARNING: Factor {factor} has no items after filtering.")

    if not lines:
        return "1"
    return "\n".join(lines)


def _align_response_categories(
    train_df: pd.DataFrame,
    test_df: pd.DataFrame,
) -> tuple[pd.DataFrame, pd.DataFrame]:
    """Align test response categories to training categories for each item."""
    train_aligned = train_df.copy()
    test_aligned = test_df.copy()

    for col in train_aligned.columns:
        train_vals = train_aligned[col].to_numpy()
        train_cats = np.unique(train_vals)

        # Ensure integer-like categories
        if not np.all(np.isclose(train_cats, np.rint(train_cats))):
            train_cats = np.unique(np.rint(train_vals).astype(int))
            train_aligned[col] = np.rint(train_aligned[col]).astype(int)
            test_aligned[col] = np.rint(test_aligned[col]).astype(int)

        train_cats = np.unique(train_aligned[col].to_numpy())
        if train_cats.size <= 1:
            continue

        expected = np.arange(train_cats.size)
        needs_remap = not np.array_equal(train_cats, expected)

        cat_to_idx = {cat: idx for idx, cat in enumerate(train_cats)}

        if needs_remap:
            print(
                f"WARNING: {col} has non-contiguous categories {train_cats.tolist()} "
                "in training; remapping to consecutive integers."
            )

        if needs_remap:
            train_aligned[col] = train_aligned[col].map(cat_to_idx).astype(int)

        test_vals = test_aligned[col].to_numpy()
        unseen_mask = ~np.isin(test_vals, train_cats)
        unseen_count = int(unseen_mask.sum())
        if unseen_count:
            print(
                f"WARNING: {col} has {unseen_count} test responses outside training categories; "
                "mapping to nearest observed category."
            )

        def _map_val(x):
            if x in cat_to_idx:
                return cat_to_idx[x]
            nearest = min(train_cats, key=lambda c: (abs(c - x), c))
            return cat_to_idx[nearest]

        test_aligned[col] = [int(_map_val(x)) for x in test_vals]

    return train_aligned, test_aligned


def _find_latest_discretized_dir(base_dir: Path) -> Path | None:
    candidates = [d for d in base_dir.iterdir() if d.is_dir() and d.name.endswith("_discretized")]
    if not candidates:
        return None
    return max(candidates, key=lambda p: p.stat().st_mtime)


def _resolve_input_dir(input_dir: str | None) -> Path | None:
    if input_dir:
        path = Path(input_dir)
        return path if path.is_dir() else None
    return _find_latest_discretized_dir(POLYTOMIZED_DATA)

def adaptive_testing(
    train_data,
    test_data,
    question_order: list[str] | None = None,
    model_type="QMCEM",
    adaptive_order="Drule",
    output_dir="outputs/",
):
    """
    Perform multidimensional adaptive testing using MIRT and mirtCAT via R.

    Args:
        train_data (pd.DataFrame): Training data with 'user_id' and item responses.
        test_data (pd.DataFrame): Test data with 'user_id' and item responses.
        model_type (str): IRT estimation method, e.g., "QMCEM" or "EM"
        adaptive_order (str): Adaptive item selection criterion, e.g., "Drule" for D-optimality.
        output_dir (str): Directory to save output CSVs.

    Saves:
        - loadings.csv: Factor loadings for the trained model.
        - thetas_N.csv: Estimated thetas (trait scores) for N adaptively selected items per user.
    """
    pandas2ri, default_converter, conversion, importr, robjects, _, _ = _import_rpy2()
    # Preprocess: Fill missing entries and separate user IDs
    train_data_filled = train_data.fillna(0).copy()
    test_data_filled = test_data.fillna(0).copy()
    train_user = train_data_filled["user_id"].tolist()
    test_user = test_data_filled["user_id"].tolist()
    train_data_filled.drop(columns=["user_id"], inplace=True)
    test_data_filled.drop(columns=["user_id"], inplace=True)

    if question_order is None:
        question_order = list(train_data_filled.columns)

    constant_cols = [
        col for col in train_data_filled.columns
        if train_data_filled[col].nunique(dropna=False) <= 1
    ]
    if constant_cols:
        print(
            "WARNING: Dropping items with only one response category in training data: "
            + ", ".join(constant_cols)
        )
        train_data_filled.drop(columns=constant_cols, inplace=True)
        test_data_filled.drop(columns=constant_cols, inplace=True, errors="ignore")
        question_order = [code for code in question_order if code not in constant_cols]

    if train_data_filled.empty:
        raise ValueError("No valid item columns remain after filtering.")

    train_data_filled, test_data_filled = _align_response_categories(
        train_data_filled, test_data_filled
    )

    print(f"Number of items: {len(train_data_filled.columns)}")

    cmodel = _build_cmodel(question_order, set(train_data_filled.columns))

    # Import necessary R packages
    mirt = importr('mirt')
    mirtCAT = importr('mirtCAT')

    # Fit multidimensional IRT model
    converter = default_converter + pandas2ri.converter

    if model_type == "QMCEM":
        model_spec = 1 if cmodel.strip() == "1" else cmodel
        with conversion.localconverter(converter):
            train_data_r = conversion.py2rpy(train_data_filled)
        model = mirt.mirt(train_data_r, model=model_spec, method=model_type)
    elif model_type == "EM":
        with conversion.localconverter(converter):
            train_data_r = conversion.py2rpy(train_data_filled)
        model = mirt.mirt(train_data_r, model=1, method=model_type)
    else:
        raise ValueError(f"Unsupported model_type {model_type}")

    # Extract and save factor loadings
    loadings = mirt.extract_mirt(model, "F")
    loadings_np = np.array(loadings)
    os.makedirs(output_dir, exist_ok=True)
    loadings_filepath = os.path.join(output_dir, "loadings.csv")
    np.savetxt(loadings_filepath, loadings_np, delimiter=",")
    print(f"Saved factor loadings to {loadings_filepath}")

    # Adaptive testing per test user for each item count
    start_theta = 0
    for num_item in range(1, len(train_data_filled.columns) + 1):
        thetas_filepath = os.path.join(output_dir, f"thetas_{num_item}.csv")
        with open(thetas_filepath, "w", newline='') as csvfile:
            writer = csv.writer(csvfile)
            for i in range(test_data_filled.shape[0]):
                design = robjects.ListVector({"min_SEM": 0.01, "max_items": num_item})
                with conversion.localconverter(converter):
                    local_pattern_r = conversion.py2rpy(test_data_filled.iloc[i])
                result = mirtCAT.mirtCAT(
                    mo=model,
                    local_pattern=local_pattern_r,
                    criteria=adaptive_order,
                    method='MAP',
                    start_theta=start_theta,
                    design=design,
                    start_item=adaptive_order
                )
                theta = result.rx2("thetas")[0]
                theta = np.array(theta).flatten().tolist()
                items = result.rx2("items_answered")
                items = np.array(items).flatten().tolist()
                writer.writerow([test_user[i]] + theta + items)
    print("Adaptive testing complete.")


def main():
    parser = argparse.ArgumentParser(description="Run adaptive testing (mirt/mirtCAT) on discretized data.")
    parser.add_argument("--input-dir", type=str, default=None, help="Directory with dev_*.csv/test_*.csv files.")
    parser.add_argument("--output-dir", type=str, default=None, help="Base output directory for loadings/thetas.")
    parser.add_argument("--model-type", type=str, default="QMCEM", choices=["QMCEM", "EM"])
    parser.add_argument("--adaptive-order", type=str, default="Drule", choices=["Drule", "random"])
    parser.add_argument("--n-folds", type=int, default=9)
    parser.add_argument("--run-both", action="store_true", help="Run both random and Drule orders.")
    parser.add_argument("--flat-output", action="store_true", help="Write outputs directly under order dir (legacy).")
    args = parser.parse_args()

    pandas2ri, _, _, importr, robjects, rpackages, StrVector = _import_rpy2()

    # get the adaptive testing packages on R
    base = importr('base')
    utils = importr('utils')
    utils.chooseCRANmirror(ind=1)
    packnames = ('mirt', 'stringr', 'mirtCAT')

    # Use a writable user library for package installation on shared systems.
    user_r_lib = Path.home() / ".R" / "library"
    user_r_lib.mkdir(parents=True, exist_ok=True)
    robjects.r(f'.libPaths(c("{str(user_r_lib)}", .libPaths()))')

    names_to_install = [x for x in packnames if not rpackages.isinstalled(x)]

    if len(names_to_install) > 0:
        print(f"Installing missing R packages into: {user_r_lib}")
        utils.install_packages(StrVector(names_to_install), lib=str(user_r_lib))

    try:
        mirt = rpackages.importr('mirt')
        stringr = rpackages.importr('stringr')
        mirtCAT = rpackages.importr('mirtCAT')
    except Exception as e:
        msg = str(e)
        if "CXXABI_1.3.15" in msg or "libstdc++.so.6" in msg:
            print("\nERROR: Failed to load R package 'mirtCAT' due to libstdc++ ABI mismatch.")
            print("This usually means your runtime libstdc++ is older than the one used to build mirtCAT.")
            print("\nSuggested fixes (pick one):")
            print("1) Use conda's libstdc++ for R at runtime:")
            print("   conda install -c conda-forge libstdcxx-ng libgcc-ng")
            print("   export LD_LIBRARY_PATH=\"$CONDA_PREFIX/lib:$LD_LIBRARY_PATH\"")
            print("\n2) Reinstall R packages from conda-forge:")
            print("   conda install -c conda-forge r-mirtcat r-mirt")
            print("\n3) Upgrade system libstdc++ (requires sudo):")
            print("   sudo apt-get install libstdc++6")
            print("\nThen re-run adaptive_testing.py.")
        else:
            print("\nERROR: Failed to load required R packages (mirt/mirtCAT).")
            print(f"Details: {e}")
        raise


    input_dir = _resolve_input_dir(args.input_dir)
    if input_dir is None:
        print("No discretized input directory found. Provide --input-dir or run discretize_question_scores.py first.")
        return 1

    question_order = _get_question_code_order()
    print(f"Question order length: {len(question_order)} items")

    output_base = Path(args.output_dir) if args.output_dir else (ensure_dir(ADAPTIVE_OUTPUTS) / input_dir.name)

    orders = ["random", "Drule"] if args.run_both else [args.adaptive_order]
    for order in orders:
        output_dir = ensure_dir(output_base / order)
        print(f"\nRunning adaptive testing: input={input_dir}, order={order}, output={output_dir}")

        for i in range(args.n_folds):
            dev_path = input_dir / f"dev_{i}.csv"
            test_path = input_dir / f"test_{i}.csv"
            if not dev_path.exists() or not test_path.exists():
                print(f"Fold {i}: missing dev/test files, skipping.")
                continue

            dev_data = pd.read_csv(dev_path)
            dev_data = dev_data[[col for col in dev_data.columns if not col.startswith("pred_score_")]]
            test_data = pd.read_csv(test_path)
            test_data = test_data[[col for col in test_data.columns if not col.startswith("pred_score_")]]

            dev_data = _prepare_irt_matrix(dev_data, f"dev fold {i}", question_order)
            test_data = _prepare_irt_matrix(test_data, f"test fold {i}", question_order)

            fold_output_dir = output_dir if args.flat_output else ensure_dir(output_dir / f"fold_{i}")

            adaptive_testing(
                dev_data,
                test_data,
                question_order=question_order,
                output_dir=str(fold_output_dir),
                adaptive_order=order,
                model_type=args.model_type,
            )

    return 0


if __name__ == "__main__":
    raise SystemExit(main())