"""
척추 라벨링 데이터 뷰어
- NII.GZ: 3D 세그멘테이션 마스크
- PF: MITK PlanarFigure (2D 측정)
- DCM: DICOM 원본 이미지
"""

import nibabel as nib
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider
from matplotlib.lines import Line2D
import xml.etree.ElementTree as ET
from pathlib import Path
import re

# --- 설정 (사용자 수정 가능) ---
DATA_ROOT = Path('.') # 데이터가 있는 최상위 경로
# -----------------------------

BASE_PATH = None
DCM_FILE = None
MITK_ORIGIN_Y = 966.81
SPACING = 0.148


def load_nii(filepath):
    img = nib.load(filepath)
    return img.get_fdata(), img.header


def load_dcm(filepath):
    import pydicom
    dcm = pydicom.dcmread(filepath)
    return dcm.pixel_array


def parse_pf(filepath):
    with open(filepath, 'r', encoding='utf-8') as f:
        content = f.read()
    content = re.sub(r'<\?xml[^?]*\?>', '', content)
    root = ET.fromstring(f'<root>{content}</root>')
    pf = root.find('PlanarFigure')

    return {
        'type': pf.get('type'),
        'vertices': [(float(v.get('x')), float(v.get('y'))) for v in pf.findall('.//Vertex')]
    }


def pf_to_pixel(x, y):
    """MITK 좌표 -> 픽셀 좌표 변환"""
    return x / SPACING, (MITK_ORIGIN_Y - y) / SPACING


def draw_pf(ax, pf_files, colors):
    labels = []
    for i, pf_file in enumerate(pf_files):
        data = parse_pf(str(pf_file))
        color = colors[i % len(colors)]
        transformed = [pf_to_pixel(x, y) for x, y in data['vertices']]

        if 'Line' in data['type']:
            xs, ys = zip(*transformed)
            ax.plot(xs, ys, '-o', color=color, linewidth=3, markersize=8)
        elif 'Circle' in data['type'] and len(transformed) >= 2:
            center, edge = transformed[0], transformed[1]
            radius = np.hypot(edge[0]-center[0], edge[1]-center[1])
            ax.add_patch(plt.Circle(center, radius, fill=False, color=color, linewidth=3))
            ax.plot(center[0], center[1], 'o', color=color, markersize=8)

        labels.append((color, pf_file.stem))
    return labels


def view_nii_overlay():
    """모든 NII 마스크를 색상별로 겹쳐서 표시"""
    nii_files = list(BASE_PATH.glob('*.nii.gz'))
    first_data, _ = load_nii(str(nii_files[0]))
    combined = np.zeros(first_data.shape)

    for i, f in enumerate(nii_files, 1):
        data, _ = load_nii(str(f))
        combined[data > 0] = i

    fig, ax = plt.subplots(figsize=(10, 10))
    mid_z = combined.shape[2] // 2

    slider_ax = plt.axes([0.2, 0.02, 0.6, 0.03])
    slider = Slider(slider_ax, 'Slice', 0, combined.shape[2]-1, valinit=mid_z, valstep=1)

    plt.sca(ax)
    img_plot = ax.imshow(combined[:, :, mid_z].T, cmap='nipy_spectral', origin='lower')
    ax.set_title(f'All Vertebrae - Slice {mid_z}')
    plt.colorbar(img_plot, ax=ax)

    def update(val):
        idx = int(slider.val)
        img_plot.set_data(combined[:, :, idx].T)
        ax.set_title(f'All Vertebrae - Slice {idx}')
        fig.canvas.draw_idle()

    slider.on_changed(update)
    plt.show()


def view_all_on_dcm():
    """DCM + NII 마스크 + PF 측정 오버레이"""
    img = load_dcm(str(DCM_FILE))
    nii_files = list(BASE_PATH.glob('*.nii.gz'))
    pf_files = list(BASE_PATH.glob('*.pf'))

    fig, ax = plt.subplots(figsize=(10, 20))
    ax.imshow(img, cmap='gray', origin='upper')

    # NII 마스크 컨투어
    cmap = plt.colormaps.get_cmap('tab20').resampled(len(nii_files))
    nii_labels = []
    for i, f in enumerate(nii_files):
        data, _ = load_nii(str(f))
        proj = np.max(data, axis=2).T  # (W,H,D) -> (H,W)
        if proj.max() > 0:
            ax.contour(proj, levels=[0.5], colors=[cmap(i)], linewidths=1.5, alpha=0.8)
        nii_labels.append((cmap(i), f.stem.split('.')[0]))

    # PF 오버레이
    pf_labels = draw_pf(ax, pf_files, ['red', 'yellow', 'cyan', 'lime', 'magenta'])

    # 범례
    legend = [Line2D([0], [0], color=c, label=n) for c, n in nii_labels]
    legend += [Line2D([0], [0], color=c, marker='o', label=f'[PF] {n}') for c, n in pf_labels]
    ax.legend(handles=legend, loc='upper left', bbox_to_anchor=(1.02, 1), fontsize=8)
    ax.set_title('DCM + NII Masks + PF Annotations')
    plt.tight_layout()
    plt.show()


def select_directory():
    """디렉토리 선택"""
    global BASE_PATH, DCM_FILE
    dirs = sorted([d for d in DATA_ROOT.iterdir() if d.is_dir() and not d.name.startswith('.')])

    if not dirs:
        print("하위 디렉토리가 없습니다.")
        return False

    print("\n=== 디렉토리 선택 ===")
    for i, d in enumerate(dirs, 1):
        print(f"{i}. {d.name}")

    try:
        idx = int(input(f"\n선택 (1-{len(dirs)}): ").strip()) - 1
        if 0 <= idx < len(dirs):
            BASE_PATH = dirs[idx]
            dcm_files = list(BASE_PATH.glob('*.dcm'))
            DCM_FILE = dcm_files[0] if dcm_files else None
            print(f"선택됨: {BASE_PATH.name}")
            return True
    except (ValueError, IndexError):
        pass
    print("잘못된 선택")
    return False


if __name__ == '__main__':
    print("=== 척추 라벨링 뷰어 ===")
    if not select_directory():
        exit()

    while True:
        print(f"\n=== {BASE_PATH.name} ===")
        print("1. NII 마스크 오버레이 (슬라이더)")
        print("2. DCM + NII + PF 전체 오버레이")
        print("q. 종료")

        choice = input("\n선택: ").strip().lower()
        if choice == 'q':
            break
        elif choice == '1':
            view_nii_overlay()
        elif choice == '2':
            if DCM_FILE:
                view_all_on_dcm()
            else:
                print("DCM 파일이 없습니다.")