src/modules/pubtime/common.py

from __future__ import annotations

import re
import zipfile
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import pandas as pd


IMPORTED_PROJECT_ROOT = Path("modules") / "pubtime"
ZIP_DATASET_ROOT = "Fina Trial Publication"
CENSOR_DATE = pd.Timestamp("2025-01-01")

DATE_COLUMNS = (
    "completion_date",
    "primary_completion_date",
    "study_first_submitted_date",
    "last_update_submitted_date",
    "study_first_posted_date",
)

BLANK_TO_MISSING_COLUMNS = (
    "has_dmc",
    "were_results_reported",
    "has_expanded_access",
    "has_us_facility",
    "has_single_facility",
    "healthy_volunteers",
    "adult",
    "child",
    "older_adult",
    "source_class",
    "phase",
    "gender",
    "primary_purpose",
    "intervention_model",
    "allocation",
    "masking",
)

RAW_PREDICTORS = (
    "has_dmc",
    "source_class",
    "phase",
    "actual_duration",
    "were_results_reported",
    "enrollment",
    "gender",
    "has_expanded_access",
    "number_of_facilities",
    "has_us_facility",
    "has_single_facility",
    "primary_purpose",
    "number_of_arms",
    "intervention_model",
    "allocation",
    "healthy_volunteers",
    "months_to_completion_date",
    "months_to_primary_completion_date",
    "months_to_study_first_submitted_date",
    "months_to_last_update_submitted_date",
    "months_to_study_first_posted_date",
    "number_of_primary_outcomes_to_measure",
    "number_of_secondary_outcomes_to_measure",
    "masking",
    "minimum_age_days",
    "adult",
    "child",
    "older_adult",
)

CONTINUOUS_VARS = (
    "actual_duration",
    "enrollment",
    "number_of_facilities",
    "number_of_arms",
    "months_to_completion_date",
    "months_to_study_first_submitted_date",
    "months_to_last_update_submitted_date",
    "months_to_study_first_posted_date",
    "number_of_primary_outcomes_to_measure",
    "number_of_secondary_outcomes_to_measure",
    "minimum_age_days",
)

DOMAIN_DATASETS = {
    "cancer": Path("Data/Cancer_dataset/cancer_data_w_pub_date.csv"),
    "covid": Path("Data/Covid_dataset/covid_data_w_pub_date.csv"),
    "cvd": Path("Data/CVD_dataset/cvd_data_w_pub_date.csv"),
}


@dataclass(frozen=True)
class PubTimePreparedData:
    domain: str
    raw_rows: int
    analytic_rows: int
    survival_rows: int
    predictors: tuple[str, ...]
    high_missing_predictors: tuple[str, ...]
    missing_percentages: dict[str, float]
    survival_df: pd.DataFrame
    scaled_survival_df: pd.DataFrame
    cox_model_status: dict[str, Any]


def resolve_source(project_root: Path) -> Path:
    module_root = project_root / IMPORTED_PROJECT_ROOT
    if module_root.exists():
        return module_root
    return project_root / "Fina Trial Publication.zip"


def source_label(source_path: Path, project_root: Path) -> str:
    try:
        return str(source_path.relative_to(project_root))
    except ValueError:
        return str(source_path)


def load_domain_csv(source_path: Path, domain: str) -> pd.DataFrame:
    if domain not in DOMAIN_DATASETS:
        return pd.DataFrame()

    dataset_path = DOMAIN_DATASETS[domain]
    if source_path.is_dir():
        return pd.read_csv(source_path / dataset_path, encoding="utf-8-sig", low_memory=False)

    with zipfile.ZipFile(source_path) as archive:
        zipped_path = f"{ZIP_DATASET_ROOT}/{dataset_path.as_posix()}"
        with archive.open(zipped_path) as raw_file:
            return pd.read_csv(raw_file, encoding="utf-8-sig", low_memory=False)


def prepare_domain_data(source_path: Path, domain: str) -> PubTimePreparedData:
    raw_df = load_domain_csv(source_path, domain)
    if raw_df.empty:
        return 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=_cox_unavailable(),
        )

    raw_df = _add_r_derived_columns(raw_df)
    predictors, high_missing, missing = _select_predictors(raw_df)
    analytic_df = raw_df.dropna(subset=list(predictors)).copy()
    analytic_df = _regroup_factors(analytic_df)
    survival_df = _add_time_to_publication(analytic_df)
    survival_df = survival_df[survival_df["registered_in_calendar_year"] >= 2010].copy()
    if domain == "covid" and not survival_df.empty:
        survival_df["year_group"] = (survival_df["registered_in_calendar_year"] > 2021).map(
            {False: 1, True: 2}
        )

    scaled = survival_df.copy()
    for column in CONTINUOUS_VARS:
        if column in scaled.columns:
            scaled[column] = scale_within_percentile(scaled[column], lower=1, upper=90)

    return PubTimePreparedData(
        domain=domain,
        raw_rows=len(raw_df),
        analytic_rows=len(analytic_df),
        survival_rows=len(survival_df),
        predictors=predictors,
        high_missing_predictors=high_missing,
        missing_percentages=missing,
        survival_df=survival_df,
        scaled_survival_df=scaled,
        cox_model_status=_cox_unavailable(),
    )


def profile_for_pubtime(profile: dict[str, Any]) -> dict[str, Any]:
    normalized = dict(profile)
    normalized["phase"] = regroup_phase(profile.get("phase"))
    normalized["primary_purpose"] = regroup_primary_purpose(profile.get("primary_purpose"))
    normalized["intervention_model"] = regroup_intervention_model(profile.get("intervention_model"))
    return normalized


def _add_r_derived_columns(df: pd.DataFrame) -> pd.DataFrame:
    result = df.copy()
    for column in ("start_date", *DATE_COLUMNS):
        if column in result.columns:
            result[column] = pd.to_datetime(result[column], errors="coerce")

    if "start_date" in result.columns:
        for column in DATE_COLUMNS:
            if column in result.columns:
                result[f"months_to_{column}"] = (result[column] - result["start_date"]).dt.days / 30.44

    if "pubmed_link" in result.columns:
        result["has_link"] = (result["pubmed_link"].fillna("No") != "No").astype(int)

    for column in BLANK_TO_MISSING_COLUMNS:
        if column in result.columns:
            result[column] = result[column].replace("", pd.NA)

    if "minimum_age" in result.columns:
        result["minimum_age_days"] = result["minimum_age"].map(convert_to_days)

    return result


def _select_predictors(df: pd.DataFrame) -> tuple[tuple[str, ...], tuple[str, ...], dict[str, float]]:
    available = [column for column in RAW_PREDICTORS if column in df.columns]
    missing = (df[available].isna().mean() * 100).to_dict()
    high_missing = tuple(column for column in available if missing[column] > 30)
    predictors = tuple(column for column in available if column not in high_missing)
    rounded_missing = {column: round(float(value), 3) for column, value in missing.items()}
    return predictors, high_missing, rounded_missing


def _regroup_factors(df: pd.DataFrame) -> pd.DataFrame:
    result = df.copy()
    if "source_class" in result.columns:
        result["source_class"] = result["source_class"].map(regroup_source_class)
    if "phase" in result.columns:
        result["phase"] = result["phase"].map(regroup_phase)
    if "primary_purpose" in result.columns:
        result["primary_purpose"] = result["primary_purpose"].map(regroup_primary_purpose)
    if "intervention_model" in result.columns:
        result["intervention_model"] = result["intervention_model"].map(regroup_intervention_model)
    return result


def _add_time_to_publication(df: pd.DataFrame) -> pd.DataFrame:
    survival = df.copy()
    survival["completion_date"] = pd.to_datetime(survival["completion_date"], errors="coerce")
    survival["pub_date"] = parse_publication_dates(survival.get("pub_date"))
    survival.loc[survival["pub_date"] < survival["completion_date"], "pub_date"] = pd.NaT
    survival["time_to_pub"] = (
        survival["pub_date"].fillna(CENSOR_DATE) - survival["completion_date"]
    ).dt.days
    survival["result_count"] = survival["pub_date"].notna().astype(int)

    first_rows = survival.groupby("nct_id", sort=False).head(1).reset_index(drop=True)
    min_time = survival.groupby("nct_id", sort=False)["time_to_pub"].min().reset_index()
    first_event = survival.groupby("nct_id", sort=False)["result_count"].first().reset_index()
    deduped = first_rows.drop(columns=["time_to_pub", "result_count"], errors="ignore")
    deduped = deduped.merge(min_time, on="nct_id", how="left")
    deduped = deduped.merge(first_event.rename(columns={"result_count": "event_pub"}), on="nct_id", how="left")
    return deduped


def convert_to_days(age: Any) -> float | None:
    if pd.isna(age):
        return None
    text = str(age)
    match = re.search(r"\d+", text)
    if not match:
        return None
    value = float(match.group(0))
    unit = re.sub(r"\d+\s*", "", text).strip().lower()
    if "minute" in unit:
        return value / 1440
    if "hour" in unit:
        return value / 24
    if "day" in unit:
        return value
    if "week" in unit:
        return value * 7
    if "month" in unit:
        return value * 30.44
    if "year" in unit:
        return value * 365.25
    return None


def parse_publication_dates(values: Any) -> pd.Series:
    if values is None:
        return pd.Series(dtype="datetime64[ns]")

    text = pd.Series(values).astype("string").str.strip().str.rstrip(".")
    text = text.replace({"": pd.NA, "No": pd.NA, "NA": pd.NA})
    parsed = pd.to_datetime(text, errors="coerce", format="mixed")

    year_month = text.str.extract(r"^(\d{4})[-\s]+([A-Za-z]{3,9}|\d{1,2})$").dropna(how="all")
    for index, row in year_month.iterrows():
        if pd.isna(parsed.loc[index]):
            parsed.loc[index] = pd.to_datetime(f"{row[0]} {row[1]} 01", errors="coerce")

    month_year = text.str.extract(r"^([A-Za-z]{3,9})\s+(\d{4})$").dropna(how="all")
    for index, row in month_year.iterrows():
        if pd.isna(parsed.loc[index]):
            parsed.loc[index] = pd.to_datetime(f"{row[0]} 01 {row[1]}", errors="coerce")

    # Month ranges such as "2021 Nov-Dec" or "2016 Nov/Dec": resolve to the first
    # month, day 1 (matching R's lubridate, which also keeps the leading month).
    # Calendar-season strings ("2009 Fall") are deliberately left unparsed: R either
    # drops them or mis-parses them to garbage, so NaT here is at least as correct.
    month_range = text.str.extract(
        r"^(\d{4})\s+([A-Za-z]{3,9})\s*[-/]\s*[A-Za-z]{3,9}$"
    ).dropna(how="all")
    for index, row in month_range.iterrows():
        if pd.isna(parsed.loc[index]):
            parsed.loc[index] = pd.to_datetime(f"{row[0]} {row[1]} 01", errors="coerce")

    return parsed


def regroup_source_class(value: Any) -> str | None:
    if pd.isna(value):
        return None
    value = str(value)
    if value in {"FED", "NIH", "OTHER_GOV"}:
        return "Government"
    if value in {"INDIV", "INDUSTRY", "NETWORK"}:
        return "Private"
    if value in {"OTHER", "UNKNOWN"}:
        return "Other"
    return None


def regroup_phase(value: Any) -> str | None:
    if pd.isna(value):
        return None
    value = str(value).upper()
    if value in {"EARLY_PHASE1", "PHASE1", "PHASE1/PHASE2"}:
        return "Early Phase"
    if value in {"PHASE2", "PHASE2/PHASE3"}:
        return "Phase 2"
    if value == "PHASE3":
        return "Phase 3"
    if value == "PHASE4":
        return "Phase 4"
    return None


def regroup_primary_purpose(value: Any) -> str | None:
    if pd.isna(value):
        return None
    value = str(value).upper()
    if value in {
        "TREATMENT",
        "SUPPORTIVE_CARE",
        "PREVENTION",
        "DIAGNOSTIC",
        "BASIC_SCIENCE",
        "HEALTH_SERVICES_RESEARCH",
    }:
        return value
    if value in {"OTHER", "SCREENING", "DEVICE_FEASIBILITY"}:
        return "OTHER"
    return None


def regroup_intervention_model(value: Any) -> str | None:
    if pd.isna(value):
        return None
    value = str(value).upper()
    if value == "PARALLEL":
        return "Parallel"
    if value == "CROSSOVER":
        return "Crossover"
    if value == "SINGLE_GROUP":
        return "Single Group"
    if value in {"SEQUENTIAL", "FACTORIAL"}:
        return "Other"
    return None


def scale_within_percentile(series: pd.Series, lower: int = 1, upper: int = 90) -> pd.Series:
    numeric = pd.to_numeric(series, errors="coerce")
    p_lower = numeric.quantile(lower / 100)
    p_upper = numeric.quantile(upper / 100)
    within = numeric.where((numeric >= p_lower) & (numeric <= p_upper))
    mean = within.mean()
    std = within.std()
    if pd.isna(std) or std == 0:
        return pd.Series(0, index=series.index, dtype="float64")
    scaled = (within - mean) / std
    minimum = scaled.min()
    maximum = scaled.max()
    scaled = scaled.mask(numeric < p_lower, minimum)
    scaled = scaled.mask(numeric > p_upper, maximum)
    return scaled


def _cox_unavailable() -> dict[str, Any]:
    return {
        "status": "unavailable",
        "reason": (
            "The original paper fits Cox proportional hazards models with R survival::coxph. "
            "This Python runtime does not add survival-model dependencies, so it exposes "
            "prepared analysis data but does not fabricate coefficients."
        ),
    }