package_code_imports/v5_scripts/config_v5_train.py

from __future__ import annotations

from pathlib import Path


SCRIPT_DIR = Path(__file__).resolve().parent
V5_ROOT = SCRIPT_DIR.parent
PACKAGE_ROOT = V5_ROOT

DEFAULT_TENSOR_DIR = Path(
    "/Users/yang/Downloads/精卫_SCTM_数据交付包/60_pretraining/"
    "sctm_v3p1/tensor_shards_v3p5_static_collapsed_seq128_stride64"
)
DEFAULT_STAGE0_DIR = Path("/Users/yang/Downloads/精卫_SCTM_数据交付包/60_pretraining/sctm_v3p1/stage0")
DEFAULT_SAMPLE_TENSOR_DIR = PACKAGE_ROOT / "data_samples" / "tensor_shards_sample"
DEFAULT_SAMPLE_STAGE0_DIR = PACKAGE_ROOT / "data_samples" / "stage0"

ACTION_AVAILABILITY = {
    "not_available": 0,
    "history": 1,
    "contact": 2,
    "post": 3,
}

ACTION_BLOCKS = (
    "medication",
    "service_pathway",
    "observation_management",
    "recovery_support",
)

V5_LOSS_WEIGHTS = {
    "behavior_policy_weight": 0.20,
    "action_pwe_weight": 1.00,
    "action_service_weight": 1.00,
    "action_missingness_weight": 0.40,
    "action_field_weight": 0.30,
    "action_numeric_weight": 0.10,
    "action_ordinal_weight": 0.20,
    "action_event_generation_weight": 0.20,
    "residual_l2_weight": 1.0e-4,
}

TARGET_TRIAL_DEFAULTS = {
    "name": "lai_initiation_vs_continued_oral",
    "primary_grace_days": 60.0,
    "sensitivity_grace_days": (30.0, 90.0),
    "primary_followup_days": 365.0,
    "sensitivity_followup_days": (180.0, 730.0),
    "primary_estimand": "ATT among eligible oral-treated patient-windows with overlap",
}

V5_ACTION_SLOT_NAMES = (
    "treatment_mode",
    "medication_administration",
    "medication_timing",
    "free_medication",
    "referral_flag",
    "referral_type",
    "followup_type",
    "visit_mode",
    "visit_target",
    "emergency_response",
    "health_check",
    "specialist_treatment_recommendation",
    "inpatient_signal",
    "management_tier_transition",
    "lai_present",
    "oral_present",
)

DERIVED_BINARY_LABELS = ("not_present", "present", "unknown")
MGMT_TRANSITION_LABELS = ("down_2", "down_1", "no_change", "up_1", "up_2", "unknown")


__all__ = [
    "ACTION_AVAILABILITY",
    "ACTION_BLOCKS",
    "DEFAULT_SAMPLE_STAGE0_DIR",
    "DEFAULT_SAMPLE_TENSOR_DIR",
    "DEFAULT_STAGE0_DIR",
    "DEFAULT_TENSOR_DIR",
    "DERIVED_BINARY_LABELS",
    "MGMT_TRANSITION_LABELS",
    "TARGET_TRIAL_DEFAULTS",
    "V5_ACTION_SLOT_NAMES",
    "V5_LOSS_WEIGHTS",
]