#!/usr/bin/env python3
"""Apply the preregistered inclusion filter to the crawled NeuroVault metadata and freeze the list.
Input: crawl/collections/*.json.gz and crawl/images/*.json.gz (raw API pages).
Output: frozen/included_maps.csv (+ exclusion counts). Deterministic: rows sorted by image_id."""
import csv, glob, gzip, json, os, re, sys, hashlib
from collections import Counter
ROOT = sys.argv[1] if len(sys.argv) > 1 else os.path.join(os.path.dirname(__file__), "..")
CRAWL, OUT = os.path.join(ROOT, "crawl"), os.path.join(ROOT, "frozen")

# Between-group / correlational contrasts, judged from image text fields only (case-insensitive).
GROUP_RE = re.compile(r"patients?\b|\bcontrols\b|healthy (controls?|volunteers? vs)|\bHCs?\b|\bTD\b|typically[- ]developing|"
    r"between[- ]group|group[- ]?(difference|comparison|effect|contrast|interaction)|\bgroup\s*[x×*]|[x×*]\s*group\b|"
    r"(diagnos|clinical|disorder\b|\bADHD|schizo|psychos|depress|\bMDD\b|autis|\bASD\b|PTSD|bipolar|alzheimer|parkinson|dementia|\bMCI\b|"
    r"\bOCD\b|anorexi|bulimi|addict|\bAUD\b|\bSUD\b|dyslexi|epilep|stroke)|"
    r"\bolder\b.*\byounger\b|\byounger\b.*\bolder\b|\b(adults?|adolescents?|children)\s*(vs\.?|versus|>|<|-)\s*(adults?|adolescents?|children)\b|"
    r"\b(males?|men)\s*(vs\.?|versus|>|<|-)\s*(females?|women)\b|\b(females?|women)\s*(vs\.?|versus|>|<|-)\s*(males?|men)\b|sex differences?|gender differences?|"
    r"two[- ]sample|independent[- ]samples?|unpaired|\bgroups?\s*(vs\.?|versus|>|<)|(vs\.?|versus|>|<)\s*(\w+\s+)?groups?\b|congenital", re.I)
CORR_RE = re.compile(r"\bBMI\b|correlat|\bcorr\b|^\s*regression\b|regression:|(positive|negative)? ?relation (with|to|between)|(positive|negative) relation|covariat|regression (on|with|of)|brain[- ]behaviou?r|individual differences|"
    r"associat(ed|ion)s? (with|between)|predict(s|ed|ion) (of|by)|\bmodulated by (age|score|trait)|\b(age|IQ|BMI|score|trait|symptom)s? effect", re.I)

# Not a task-activation contrast map (resting-state / connectivity / ICA / F-test), judged from metadata text.
NONTASK_RE = re.compile(r"\bICA\b|independent component|\bcomponent|\bIC_?\d|melodic|dual[_ ]?regression|\bseed|connectivity|\bFC\b|"
    r"\bPPI\b|resting|rs-?fMRI|rsfc|\bf?ALFF\b|ReHo|\bDC:|degree centrality|\bz?fstat|\bF[- ]?(test|stat|map|contrast)\b|"
    r"multivariate|connectedness|searchlight|decod|\bMVPA\b|classifi|\bRSA\b|representational similarity|\bsub(ject)?[ -]?\d+\b", re.I)

def text_of(im):  # underscores/dots -> spaces so word boundaries work on file-style names
    return re.sub(r"[_.]+", " ", _text_of(im))

def _text_of(im):
    return " | ".join(str(im.get(k) or "") for k in ("name", "description", "contrast_definition", "cognitive_contrast_cogatlas")) + " | " + os.path.basename(im.get("file") or "")

def fnum(x):
    try:
        v = float(x)
        return v if v == v else None
    except (TypeError, ValueError):
        return None

def main():
    cols = {}
    for f in sorted(glob.glob(f"{CRAWL}/collections/p*.json.gz")):
        for c in json.load(gzip.open(f))["results"]:
            cols[c["id"]] = c
    images = {}
    for f in sorted(glob.glob(f"{CRAWL}/images/c*_p*.json.gz")):
        for im in json.load(gzip.open(f))["results"]:
            images[im["id"]] = im
    ex = Counter(); rows = []
    for iid in sorted(images):
        im = images[iid]; c = cols.get(im["collection_id"], {})
        def drop(r): ex[r] += 1
        if not (c.get("DOI") or "").strip(): drop("01_no_publication_DOI"); continue
        if im.get("image_type") != "statistic_map": drop("02_not_statistic_map"); continue
        if im.get("modality") != "fMRI-BOLD": drop("03_not_fMRI_BOLD"); continue
        if im.get("analysis_level") != "group": drop("04_not_group_level"); continue
        if im.get("map_type") not in ("T map", "Z map"): drop("05_not_T_or_Z"); continue
        if im.get("is_thresholded") is not False: drop("06_thresholded_or_unknown"); continue
        if im.get("not_mni") is not False: drop("07_not_MNI_or_unknown"); continue
        if im.get("data_origin") not in (None, "volume"): drop("08_not_volume"); continue
        if (im.get("subject_species") or "homo sapiens").lower() != "homo sapiens": drop("09_not_human"); continue
        f = im.get("file") or ""
        if not (f.endswith(".nii.gz") or f.endswith(".nii")): drop("10_not_nifti_file"); continue
        n = fnum(im.get("number_of_subjects"))
        if n is None: drop("11_N_missing"); continue
        if n < 10 or n != int(n): drop("12_N_lt_10_or_noninteger"); continue
        if c.get("group_comparison") is True: drop("13_collection_declares_group_comparison"); continue
        t = text_of(im)
        if GROUP_RE.search(t): drop("14_between_group_text"); continue
        if CORR_RE.search(t): drop("15_correlation_text"); continue
        if NONTASK_RE.search(t) or re.search(r"\brest", im.get("cognitive_paradigm_cogatlas") or "", re.I):
            drop("16_not_task_activation_map"); continue
        n = int(n); mt = "T" if im["map_type"] == "T map" else "Z"
        df = fnum(im.get("statistic_parameters"))
        if mt == "T" and df is not None and (df > n - 1 + 0.5 or df < 0.5 * (n - 1)):
            drop("17_T_df_inconsistent_with_N"); continue
        fwhm = fnum(im.get("smoothness_fwhm"))
        if fwhm is None: fwhm = fnum(c.get("smoothing_fwhm"))
        if fwhm is not None and not (0 <= fwhm <= 30): fwhm = None
        task = (im.get("cognitive_paradigm_cogatlas") or "").strip()
        if task in ("", "None", "None / Other"): task = "unknown"
        rows.append({"image_id": iid, "collection_id": im["collection_id"], "file_url": f.replace("http://", "https://", 1),
                     "N": n, "map_type": mt, "df_meta": "" if df is None else f"{df:g}",
                     "fwhm_mm": "" if fwhm is None else f"{fwhm:g}", "upload_year": (im.get("add_date") or "")[:4],
                     "task": task, "task_id": im.get("cognitive_paradigm_cogatlas_id") or "",
                     "doi": c["DOI"].strip(), "name": re.sub(r"[\r\n\t]+", " ", im.get("name") or "")[:200],
                     "contrast_definition": re.sub(r"[\r\n\t]+", " ", im.get("contrast_definition") or "")[:200]})
    os.makedirs(OUT, exist_ok=True)
    path = f"{OUT}/included_maps.csv"
    with open(path, "w", newline="") as fh:
        w = csv.DictWriter(fh, fieldnames=list(rows[0].keys()), lineterminator="\n"); w.writeheader(); w.writerows(rows)
    sha = hashlib.sha256(open(path, "rb").read()).hexdigest()
    summ = {"n_collections_crawled": len(cols), "n_images_crawled_in_DOI_collections": len(images),
            "exclusions_in_order": dict(sorted(ex.items())), "n_included_maps": len(rows),
            "n_included_collections": len({r["collection_id"] for r in rows}), "included_maps_csv_sha256": sha}
    json.dump(summ, open(f"{OUT}/freeze_summary.json", "w"), indent=1, sort_keys=True)
    print(json.dumps(summ, indent=1))

if __name__ == "__main__":
    main()