src/trial_profile.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Any


DOMAIN_LABELS = {
    "cancer": "Cancer",
    "covid": "COVID",
    "cvd": "Cardiovascular",
}

ENUM_FIELDS = {
    "study_type",
    "overall_status",
    "source_class",
    "responsible_party_type",
    "intervention_type",
    "phase",
    "primary_purpose",
    "intervention_model",
    "allocation",
    "masking",
    "gender",
    "enrollment_type",
    "ipd_sharing_plan",
}

NUMERIC_FIELDS = {
    "enrollment": int,
    "number_of_arms": int,
    "number_of_facilities": int,
    "number_of_primary_outcomes": int,
    "number_of_secondary_outcomes": int,
    "minimum_age_years": float,
    "maximum_age_years": float,
}

BOOLEAN_FIELDS = {
    "has_dmc",
    "has_us_facility",
    "healthy_volunteers",
}

REQUIRED_FIELDS = {
    "domain",
    "brief_title",
    "official_title",
    "brief_summary",
    "conditions",
    "sponsor_name",
    "study_type",
    "phase",
    "enrollment",
    "number_of_arms",
    "number_of_facilities",
    "primary_purpose",
    "intervention_model",
    "allocation",
    "masking",
    "gender",
    "criteria",
}

DATE_FIELDS = {
    "start_date",
    "primary_completion_date",
    "completion_date",
}


@dataclass(frozen=True)
class ValidationResult:
    errors: list[str]
    warnings: list[str]


def normalize_trial(raw: dict[str, Any]) -> dict[str, Any]:
    profile: dict[str, Any] = {}

    for key, value in raw.items():
        if isinstance(value, str):
            value = value.strip()
        profile[key] = value

    profile["domain"] = str(profile.get("domain", "")).strip().lower()

    for key in ENUM_FIELDS:
        value = profile.get(key)
        if isinstance(value, str):
            profile[key] = value.strip().upper()

    for key, caster in NUMERIC_FIELDS.items():
        profile[key] = _parse_number(profile.get(key), caster)

    for key in BOOLEAN_FIELDS:
        profile[key] = _parse_bool(profile.get(key))

    profile["domain_label"] = DOMAIN_LABELS.get(profile.get("domain"), "Unknown")
    profile["protocol_sections"] = protocol_section_status(profile)
    profile["validation"] = validate_profile(profile).__dict__
    return profile


def validate_profile(profile: dict[str, Any]) -> ValidationResult:
    errors: list[str] = []
    warnings: list[str] = []

    for field in sorted(REQUIRED_FIELDS):
        value = profile.get(field)
        if value is None or value == "":
            errors.append(f"Missing required field: {field}")

    if profile.get("domain") not in DOMAIN_LABELS:
        errors.append("Domain must be cancer, covid, or cvd.")

    enrollment = profile.get("enrollment")
    if enrollment is not None and enrollment <= 0:
        errors.append("Enrollment must be greater than zero.")

    arms = profile.get("number_of_arms")
    if arms is not None and arms <= 0:
        errors.append("Number of arms must be greater than zero.")

    facilities = profile.get("number_of_facilities")
    if facilities is not None and facilities <= 0:
        errors.append("Number of facilities must be greater than zero.")

    criteria = str(profile.get("criteria", "") or "")
    if criteria and len(criteria.split()) < 20:
        warnings.append("Eligibility criteria are short; downstream extraction may be weak.")

    for date_field in DATE_FIELDS:
        value = profile.get(date_field)
        if value and not _looks_like_iso_date(str(value)):
            warnings.append(f"{date_field} should use YYYY-MM-DD format.")

    if profile.get("has_dmc") is False and enrollment and enrollment >= 500:
        warnings.append("Large planned enrollment without a DMC should be reviewed.")

    if profile.get("masking") == "NONE" and profile.get("allocation") == "RANDOMIZED":
        warnings.append("Randomized open-label design may need bias/ascertainment review.")

    min_age = profile.get("minimum_age_years")
    max_age = profile.get("maximum_age_years")
    if min_age is not None and max_age is not None and min_age > max_age:
        errors.append("Minimum age cannot be greater than maximum age.")

    return ValidationResult(errors=errors, warnings=warnings)


def protocol_section_status(profile: dict[str, Any]) -> dict[str, dict[str, Any]]:
    sections = {
        "identification": ["brief_title", "official_title", "brief_summary"],
        "sponsor": ["sponsor_name", "source_class", "responsible_party_type"],
        "conditions": ["conditions", "domain"],
        "status_and_dates": ["overall_status", "start_date", "primary_completion_date", "completion_date"],
        "design": [
            "study_type",
            "phase",
            "primary_purpose",
            "intervention_model",
            "allocation",
            "masking",
            "number_of_arms",
            "enrollment",
            "enrollment_type",
        ],
        "oversight": ["has_dmc"],
        "outcomes": [
            "primary_outcome_title",
            "primary_outcome_time_frame",
            "number_of_primary_outcomes",
            "number_of_secondary_outcomes",
        ],
        "eligibility": ["gender", "minimum_age_years", "maximum_age_years", "healthy_volunteers", "criteria"],
        "locations": ["number_of_facilities", "has_us_facility"],
        "sharing": ["ipd_sharing_plan"],
    }

    status: dict[str, dict[str, Any]] = {}
    for section, fields in sections.items():
        filled = [field for field in fields if _has_value(profile.get(field))]
        status[section] = {
            "filled": len(filled),
            "total": len(fields),
            "missing": [field for field in fields if field not in filled],
        }
    return status


def _parse_number(value: Any, caster: type[int] | type[float]) -> int | float | None:
    if value is None or value == "":
        return None
    try:
        number = caster(float(value)) if caster is int else caster(value)
    except (TypeError, ValueError):
        return None
    return number


def _parse_bool(value: Any) -> bool | None:
    if isinstance(value, bool):
        return value
    if value is None or value == "":
        return None
    normalized = str(value).strip().lower()
    if normalized in {"true", "yes", "1", "on"}:
        return True
    if normalized in {"false", "no", "0", "off"}:
        return False
    return None


def _has_value(value: Any) -> bool:
    return value is not None and value != ""


def _looks_like_iso_date(value: str) -> bool:
    parts = value.split("-")
    if len(parts) != 3:
        return False
    year, month, day = parts
    return len(year) == 4 and len(month) == 2 and len(day) == 2 and all(part.isdigit() for part in parts)