preprocess_individual.py

# scripts/preprocess_individual.py
import argparse
import csv
import json
import sys
from pathlib import Path
from typing import Dict, Optional, List, Tuple, Any

import numpy as np
import pandas as pd

HERE = Path(__file__).resolve().parent
CONFIG_DIR = HERE.parent / "config"

TRUST_DICT_PATH = CONFIG_DIR / "Trust-DataDictionaryfMRI.csv"
SOCIAL_DICT_PATH = CONFIG_DIR / "SRA-DataDictionaryfMRI.csv"
OL_DICT_PATH = CONFIG_DIR / "OL-DataDictionaryfMRI.csv"

def log(msg: str) -> None:
    print(msg, flush=True)

# ------------------------------
# Robust CSV reader (NO SKIPS)
# ------------------------------
def read_trials_csv(path: Path) -> pd.DataFrame:
    try:
        return pd.read_csv(path)
    except Exception as e:
        log(f"⚠️ read_csv default failed: {e}")
        log("🔧 Falling back to row-repair parser (no skips).")

    text = path.read_text(encoding="utf-8", errors="replace")
    lines = text.splitlines()
    if not lines:
        return pd.DataFrame()

    header = next(csv.reader([lines[0]]))
    expected = len(header)
    repaired_rows: List[List[str]] = []

    def looks_like_listish(s: str) -> bool:
        s = (s or "").strip()
        return ("[" in s and "]" in s) or ("{" in s and "}" in s) or ("'" in s) or ('"' in s)

    for _lineno, line in enumerate(lines[1:], start=2):
        if not line.strip():
            repaired_rows.append([""] * expected)
            continue

        tokens = next(csv.reader([line]))

        if len(tokens) == expected:
            repaired_rows.append(tokens)
            continue

        if len(tokens) < expected:
            repaired_rows.append(tokens + [""] * (expected - len(tokens)))
            continue

        extra = len(tokens) - expected
        merge_col = None
        best_score = -1.0

        for k in range(0, expected):
            left = tokens[:k]
            mid = tokens[k : k + extra + 1]
            right = tokens[k + extra + 1 :]

            if len(left) + 1 + len(right) != expected:
                continue

            merged = ",".join(mid)
            score = 0.0
            if looks_like_listish(merged):
                score += 5.0
            score += k / max(1, expected - 1)

            if score > best_score:
                best_score = score
                merge_col = k

        if merge_col is None:
            merge_col = expected - 1

        left = tokens[:merge_col]
        mid = tokens[merge_col : merge_col + extra + 1]
        right = tokens[merge_col + extra + 1 :]
        merged = ",".join(mid)

        fixed = left + [merged] + right

        if len(fixed) > expected:
            fixed = fixed[: expected - 1] + [",".join(fixed[expected - 1 :])]
        elif len(fixed) < expected:
            fixed = fixed + [""] * (expected - len(fixed))

        repaired_rows.append(fixed)

    return pd.DataFrame(repaired_rows, columns=header)

# ------------------------------
# Utilities
# ------------------------------
def trim_extrainfo(df: pd.DataFrame) -> pd.DataFrame:
    if df.empty:
        return df
    first = df.columns[0]
    series = df[first].astype(str).str.strip().str.lower()
    idx = series[series.eq("extrainfo")].index
    if len(idx) > 0:
        return df.loc[: idx[0] - 1].copy()
    return df

def integer_like(series: pd.Series) -> bool:
    try:
        s = pd.to_numeric(series, errors="coerce").dropna()
        if s.empty:
            return False
        arr = s.to_numpy(dtype="float64", copy=False)
        return np.allclose(arr, np.round(arr))
    except Exception:
        return False

def cast_integer_like(df: pd.DataFrame) -> pd.DataFrame:
    out = df.copy()
    for col in out.columns:
        try:
            if integer_like(out[col]):
                out[col] = pd.to_numeric(out[col], errors="coerce").round(0).astype("Int64")
        except Exception:
            pass
    return out

def drop_metadata_rows(out: pd.DataFrame, meta_keep=("pOrder",)) -> pd.DataFrame:
    if out.empty:
        return out
    work = out.replace(r"^\s*$", pd.NA, regex=True).infer_objects(copy=False)
    non_meta = [c for c in work.columns if c not in meta_keep]
    if not non_meta:
        return out
    cleaned = work.dropna(how="all", subset=non_meta)
    cleaned.reset_index(drop=True, inplace=True)
    return cleaned

def detect_task_from_filename(infile: Path) -> str:
    name = infile.name.lower()
    if "_socialra_" in name or "socialra" in name or "_sra_" in name:
        return "socialRA"
    if "_th_" in name:
        return "th"
    if "_tm_" in name:
        return "tm"
    if "_ol_" in name:
        return "ol"
    parent = infile.parent.name.lower()
    if parent in {"th", "tm", "ol", "socialra"}:
        return parent
    return "th"

# ------------------------------
# Fuzzy column resolver (fixes empty timers)
# ------------------------------
def _canon(s: str) -> str:
    return "".join(ch.lower() for ch in str(s) if ch.isalnum())

def resolve_col(df: pd.DataFrame, desired: str) -> Optional[str]:
    """
    Find an existing column that matches desired ignoring case, underscores, dots, etc.
    """
    if desired in df.columns:
        return desired
    want = _canon(desired)
    for c in df.columns:
        if _canon(c) == want:
            return c
    # fallback: sometimes extra suffix/prefix
    for c in df.columns:
        if want in _canon(c) or _canon(c) in want:
            return c
    return None

def get_col(df: pd.DataFrame, desired: str) -> pd.Series:
    c = resolve_col(df, desired)
    if c is None:
        return pd.Series(pd.NA, index=df.index)
    return df[c]

# ------------------------------
# Dictionary keep list (Required==yes)
# ------------------------------
def _normalize_cols(cols: List[str]) -> Dict[str, str]:
    return {c: str(c).strip().lower() for c in cols}

def load_keep_list_from_dict(path: Path, required_col: str = "Required") -> List[str]:
    if not path.exists():
        raise FileNotFoundError(f"Dictionary not found: {path}")
    dd = pd.read_csv(path, engine="python", on_bad_lines="skip")
    cols = list(dd.columns)
    norm = _normalize_cols(cols)

    key_col_real = next((c for c in cols if norm[c] in {"variable", "name"}), None)
    if key_col_real is None:
        raise ValueError(f"{path.name} missing key column (Variable/Name). Found columns: {cols}")

    req_real = next((c for c in cols if norm[c] == required_col.lower()), None)
    if req_real is not None:
        dd = dd[dd[req_real].astype(str).str.strip().str.lower().isin(
            ["yes", "y", "1", "true", "required"]
        )].copy()

    keep: List[str] = []
    for v in dd[key_col_real].tolist():
        s = str(v).strip()
        if s and s.lower() != "nan":
            keep.append(s)

    if not keep:
        raise ValueError(f"{path.name} produced empty keep-list.")
    return keep

# ------------------------------
# TRUST (th/tm)
# ------------------------------
def preprocess_trust(infile: Path) -> pd.DataFrame:
    df_raw = trim_extrainfo(read_trials_csv(infile))
    keep = load_keep_list_from_dict(TRUST_DICT_PATH)

    out = pd.DataFrame(index=df_raw.index)
    for col in keep:
        out[col] = get_col(df_raw, col)

    # trust wants trials.thisN 1-based
    if "trials.thisN" in out.columns:
        out["trials.thisN"] = (pd.to_numeric(out["trials.thisN"], errors="coerce") + 1).round(0).astype("Int64")

    out = drop_metadata_rows(out, meta_keep=("pOrder",))
    out = cast_integer_like(out)
    return out

# ------------------------------
# OL (fix missing cols + empty scannerTimer_exp_End)
# ------------------------------
def preprocess_ol(infile: Path) -> pd.DataFrame:
    df_raw = trim_extrainfo(read_trials_csv(infile))
    keep = load_keep_list_from_dict(OL_DICT_PATH)

    out = pd.DataFrame(index=df_raw.index)
    for col in keep:
        out[col] = get_col(df_raw, col)  # <-- fuzzy match fixes scannerTimer_exp_End

    # ---- add your requested extra columns (even if not "Required") ----
    # runID: if absent, set = runNb (or keep if exists)
    if "runID" not in out.columns:
        out["runID"] = get_col(df_raw, "runID")
        if out["runID"].isna().all():
            out["runID"] = out.get("runNb", get_col(df_raw, "runNb"))

    # trials.thisN (raw) and trials.thisIindex (raw)
    if "trials.thisN" not in out.columns:
        out["trials.thisN"] = get_col(df_raw, "trials.thisN")
        if out["trials.thisN"].isna().all():
            # derive from trials.thisTrialN if present (often 1-based already)
            ttn = pd.to_numeric(get_col(df_raw, "trials.thisTrialN"), errors="coerce")
            if not ttn.isna().all():
                out["trials.thisN"] = (ttn - 1).astype("Int64")

    if "trials.thisIindex" not in out.columns:
        out["trials.thisIindex"] = get_col(df_raw, "trials.thisIindex")
        if out["trials.thisIindex"].isna().all():
            # sometimes stimulus index is in 'X'
            out["trials.thisIindex"] = get_col(df_raw, "X")

    # choice_keys (usually key_resp.keys)
    if "choice_keys" not in out.columns:
        out["choice_keys"] = get_col(df_raw, "choice_keys")
        if out["choice_keys"].isna().all():
            out["choice_keys"] = get_col(df_raw, "key_resp.keys")
            if out["choice_keys"].isna().all():
                out["choice_keys"] = get_col(df_raw, "key_resp_keys")

    # if trials.thisN exists and you want it 1-based like trust, do it:
    if "trials.thisN" in out.columns:
        tn = pd.to_numeric(out["trials.thisN"], errors="coerce")
        # only shift if looks 0-based (contains 0)
        if (tn == 0).any():
            out["trials.thisN"] = (tn + 1).round(0).astype("Int64")

    out = drop_metadata_rows(out, meta_keep=("pOrder",))
    out = cast_integer_like(out)

    # Put the requested added columns early (nice UX)
    front = [c for c in ["runID", "runNb", "trials.thisN", "trials.thisIindex", "choice_keys"] if c in out.columns]
    rest = [c for c in out.columns if c not in front]
    out = out[front + rest]

    return out

# ---------------------------------
# SocialRA — helpers (unchanged core)
# ---------------------------------
def _derive_run_from_expname(expname: pd.Series) -> pd.Series:
    exp = expname.astype(str).str.lower()
    run = pd.Series(pd.NA, index=expname.index, dtype="Int64")
    run = run.mask(exp.str.contains("social_2"), 3)
    run = run.mask(exp.str.contains("social_1"), 2)
    run = run.mask(exp.str.contains("self"), 1)
    return run

def _as_list(x: Any) -> List[Any]:
    if isinstance(x, list):
        return x
    if pd.isna(x):
        return []
    s = str(x).strip()
    if s.startswith("[") and s.endswith("]"):
        try:
            val = json.loads(s)
            return val if isinstance(val, list) else []
        except Exception:
            s2 = s.strip("[]").strip()
            return [t.strip().strip("'").strip('"') for t in s2.split(",")] if s2 else []
    return [s]

def _first_number(x: Any) -> Optional[float]:
    if isinstance(x, (int, float, np.integer, np.floating)) and not pd.isna(x):
        return float(x)
    for t in _as_list(x):
        try:
            return float(t)
        except Exception:
            continue
    return np.nan

def _norm_choice(x: pd.Series) -> pd.Series:
    result = []
    for v in x:
        lis = _as_list(v)
        tok = str(lis[0]).lower() if lis else ""
        if "left" in tok:
            result.append("left")
        elif "right" in tok:
            result.append("right")
        else:
            result.append(pd.NA)
    return pd.Series(result, index=x.index, dtype="string")

def _max_from_list_column(col: pd.Series) -> pd.Series:
    out = []
    for v in col:
        nums = [_first_number(t) for t in _as_list(v)]
        nums = [n for n in nums if pd.notna(n)]
        out.append(max(nums) if nums else pd.NA)
    return pd.Series(out, index=col.index, dtype="Float64")

def _choice_to_gamble(choice: pd.Series, safe_left_flag: pd.Series) -> pd.Series:
    choice = choice.astype("string")
    safe_left = pd.to_numeric(safe_left_flag, errors="coerce")
    gamble = pd.Series(pd.NA, index=choice.index, dtype="Int64")
    mask_sl = safe_left.eq(0)
    gamble[mask_sl.fillna(False) & choice.eq("right")] = 1
    gamble[mask_sl.fillna(False) & choice.eq("left")]  = 0
    mask_sr = safe_left.eq(1)
    gamble[mask_sr.fillna(False) & choice.eq("left")]  = 1
    gamble[mask_sr.fillna(False) & choice.eq("right")] = 0
    return gamble

def _get_first_existing(df: pd.DataFrame, candidates: List[str]) -> pd.Series:
    for name in candidates:
        c = resolve_col(df, name)
        if c is not None:
            return df[c]
    return pd.Series(pd.NA, index=df.index)

# ---------------------------------
# SocialRA — preprocess (fix ITI timer)
# ---------------------------------
def preprocess_social(infile: Path) -> pd.DataFrame:
    df_raw = trim_extrainfo(read_trials_csv(infile))
    keep = load_keep_list_from_dict(SOCIAL_DICT_PATH)

    out = pd.DataFrame(index=df_raw.index)
    for col in keep:
        out[col] = get_col(df_raw, col)

    # ✅ Fix scannerTimer_jitterself_Start empty (fuzzy fill)
    if "scannerTimer_jitterself_Start" in out.columns:
        if out["scannerTimer_jitterself_Start"].isna().all():
            out["scannerTimer_jitterself_Start"] = get_col(df_raw, "scannerTimer_jitterself_Start")

    block = pd.to_numeric(out.get("block", _get_first_existing(df_raw, ["block"])), errors="coerce")
    expName = out.get("expName", _get_first_existing(df_raw, ["expName"]))

    if "run" in out.columns:
        need = out["run"].isna()
        if need.any():
            out.loc[need, "run"] = _derive_run_from_expname(expName)[need]
    else:
        out["run"] = _derive_run_from_expname(expName)

    if "trials_self" not in out.columns or out["trials_self"].isna().all():
        out["trials_self"] = pd.to_numeric(_get_first_existing(df_raw, ["trials.thisRepN", "trials_thisRepN"]), errors="coerce")
    if "trials_social" not in out.columns or out["trials_social"].isna().all():
        out["trials_social"] = pd.to_numeric(_get_first_existing(df_raw, ["trials_2.thisRepN", "trials_2_thisRepN"]), errors="coerce")

    if "choice_self" not in out.columns or out["choice_self"].isna().all():
        out["choice_self"] = _norm_choice(_get_first_existing(df_raw, ["key_resp.keys", "key_resp_keys", "choice_self"]))
    else:
        out["choice_self"] = _norm_choice(out["choice_self"])
    if "choice_social" not in out.columns or out["choice_social"].isna().all():
        out["choice_social"] = _norm_choice(_get_first_existing(df_raw, ["key_resp_3.keys", "key_resp_3_keys", "choice_social"]))
    else:
        out["choice_social"] = _norm_choice(out["choice_social"])

    if "rt_self" not in out.columns or out["rt_self"].isna().all():
        raw = _get_first_existing(df_raw, ["key_resp.rt", "key_resp_rt", "rt_self"])
        out["rt_self"] = pd.Series([_first_number(v) for v in raw], index=raw.index, dtype="Float64")
    else:
        out["rt_self"] = pd.Series([_first_number(v) for v in out["rt_self"]], index=out.index, dtype="Float64")

    if "rt_social" not in out.columns or out["rt_social"].isna().all():
        raw = _get_first_existing(df_raw, ["key_resp_3.rt", "key_resp_3_rt", "rt_social"])
        out["rt_social"] = pd.Series([_first_number(v) for v in raw], index=raw.index, dtype="Float64")
    else:
        out["rt_social"] = pd.Series([_first_number(v) for v in out["rt_social"]], index=out.index, dtype="Float64")

    if "expName" in out.columns:
        out["trials_inblock"] = (out.groupby("expName", dropna=False).cumcount() + 1).astype("Int64")
    else:
        out["trials_inblock"] = (out.groupby(["block", "run"], dropna=False).cumcount() + 1).astype("Int64")

    if "slider_key_list" not in out.columns or out["slider_key_list"].isna().all():
        out["slider_key_list"] = _get_first_existing(df_raw, ["key_resp_2.keys", "key_resp_2_keys", "slider_key_list"])
    if "slider_rt_list" not in out.columns or out["slider_rt_list"].isna().all():
        out["slider_rt_list"] = _get_first_existing(df_raw, ["key_resp_2rt", "key_resp_2.rt", "slider_rt_list"])
    out["slider_rt"] = _max_from_list_column(out["slider_rt_list"]).where(block == 2, pd.NA)

    gl_self_raw = _get_first_existing(df_raw, ["gamble_left_self", "gamble_left"])
    gl_social_raw = _get_first_existing(df_raw, ["gamble_left_social", "gamble_left"])
    gl_self = pd.to_numeric(gl_self_raw, errors="coerce").where(block == 1, pd.NA)
    gl_soc = pd.to_numeric(gl_social_raw, errors="coerce").where(block == 2, pd.NA)
    out["gamble_left_self"] = pd.to_numeric(out.get("gamble_left_self", gl_self), errors="coerce").where(block == 1, pd.NA).astype("Float64")
    out["gamble_left_social"] = pd.to_numeric(out.get("gamble_left_social", gl_soc), errors="coerce").where(block == 2, pd.NA).astype("Float64")

    gl = pd.Series(pd.NA, index=df_raw.index)
    gl = gl.mask(block == 1, out["gamble_left_self"])
    gl = gl.mask(block == 2, out["gamble_left_social"])
    out["gamble_left"] = pd.to_numeric(gl, errors="coerce").round(0).astype("Int64")

    sr = pd.to_numeric(out.get("slider_response", _get_first_existing(df_raw, ["slider_response"])), errors="coerce")
    map_prob = {1: 0.10, 2: 0.26, 3: 0.42, 4: 0.58, 5: 0.74, 6: 0.90}
    out["info_probability"] = sr.map(map_prob).where(block == 2, pd.NA)

    cs = pd.to_numeric(out.get("ch_safe", _get_first_existing(df_raw, ["ch_safe"])), errors="coerce")
    cr = pd.to_numeric(out.get("ch_risky", _get_first_existing(df_raw, ["ch_risky"])), errors="coerce")
    si = pd.to_numeric(out.get("seeInfo", _get_first_existing(df_raw, ["seeInfo"])), errors="coerce")
    out["ch_safe_seen"] = cs.where((si == 1) & (block == 2), pd.NA)
    out["ch_risky_seen"] = cr.where((si == 1) & (block == 2), pd.NA)

    gamble_self = _choice_to_gamble(out["choice_self"], out["gamble_left_self"]).where(block == 1, pd.NA)
    gamble_social = _choice_to_gamble(out["choice_social"], out["gamble_left_social"]).where(block == 2, pd.NA)

    combined_choice = pd.Series(pd.NA, index=out.index, dtype="Int64")
    combined_choice = combined_choice.mask(block == 1, gamble_self)
    combined_choice = combined_choice.mask(block == 2, gamble_social)
    out["gamble_choice"] = combined_choice.astype("Int64")

    grt = pd.Series(pd.NA, index=out.index, dtype="Float64")
    grt = grt.mask(block == 1, out["rt_self"])
    grt = grt.mask(block == 2, out["rt_social"])
    grt = grt.where(grt > 0, pd.NA)
    out["gamble_rt"] = grt

    out["missed_gamble"] = out["gamble_choice"].isna().astype("Int64")
    sr_missing = out.get("slider_response", pd.Series(pd.NA, index=df_raw.index)).isna()
    ms = pd.Series(pd.NA, index=df_raw.index)
    ms[block == 2] = sr_missing[block == 2].astype("Int64")
    out["missed_slider"] = ms

    for c in ["block", "run", "trials_inblock", "gamble_left", "gamble_choice"]:
        if c in out.columns:
            out[c] = pd.to_numeric(out[c], errors="coerce").astype("Int64")

    out = drop_metadata_rows(out, meta_keep=("pOrder",))
    out = cast_integer_like(out)
    return out

# ------------------------------
# main
# ------------------------------
def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--infile", required=True)
    parser.add_argument("--outfile", required=True)
    args = parser.parse_args()

    infile = Path(args.infile)
    outfile = Path(args.outfile)

    task = detect_task_from_filename(infile)

    try:
        if task in {"th", "tm"}:
            out = preprocess_trust(infile)
        elif task == "ol":
            out = preprocess_ol(infile)
        elif task == "socialRA":
            out = preprocess_social(infile)
        else:
            out = preprocess_trust(infile)

        outfile.parent.mkdir(parents=True, exist_ok=True)
        out.to_csv(outfile, index=False)
        log(f"✅ Saved preprocessed → {outfile}  (rows={len(out):,}, cols={out.shape[1]})")
    except Exception as e:
        log(f"❌ Preprocessing failed for {infile.name}: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()