"""UI preset definitions: anatomy choices, sample-input chips, body regions."""
from __future__ import annotations

# Body regions accepted by the upstream CT pipeline (input to LDMSampler).
CT_BODY_REGIONS = ["head", "chest", "thorax", "abdomen", "pelvis", "lower"]

# Controllable anatomies for paired CT generation. Subset of label_dict.json that
# the upstream pipeline knows how to condition on (per docs/inference.md).
CT_ANATOMY_CHOICES = [
    "liver",
    "spleen",
    "pancreas",
    "right kidney",
    "left kidney",
    "gallbladder",
    "stomach",
    "duodenum",
    "bladder",
    "colon",
    "small bowel",
    "left lung upper lobe",
    "left lung lower lobe",
    "right lung upper lobe",
    "right lung middle lobe",
    "right lung lower lobe",
    "lung tumor",
    "hepatic tumor",
    "pancreatic tumor",
    "bone lesion",
]

# MR contrast options exposed in the dropdown (maps to modality int)
MR_CONTRAST_CHOICES = [
    ("T1 brain", 9),
    ("T2 brain", 10),
    ("FLAIR (skull-stripped brain)", 11),
    ("T2 prostate", 10),
    ("T1 breast", 9),
    ("T1 abdomen", 9),
    ("T2 abdomen", 10),
    ("Generic MRI", 8),
]

MR_BRAIN_CONTRASTS = ["T1", "T2", "FLAIR", "SWI", "T1_skull_stripped", "T2_skull_stripped", "FLAIR_skull_stripped", "SWI_skull_stripped"]

# Paired CT requires X==Y from {256, 384, 512} and Z from {128, 256, 384, 512, 640, 768}.
# We expose 256 and 384 by default (skipping 512 to stay within ~70 GB VRAM).
XY_CHOICES = [256, 384]
Z_CHOICES = [128, 256, 384]


CT_SAMPLES = [
    {
        "label": "Whole abdomen",
        "body_region": ["abdomen"],
        "anatomy_list": [
            "liver", "spleen", "pancreas", "right kidney", "left kidney",
            "gallbladder", "stomach", "duodenum", "colon", "small bowel",
        ],
        "xy": 384, "z": 256, "spacing": [1.5, 1.5, 1.5],
    },
    {
        "label": "Full chest",
        "body_region": ["chest"],
        "anatomy_list": [
            "left lung upper lobe", "left lung lower lobe",
            "right lung upper lobe", "right lung middle lobe", "right lung lower lobe",
        ],
        "xy": 384, "z": 256, "spacing": [1.5, 1.5, 2.0],
    },
    {
        "label": "Liver + lesion",
        "body_region": ["abdomen"],
        "anatomy_list": ["liver", "hepatic tumor", "spleen", "gallbladder"],
        "xy": 256, "z": 256, "spacing": [1.5, 1.5, 1.5],
    },
    {
        "label": "Pelvis organs",
        "body_region": ["pelvis"],
        "anatomy_list": ["bladder", "colon", "small bowel"],
        "xy": 256, "z": 256, "spacing": [1.5, 1.5, 1.5],
    },
]

# FOV-derived from rflow-mr training-data medians (FOV mm = voxels × spacing);
# the three highest-volume body regions in the training set:
#   T1 brain    160 × 256 × 256 mm   (n=4,659)
#   T1 breast   174 × 200 × 200 mm   (n=2,162)
#   T2 prostate 170 × 170 ×  90 mm   (n=898)
MR_SAMPLES = [
    {"label": "T1 brain", "modality_label": "T1 brain", "xy": 256, "z": 256, "spacing": [0.65, 1.0, 1.0]},
    {"label": "T1 breast", "modality_label": "T1 breast", "xy": 256, "z": 256, "spacing": [0.70, 0.80, 0.80]},
    {"label": "T2 prostate", "modality_label": "T2 prostate", "xy": 256, "z": 128, "spacing": [0.65, 0.65, 0.70]},
]

MR_BRAIN_SAMPLES = [
    {"label": "T1 whole brain", "contrast": "T1", "xy": 256, "z": 256, "spacing": [1.0, 1.0, 1.0]},
    {"label": "T2 whole brain", "contrast": "T2", "xy": 256, "z": 256, "spacing": [1.0, 1.0, 1.0]},
    {"label": "FLAIR", "contrast": "FLAIR", "xy": 256, "z": 256, "spacing": [1.0, 1.0, 1.0]},
    {"label": "SWI", "contrast": "SWI", "xy": 256, "z": 256, "spacing": [1.0, 1.0, 1.0]},
]