# 104 Anatomical Structure Labels for TotalSegmentator / MONAI wholeBody_ct_segmentation
# Based on TotalSegmentator v1 class definitions

LABEL_NAMES = {
    0: "background",
    1: "spleen",
    2: "kidney_right",
    3: "kidney_left",
    4: "gallbladder",
    5: "liver",
    6: "stomach",
    7: "aorta",
    8: "inferior_vena_cava",
    9: "portal_vein_and_splenic_vein",
    10: "pancreas",
    11: "adrenal_gland_right",
    12: "adrenal_gland_left",
    13: "lung_upper_lobe_left",
    14: "lung_lower_lobe_left",
    15: "lung_upper_lobe_right",
    16: "lung_middle_lobe_right",
    17: "lung_lower_lobe_right",
    18: "vertebrae_L5",
    19: "vertebrae_L4",
    20: "vertebrae_L3",
    21: "vertebrae_L2",
    22: "vertebrae_L1",
    23: "vertebrae_T12",
    24: "vertebrae_T11",
    25: "vertebrae_T10",
    26: "vertebrae_T9",
    27: "vertebrae_T8",
    28: "vertebrae_T7",
    29: "vertebrae_T6",
    30: "vertebrae_T5",
    31: "vertebrae_T4",
    32: "vertebrae_T3",
    33: "vertebrae_T2",
    34: "vertebrae_T1",
    35: "vertebrae_C7",
    36: "vertebrae_C6",
    37: "vertebrae_C5",
    38: "vertebrae_C4",
    39: "vertebrae_C3",
    40: "vertebrae_C2",
    41: "vertebrae_C1",
    42: "esophagus",
    43: "trachea",
    44: "heart_myocardium",
    45: "heart_atrium_left",
    46: "heart_ventricle_left",
    47: "heart_atrium_right",
    48: "heart_ventricle_right",
    49: "pulmonary_artery",
    50: "brain",
    51: "iliac_artery_left",
    52: "iliac_artery_right",
    53: "iliac_vena_left",
    54: "iliac_vena_right",
    55: "small_bowel",
    56: "duodenum",
    57: "colon",
    58: "rib_left_1",
    59: "rib_left_2",
    60: "rib_left_3",
    61: "rib_left_4",
    62: "rib_left_5",
    63: "rib_left_6",
    64: "rib_left_7",
    65: "rib_left_8",
    66: "rib_left_9",
    67: "rib_left_10",
    68: "rib_left_11",
    69: "rib_left_12",
    70: "rib_right_1",
    71: "rib_right_2",
    72: "rib_right_3",
    73: "rib_right_4",
    74: "rib_right_5",
    75: "rib_right_6",
    76: "rib_right_7",
    77: "rib_right_8",
    78: "rib_right_9",
    79: "rib_right_10",
    80: "rib_right_11",
    81: "rib_right_12",
    82: "humerus_left",
    83: "humerus_right",
    84: "scapula_left",
    85: "scapula_right",
    86: "clavicula_left",
    87: "clavicula_right",
    88: "femur_left",
    89: "femur_right",
    90: "hip_left",
    91: "hip_right",
    92: "sacrum",
    93: "face",
    94: "gluteus_maximus_left",
    95: "gluteus_maximus_right",
    96: "gluteus_medius_left",
    97: "gluteus_medius_right",
    98: "gluteus_minimus_left",
    99: "gluteus_minimus_right",
    100: "autochthon_left",
    101: "autochthon_right",
    102: "iliopsoas_left",
    103: "iliopsoas_right",
    104: "urinary_bladder",
}

# Color map for visualization (RGB values)
# Using a custom colormap for better visualization
import numpy as np

def get_color_map():
    """Generate a color map for 105 classes (background + 104 structures)"""
    np.random.seed(42)  # For reproducibility
    colors = np.zeros((105, 3), dtype=np.uint8)
    
    # Background is black
    colors[0] = [0, 0, 0]
    
    # Assign distinct colors to different organ categories
    # Organs (1-12): Warm colors
    organ_colors = [
        [255, 99, 71],   # spleen - tomato
        [255, 165, 0],   # kidney_right - orange
        [255, 140, 0],   # kidney_left - dark orange
        [50, 205, 50],   # gallbladder - lime green
        [139, 69, 19],   # liver - saddle brown
        [255, 192, 203], # stomach - pink
        [220, 20, 60],   # aorta - crimson
        [0, 0, 139],     # inferior_vena_cava - dark blue
        [138, 43, 226],  # portal_vein_and_splenic_vein - blue violet
        [255, 215, 0],   # pancreas - gold
        [255, 255, 0],   # adrenal_gland_right - yellow
        [255, 255, 0],   # adrenal_gland_left - yellow
    ]
    colors[1:13] = organ_colors
    
    # Lungs (13-17): Light blue shades
    colors[13:18] = [[135, 206, 235], [100, 149, 237], [30, 144, 255], [0, 191, 255], [70, 130, 180]]
    
    # Vertebrae (18-41): Gradient from red to purple
    for i in range(18, 42):
        colors[i] = [200 - (i-18)*5, 100, 150 + (i-18)*3]
    
    # Other structures (42-57): Various colors
    colors[42] = [255, 182, 193]  # esophagus
    colors[43] = [176, 224, 230]  # trachea
    colors[44:49] = [[220, 20, 60], [255, 105, 180], [255, 20, 147], [255, 182, 193], [199, 21, 133]]  # heart
    colors[49] = [148, 0, 211]    # pulmonary_artery
    colors[50] = [255, 218, 185]  # brain
    colors[51:55] = [[178, 34, 34], [178, 34, 34], [70, 130, 180], [70, 130, 180]]  # iliac vessels
    colors[55:58] = [[222, 184, 135], [210, 180, 140], [188, 143, 143]]  # bowels
    
    # Ribs (58-81): Bone color variations
    for i in range(58, 82):
        colors[i] = [255, 250, 205 + (i-58) % 50]
    
    # Bones (82-92): Gray/white shades
    for i in range(82, 93):
        colors[i] = [220 + (i-82)*2, 220 + (i-82)*2, 220]
    
    # Face and muscles (93-103): Skin and muscle tones
    colors[93] = [255, 228, 196]  # face
    for i in range(94, 104):
        colors[i] = [205, 92, 92 + (i-94)*5]  # muscles - indian red variations
    
    # Bladder
    colors[104] = [255, 255, 0]  # urinary_bladder - yellow
    
    return colors

def get_label_name(label_id: int) -> str:
    """Get human-readable name for a label ID"""
    return LABEL_NAMES.get(label_id, f"unknown_{label_id}").replace("_", " ").title()

def get_organ_categories():
    """Group organs by category for UI display"""
    return {
        "Major Organs": [1, 2, 3, 4, 5, 6, 10, 104],
        "Cardiovascular": [7, 8, 9, 44, 45, 46, 47, 48, 49, 51, 52, 53, 54],
        "Respiratory": [13, 14, 15, 16, 17, 43],
        "Digestive": [42, 55, 56, 57],
        "Vertebrae": list(range(18, 42)),
        "Ribs": list(range(58, 82)),
        "Upper Body Bones": [82, 83, 84, 85, 86, 87],
        "Lower Body Bones": [88, 89, 90, 91, 92],
        "Muscles": list(range(94, 104)),
        "Other": [11, 12, 50, 93],
    }