script.py

import argparse
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
import gc
import json
import resource
import shutil
import sys
import morph_kgc
import mwtab
import numpy as np
import pandas as pd
import os
from pathlib import Path
from tqdm import tqdm
from utils.config import bind_graph
import xml.etree.ElementTree as ET


def _raise_fd_limit() -> None:
    soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
    target = hard if hard != resource.RLIM_INFINITY else 65536
    if soft < target:
        try:
            resource.setrlimit(resource.RLIMIT_NOFILE, (target, hard))
        except (ValueError, OSError):
            pass


_raise_fd_limit()


_PBAR: tqdm | None = None
_CURRENT_STAGE: str | None = None


def _set_stage(file_id: str | None, stage: str | None) -> None:
    global _CURRENT_STAGE
    _CURRENT_STAGE = stage
    if _PBAR is None:
        return
    if file_id is None or stage is None:
        _PBAR.set_postfix_str("idle", refresh=True)
    else:
        _PBAR.set_postfix_str(f"{file_id}[{stage}]", refresh=True)

MAPPING_FOLDER = "mapping/rml/"

LOGS_LEVEL = "WARNING"

GNPS_JOBS_DIR = Path("data_retriever/gnps_jobs")
MTBLS_DIR = Path("data_retriever/mtbls_study")
MW_DIR = Path("data_retriever/mw_study")

JOB_LINKS = pd.read_csv("data_retriever/job_massive_llm.csv")


@dataclass
class Metadata:
    repo: str
    origin: str
    file: Path
    id: str = field(init=False)

    def __post_init__(self):
        if self.origin == "remote":
            self.id = f"{self.file.parent.parent.name}_{self.file.parent.name}"
        else:
            self.id = f"{self.file.stem}"

    def __repr__(self) -> str:
        return f"repo={self.repo} origin={self.origin} id={self.id}"


def build_sources_df(
    id_meta: Metadata,
    sources: pd.DataFrame,
    feature_table: str | None = None,
) -> pd.DataFrame:
    """Return a sources table with one row per (collectionID, mzml).

    For each unique path in ``mzml_paths``:
      - if the path starts with ``MSV.../`` that prefix is the collectionID;
      - otherwise fall back to the JOB_LINKS ``massive_ids`` for this task,
        exploded on ``' | '``, producing one row per collectionID.
    """
    task_id = id_meta.file.parent.name
    job_row = JOB_LINKS[JOB_LINKS["task_id"] == task_id]
    ids: list[str] = []
    if not job_row.empty:
        raw = str(job_row["massive_ids"].iloc[0])
        ids = [c.strip() for c in raw.split("|") if c.strip()]
    else:
        ids = [sources["mzml"].iloc[0].split("/", 1)[0]]

    sources = sources.merge(pd.DataFrame({"collectionID": ids}), how="cross")

    sources["mzml"] = sources["mzml"].str.rsplit("/", n=1).str[-1]
    
    sources["task_id"] = task_id
    sources["feature_table"] = feature_table

    base = "https://gnps.ucsd.edu/ProteoSAFe"
    default_view = "view_all_annotations_DB"
    sources["download_url"] = (
        f"{base}/DownloadResult?task={task_id}&view={default_view}"
    )

    required = [
        "collectionID",
        "mzml",
        "task_id",
        "feature_table",
        "download_url",
        "annotation_file",
        "NumberHits",
        "id",
    ]
    missing = set(required) - set(sources.columns)
    if missing:
        raise ValueError(f"Missing columns: {missing}")

    return sources


def pan_redu_to_dfs(
    redu_path: str, id_meta: Metadata
) -> dict[str, pd.DataFrame | dict]:
    tqdm.write(f"Parsing Pan-ReDu file: {redu_path}")

    df = pd.read_csv(redu_path, sep="\t", na_values="missing value")

    df = df[df["ATTRIBUTE_DatasetAccession"].str.startswith("MSV")]
    df['filename'] = df['filename'].str.split("/").str[-1]

    df["MassSpectrometer"] = (
        df["MassSpectrometer"].astype(str).str.split("|").str[1].fillna("MS:1000463")
    )
    df["Polarity"] = (
        df["IonizationSourceAndPolarity"].astype(str).str.extract(r"\((.*?)\)")
    )
    df["IonizationSource"] = (
        df["IonizationSourceAndPolarity"]
        .astype(str)
        .str.replace(r"\s*\(.*?\)", "", regex=True)
        .fillna("unknown")
    )

    nodes_values = {
        "iri_source": [
            "NCBITaxonomy",
            "Country",
            "ENVOEnvironmentBiomeIndex",
            "ENVOEnvironmentMaterialIndex",
            "BiologicalSex",
            "LifeStage",
            "HealthStatus",
            "AgeInYears",
        ]
    }

    for key, value in nodes_values.items():
        df[key] = df[value].apply(lambda x: "_".join(x.dropna().astype(str)), axis=1)

    df["mzml"] = df["filename"].str.split("/").str[-1]
    df["YearOfAnalysis"] = df["YearOfAnalysis"].astype("Int64")

    return {"pan_redu": df}


GNPS_OPTIONAL_COLUMNS = [
    "superclass",
    "class",
    "subclass",
    "npclassifier_superclass",
    "npclassifier_class",
    "npclassifier_pathway",
    "MQScore",
    "MZErrorPPM",
    "SharedPeaks",
    "MassDiff",
    "Compound_Name",
    "Pubmed_ID",
    "IonMode",
    "Ion_Source",
    "Instrument",
    "Adduct",
    "Precursor_MZ",
    "ExactMass",
    "Charge",
    "CAS_Number",
    "Smiles",
    "INCHI",
    "LibraryQualityString",
    "Organism",
    "InChIKey",
]

# Headers in FBMN quantification CSVs that are NOT per-sample abundance.
# Anything else (and not an "Unnamed: N" artefact) is treated as a sample column.
GNPS_NON_ABUNDANCE_COLS = {
    "row ID", "row m/z", "row retention time",
    "correlation group ID", "annotation network number",
    "auto MS2 verify", "best ion", "identified by n=",
    "neutral M mass", "partners",
    "row CCS", "row ion mobility", "row ion mobility unit",
    "Compound", "Adducts", "Formula", "Score", "Fragmentation Score",
    "Mass Error (ppm)", "Isotope Similarity", "Retention Time (min)",
    "Chromatographic peak width (min)", "Identifications",
    "Anova (p)", "q Value", "Max Fold Change",
    "Highest Mean", "Lowest Mean", "Maximum Abundance", "Minimum CV%",
    "Accepted Compound ID", "Compound Link", "Isotope Distribution",
    "Molecular Formula", "Neutral mass", "Retention Time Error (mins)",
    "Low abundant", "Charge",
}


def _first_file(folder: Path, suffixes: set[str]) -> Path | None:
    if not folder.exists():
        return None
    return next((f for f in folder.iterdir() if f.suffix in suffixes), None)


def _write_full_gnps_mapping(merge_mapping: str, extra_template: str) -> str:
    with open(f"{MAPPING_FOLDER}{extra_template}") as f:
        merge_mapping += f.read()
    with open(f"{MAPPING_FOLDER}FULL_GNPS.ttl", "w") as f:
        f.write(merge_mapping)
    return "FULL_GNPS.ttl"


def _build_sources(
    base: pd.DataFrame,
    df: pd.DataFrame,
    file: str,
    id_meta: Metadata,
    rename: dict[str, str],
    feature_table: str | None = None,
) -> pd.DataFrame:
    sources_df = base.copy()
    sources_df["annotation_file"] = file
    sources_df.rename(columns=rename, inplace=True)
    sources_df = sources_df.merge(
        df[["#Scan#", "NumberHits"]],
        left_on="id",
        right_on="#Scan#",
        how="inner",
    ).drop(columns=["#Scan#"])
    return build_sources_df(id_meta, sources_df, feature_table)


def _annotate_gnps_df(df: pd.DataFrame, task_id: str) -> None:
    for col in GNPS_OPTIONAL_COLUMNS:
        if col not in df.columns:
            df[col] = None
    
    if 'NumberHits' not in df.columns:
        df['NumberHits'] = 1

    if not df["Ion_Source"].isnull().all():
        df["Chromatography"] = np.where(
            df["Ion_Source"].str.contains("-"),
            df["Ion_Source"].str.split("-").str[0],
            "",
        )
        df["Ion_Source"] = df["Ion_Source"].str.split("-").str[-1]

    df["task_id"] = task_id


def gnps_to_dfs(
    file: str, id_meta: Metadata
) -> tuple[dict[str, pd.DataFrame | dict], str]:
    with open(f"{MAPPING_FOLDER}GNPS.ttl") as f:
        merge_mapping = f.read() + "\n"

    df = pd.read_csv(file, sep="\t", na_values=["", " ", "N/A"])

    if "#Scan#" not in df.columns:
        raise ValueError(f"#Scan# column is required for {id_meta}")
    
    task_id = id_meta.file.parent.name
    _annotate_gnps_df(df, task_id)

    if id_meta.origin != "remote":
        with open(f"{MAPPING_FOLDER}FULL_GNPS.ttl", "w") as f:
            f.write(merge_mapping)
        return {"gnps_data": df}, "FULL_GNPS.ttl"

    is_mn = "METABOLOMICS-SNETS" in str(id_meta.file)

    job_folder = Path(file).parent
    summary_folder = job_folder / (
        "clusterinfosummarygroup_attributes_withIDs_withcomponentID"
        if is_mn
        else "clusterinfo_summary"
    )

    def _read_summary() -> pd.DataFrame:
        summary_path = _first_file(summary_folder, {".clustersummary", ".tsv"})
        if summary_path is None:
            raise FileNotFoundError(
                f"No summary file found in {summary_folder} for {id_meta}"
            )
        return pd.read_csv(summary_path, sep="\t", na_values=["", " ", "N/A"])

    def _read_params() -> ET.Element:
        return ET.parse(job_folder / "params.xml").getroot()

    def _read_secondary() -> pd.DataFrame:
        if is_mn:
            return pd.read_csv(
                next((job_folder / "clusterinfo").glob("*.clusterinfo")), sep="\t"
            )
        
        quantification_path = _first_file(
            job_folder / "quantification_table_reformatted", {".csv"}
        )
        if quantification_path is None:
            return pd.DataFrame()
        return pd.read_csv(quantification_path)

    with ThreadPoolExecutor(max_workers=3) as ex:
        summary_future = ex.submit(_read_summary)
        params_future = ex.submit(_read_params)
        secondary_future = ex.submit(_read_secondary)
        summary_df = summary_future.result()
        params_root = params_future.result()
        secondary_df = secondary_future.result()

    summary_df["task_id"] = task_id
    summary_df["cluster index"] = summary_df["cluster index"].astype(int)
    df["#Scan#"] = df["#Scan#"].astype(int)

    if is_mn:
        scan_filter = df[["#Scan#"]]
        summary_df = summary_df.merge(
            scan_filter, left_on="cluster index", right_on="#Scan#", how="inner"
        ).drop(columns=["#Scan#"])

        cluster_df = secondary_df.merge(
            scan_filter, left_on="#ClusterIdx", right_on="#Scan#", how="inner"
        ).drop(columns=["#Scan#"])
        cluster_df["task_id"] = task_id
        cluster_df["#ClusterIdx"] = cluster_df["#ClusterIdx"].astype(int)

        file_map = {}
        for mapping in params_root.findall(
            ".//parameter[@name='upload_file_mapping']"
        ):
            local, real = mapping.text.split("|", 1)
            file_map[local] = real
        cluster_df["#Filename"] = cluster_df["#Filename"].apply(
            lambda x: file_map[x.replace("inputspectra/", "")]
        )

        sources_df = _build_sources(
            cluster_df[["#ClusterIdx", "#Filename"]],
            df,
            file,
            id_meta,
            rename={"#Filename": "mzml", "#ClusterIdx": "id"},
        )

        template = _write_full_gnps_mapping(merge_mapping, "MN_GPNS.ttl")
        return {
            "gnps_data": df,
            "cluster_info": cluster_df,
            "cluster_summary": summary_df,
            "sources": sources_df,
        }, template

    quant_el = params_root.find(".//parameter[@name='QUANT_TABLE_SOURCE']")
    software_name = quant_el.text if quant_el is not None else None
    summary_df["software"] = software_name

    if 'SumPeakIntensity' not in summary_df.columns:
        summary_df['SumPeakIntensity'] = summary_df['sum(precursor intensity)']

    ft = params_root.find(".//parameter[@name='quantification_table']")
    feature_table = ft.text[2:] if ft is not None and ft.text is not None else None

    abundance_df = secondary_df
    if abundance_df.empty:
        sources_df = df[["#Scan#"]].copy()
        sources_df['mzml'] = None
        sources_df = _build_sources(
            sources_df,
            df,
            file,
            id_meta,
            rename={"#Scan#": "id"},
            feature_table=feature_table,
        )
        abundance_df = pd.DataFrame(columns=["row ID", "row m/z", "row retention time", "mzml","abundance"])
        
    else:
        cols = abundance_df.columns.to_series()
        keep_mask = (
            ~cols.isin(GNPS_NON_ABUNDANCE_COLS)
            & ~cols.str.match(r'^Unnamed:\s*\d+$')
        )
        abundance_cols = abundance_df.columns[keep_mask].tolist()
        id_cols = ['row ID', 'row m/z', 'row retention time']
        if len(abundance_cols) == 0:
            raise ValueError(f"No abundance columns found in {id_meta}")

        abundance_df = abundance_df[[c for c in id_cols if c in abundance_df.columns] + abundance_cols]
        abundance_df.columns = abundance_df.columns.str.replace(
            r'\s*Peak (?:area|height)\s*$', '', regex=True
        ).str.strip()
        abundance_df = abundance_df.melt(
            id_vars=["row ID", "row m/z", "row retention time"],
            var_name="mzml",
            value_name="abundance",
        )
        abundance_df["abundance"] = pd.to_numeric(abundance_df["abundance"], errors="raise")
        abundance_df = abundance_df[abundance_df["abundance"] > 0]

        sources_df = _build_sources(
            abundance_df[["row ID", "mzml"]],
            df,
            file,
            id_meta,
            rename={"row ID": "id"},
            feature_table=feature_table,
        )

        reduce = True
        if reduce:
            abundance_df["mzml"] = None
            abundance_df["abundance"] = None

    abundance_df["task_id"] = task_id
    abundance_df["software"] = software_name
    template = _write_full_gnps_mapping(merge_mapping, "FBMN_GNPS.ttl")
    return {
        "gnps_data": df,
        "fbmn_cluster_summary": summary_df,
        "fbmn_abundance": abundance_df,
        "sources": sources_df,
    }, template


def maf_to_dfs(maf_file: str, id_meta: Metadata) -> dict[str, pd.DataFrame | dict]:
    df = pd.read_csv(maf_file, sep="\t")
    df["SML_ID"] = pd.Series(range(1, len(df) + 1))
    df["annotation_file"] = str(maf_file)

    return {"mtbls_data": df}


def mw_to_dfs(mw_file: str, id_meta: Metadata) -> dict[str, pd.DataFrame | dict]:
    mwtabfile = next(mwtab.read_files([str(mw_file)]))

    header = mwtabfile.get("METABOLOMICS WORKBENCH", {})
    study_id = header.get("STUDY_ID", "")
    analysis_id = header.get("ANALYSIS_ID", "")

    ms_meta = mwtabfile.get("MS", {})
    instrument = ms_meta.get("INSTRUMENT_NAME", "")
    ion_mode = ms_meta.get("ION_MODE", "")

    metabolites = mwtabfile.get_table_as_pandas("Metabolites")
    metabolites["study_id"] = study_id
    metabolites["analysis_id"] = analysis_id
    metabolites["instrument"] = instrument
    metabolites["ion_mode"] = ion_mode
    metabolites["annotation_file"] = str(mw_file)

    data = mwtabfile.get_table_as_pandas("Data")
    sample_cols = [c for c in data.columns if c != "Metabolite"]
    data_long = data.melt(
        id_vars=["Metabolite"],
        value_vars=sample_cols,
        var_name="sample",
        value_name="abundance",
    )
    data_long["abundance"] = pd.to_numeric(data_long["abundance"], errors="coerce")
    data_long.dropna(subset=["abundance"], inplace=True)
    data_long["study_id"] = study_id
    data_long["analysis_id"] = analysis_id

    return {"mw_metabolites": metabolites, "mw_data": data_long}


def mapping_create_and_run(
    kg_path: Path,
    template: str,
    dfs: dict[str, pd.DataFrame | dict],
    file_id: str | None = None,
):
    if kg_path.exists():
        os.remove(kg_path)

    folder_path = Path("mapping/tmp/")
    if folder_path.exists():
        shutil.rmtree(folder_path)
    folder_path.mkdir()
    for var, df in dfs.items():
        if isinstance(df, dict):
            csv_path = f"{folder_path}/{var}.json"
            with open(csv_path, "w") as fp:
                json.dump(df, fp, indent=2)
        else:
            csv_path = f"{folder_path}/{var}.csv"
            df.to_csv(csv_path, index=False)

    del dfs

    config = f"""
    [CONFIGURATION]
    # INPUT
    na_values=#N/A,N/A,#N/A N/A,n/a,NA,<NA>,#NA,NULL,null,NaN,nan,None,nul,missing value,not specified,,
    # number_of_processes=8

    # LOGS
    logging_level={LOGS_LEVEL}
    logs_file=

    udfs=utils/udf.py

    [DataSource]
    mappings: {MAPPING_FOLDER}{template}
    """
    base = "https://ns.inria.fr/metaboKG/data/"

    _set_stage(file_id, "materializing")
    final_kg = morph_kgc.materialize(config)

    final_kg = bind_graph(final_kg)

    _set_stage(file_id, "serializing")
    final_kg.serialize(kg_path, format="turtle", base=base)


def router(file, id_meta: Metadata):
    repo = id_meta.repo
    if repo == "GNPS":
        return gnps_to_dfs(file, id_meta)
    elif repo == "REDU":
        return pan_redu_to_dfs(file, id_meta), "PanRedu.ttl"
    elif repo == "MW":
        return mw_to_dfs(file, id_meta), "MW.ttl"
    elif repo == "MTBLS":
        return maf_to_dfs(file, id_meta), "MTBLS.ttl"
    else:
        tqdm.write(f"Unsupported file format: {file}")
        sys.exit(1)


def process_file(
    file: Path,
    id_meta: Metadata,
    materialize: bool = True,
    save: bool = True,
    rewrite: bool = False,
) -> None:
    kg_name = f"{id_meta.id}.ttl"
    kg_path = Path(f"mapping/kg/{kg_name}")

    if (
        "38123557" in str(id_meta.file.parent.parent)
    ):  # Open access repository-scale propagated nearest neighbor suspect spectral library for untargeted metabolomics
        kg_name = f"{id_meta.id}_GNPS.ttl"
    if kg_path.exists() and not rewrite:
        return
    if "METABOLOMICS-SNETS-MZMINE" in str(id_meta.file):
        return

    _set_stage(id_meta.id, "preprocessing")
    dfs, template = router(file, id_meta)
    if save:
        _set_stage(id_meta.id, "saving")
        folder_path = f"mapping/csv_output/{id_meta.repo}/{id_meta.id}"
        os.makedirs(folder_path, exist_ok=True)
        for var, df in dfs.items():
            if isinstance(df, dict):
                csv_path = f"{folder_path}/{var}.json"
                with open(csv_path, "w") as fp:
                    json.dump(df, fp, indent=2)
            else:
                csv_path = f"{folder_path}/{var}.csv"
                df.to_csv(csv_path, index=False)
    if materialize:
        mapping_create_and_run(kg_path, template, dfs, file_id=id_meta.id)


def iterate_local_sources(
    materialize: bool = True,
    save: bool = True,
    rewrite: bool = False,
    origin: str = "remote",
):
    processed_files: list[str] = []

    repositories = [
        # {"repo": "GNPS", "root": GNPS_JOBS_DIR, "pattern": "*/*/*.tsv"},
        {"repo": "GNPS","root": GNPS_JOBS_DIR,"pattern": "*/*/METABOLOMICS-SNETS*.tsv",},
        {"repo": "GNPS", "root": GNPS_JOBS_DIR, "pattern": "*/*/FEATURE-BASED*.tsv"},
        {"repo": "MTBLS", "root": MTBLS_DIR, "pattern": "*/*maf.tsv"},
        {"repo": "MW", "root": MW_DIR, "pattern": "*/*.json"},
    ]

    linked = JOB_LINKS[JOB_LINKS["massive_ids"].astype(str).str.strip().ne("")]
    linked = linked[linked["massive_ids"].notna()]
    known_task_ids = set(linked["task_id"].astype(str))

    tasks: list[tuple[str, Path]] = []
    for repo_info in repositories:
        repo_name = repo_info["repo"]
        if repo_name != "GNPS":
            continue
        for file_path in repo_info["root"].glob(repo_info["pattern"]):
            if not file_path.is_file() or file_path.stat().st_size == 0:
                continue
            if file_path.parent.name not in known_task_ids:
                continue
            tasks.append((repo_name, file_path))

    global _PBAR
    pbar = tqdm(
        tasks,
        total=len(tasks),
        desc="GNPS jobs",
        unit="job",
        dynamic_ncols=True,
        colour="green",
        smoothing=0.1,
        bar_format=(
            "{desc}: {percentage:3.0f}%|{bar}| "
            "{n_fmt}/{total_fmt} "
            "[{elapsed}<{remaining}, {rate_fmt}] {postfix}"
        ),
    )
    _PBAR = pbar
    try:
        for repo_name, file_path in pbar:
            id_meta = Metadata(repo_name, origin, file_path)
            process_file(
                file_path,
                id_meta,
                materialize=materialize,
                save=save,
                rewrite=rewrite,
            )
            _set_stage(None, None)
            gc.collect()
            processed_files.append(str(file_path))
    finally:
        pbar.close()
        _PBAR = None

    return processed_files


def metadata_source(rewrite: bool = True):
    metadata_path = Path("test/data/redu_metadata.tsv")
    id_meta = Metadata("REDU", "local", metadata_path)
    process_file(metadata_path, id_meta, rewrite=rewrite)


SOURCE_CHOICES = ("metadata", "local", "both")


def main() -> int:
    parser = argparse.ArgumentParser(
        prog="mapping.script",
        description="Generate TTL knowledge graphs from local and/or metadata sources.",
    )
    parser.add_argument(
        "--source",
        choices=SOURCE_CHOICES,
        default="local",
        help="Which source(s) to map: 'metadata' (ReDU metadata), "
             "'local' (GNPS jobs), or 'both'. Default: local.",
    )
    parser.add_argument(
        "--rewrite",
        action="store_true",
        help="Overwrite existing KG files instead of skipping them.",
    )
    args = parser.parse_args()

    if args.source in ("metadata", "both"):
        metadata_source(rewrite=args.rewrite)
    if args.source in ("local", "both"):
        iterate_local_sources(rewrite=args.rewrite)
    return 0


if __name__ == "__main__":
    sys.exit(main())