src/modules/pubtime/adapter.py

from __future__ import annotations

from pathlib import Path

import pandas as pd

from .common import PubTimePreparedData, resolve_source, source_label
from . import cancer_analysis, covid_analysis, cvd_analysis


PREPARERS = {
    "cancer": cancer_analysis.prepare,
    "covid": covid_analysis.prepare,
    "cvd": cvd_analysis.prepare,
}


def prepare_from_archive(domain: str, project_root: Path) -> tuple[PubTimePreparedData, str, str]:
    source_path = resolve_source(project_root)
    return prepare_from_source(domain, source_path, source_label(source_path, project_root))


def prepare_from_source(domain: str, source_path: Path, label: str | None = None) -> tuple[PubTimePreparedData, str, str]:
    preparer = PREPARERS.get(domain)
    if preparer is None:
        empty = PubTimePreparedData(
            domain=domain,
            raw_rows=0,
            analytic_rows=0,
            survival_rows=0,
            predictors=tuple(),
            high_missing_predictors=tuple(),
            missing_percentages={},
            survival_df=pd.DataFrame(),
            scaled_survival_df=pd.DataFrame(),
            cox_model_status={"status": "unavailable", "reason": "Unknown domain."},
        )
        return empty, label or str(source_path), "directory" if source_path.is_dir() else "zip"

    return preparer(source_path), label or str(source_path), "directory" if source_path.is_dir() else "zip"