#!/home/xfr/.conda/envs/cbt/bin/python

import logging
import os
import re
import sys
import time
import queue as queue_module
import subprocess
import multiprocessing as mp
from concurrent.futures import ThreadPoolExecutor

import SimpleITK as sitk
import torch

# from config.device import get_device
from core.cylinder import create_coordinate_grid
from core.objective import set_global_context
from core.optimizer import run_pso_torch, run_de_torch, run_nm_torch, run_pso_torch_xfr
from imaging.orientation import azimuth_rotation, analyze_vertebral_tilt_contour
from imaging.transforms import level_file_path

import xfr_cbt_native

standardized_dir = '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr-3/'

azimuth_rotation_dir = '/mnt/1248/open2/cyrou/azimuth_rotation'
tilt_contour_dir = '/mnt/1248/open2/cyrou/tilt_contour'
Output_dir = '/mnt/1248/open/cyrou/Output'

def available_levels(volume_id):
    """該 volume 可跑 debug_pso 的全部 lumbar level：rotated/ 中輸入三件
    （_cortical / _binary_sdf / _roi）齊全的 L\\d。各 volume 層數不同
    （例：有的只有 L1~L3、有的含 L6），故不用固定 LEVELS 清單。"""
    rotated_dir = os.path.join(standardized_dir, volume_id, 'rotated')
    if not os.path.isdir(rotated_dir):
        return ()
    levels = []
    for fn in os.listdir(rotated_dir):
        m = re.fullmatch(r'(L[1-9]\d*)_cortical\.nii\.gz', fn)
        if not m:
            continue
        level = m.group(1)
        if (os.path.exists(os.path.join(rotated_dir, f'{level}_binary_sdf.nii.gz'))
                and os.path.exists(os.path.join(rotated_dir, f'{level}_roi.nii.gz'))):
            levels.append(level)
    return tuple(sorted(levels, key=lambda lv: int(lv[1:])))


LOG_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'logs')

# 單側任務的中繼結果（<run_id>/<volume_id>/<level>_<side>.json + plot lock）：
# 同一 level 的 L/R 可跑在不同 GPU，較晚完成的一側讀到兩側結果後跑合併輸出
SIDE_RESULT_DIR = os.path.join(LOG_DIR, 'side_results')

logger = logging.getLogger('xfr_debug')


# 目前任務標記（每個 worker 流程各自一份），讓交錯的 log 可以歸屬到 (volume, level, side)
_TASK_TAG = {'vid': None, 'level': None, 'side': None}

_NP_WRAP_RE = re.compile(r'\bnp\.[A-Za-z_][A-Za-z0-9_]*\(([^()]*)\)')


def set_task_tag(volume_id, level, side=None):
    _TASK_TAG['vid'] = volume_id
    _TASK_TAG['level'] = level
    _TASK_TAG['side'] = side


class _Tee:
    """同時輸出到 console 與 log 檔，逐行加上如 [0001 L1 LEFT] 的標記。
    volume 取 id 末段（1.3.6.1.4.1.9328.50.4.0001 -> 0001）；
    side 由 [LEFT]/[RIGHT]/左側/右側 段落標記判定並沿用給後續行；
    最終結果段屬於整椎，重置沿用值，其中的 Left/Right 摘要行只標該行本身；
    np.float64(...) 之類包裹會解包成裸數值。"""

    def __init__(self, console, log_fh):
        self.console = console
        self.log_fh = log_fh
        self.buf = ''

    @staticmethod
    def _marker_side(line):
        """回傳 (side, persist)：[LEFT]/[RIGHT]/左側/右側 是段落標記（persist=True），
        Left/Right 摘要行只標該行（persist=False）"""
        s = line.lstrip()
        if s.startswith('[LEFT]') or '左側' in s:
            return 'LEFT', True
        if s.startswith('[RIGHT]') or '右側' in s:
            return 'RIGHT', True
        if s.startswith('Left '):
            return 'LEFT', False
        if s.startswith('Right '):
            return 'RIGHT', False
        return None, False

    def _prefix(self, side):
        if not _TASK_TAG['vid'] or not _TASK_TAG['level']:
            return ''
        vol = _TASK_TAG['vid'].rsplit('.', 1)[-1]
        tag = f'{vol} {_TASK_TAG["level"]}'
        if side:
            tag += f' {side}'
        return f'[{tag}] '

    def _emit(self, line):
        line = _NP_WRAP_RE.sub(r'\1', line)
        if '最終結果' in line:
            _TASK_TAG['side'] = None
        side, persist = self._marker_side(line)
        if persist:
            _TASK_TAG['side'] = side
        eff_side = side if side is not None else _TASK_TAG['side']
        prefix = self._prefix(eff_side) if line.strip() else ''
        if prefix:
            # 前綴已含 side 時，去掉行首重複的 [LEFT] / [RIGHT] 標記
            marker = f'[{eff_side}] '
            if line.startswith(marker):
                line = line[len(marker):]
        self.console.write(prefix + line + '\n')
        self.log_fh.write(prefix + line + '\n')

    def write(self, data):
        if not data:
            return
        self.buf += data
        while True:
            idx_n = self.buf.find('\n')
            idx_r = self.buf.find('\r')
            candidates = [i for i in (idx_n, idx_r) if i != -1]
            if not candidates:
                break
            idx = min(candidates)
            self._emit(self.buf[:idx])
            self.buf = self.buf[idx + 1:]

    def flush(self):
        if self.buf:
            self._emit(self.buf)
            self.buf = ''
        self.console.flush()
        self.log_fh.flush()

    def isatty(self):
        return False


def setup_tee(log_path):
    """把這個流程的 stdout/stderr 同時寫到 log_path（append、line-buffered）"""
    log_fh = open(log_path, 'a', buffering=1)
    sys.stdout = _Tee(sys.stdout, log_fh)
    sys.stderr = _Tee(sys.stderr, log_fh)


def _setup_logging():
    """须在 setup_tee 之後呼叫，讓 handler 寫入 Tee（同時進 console 與 log 檔）"""
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )


def get_device(gpu_id=None):
    if torch.cuda.is_available():
        if gpu_id is None:
            max_free = -1
            gpu_id = 0
            for i in range(torch.cuda.device_count()):
                free_mem, _ = torch.cuda.mem_get_info(i)
                # print(f'GPU {gpu_id}: {torch.cuda.get_device_name(gpu_id)} {free_mem}')
                if free_mem > max_free:
                    max_free = free_mem
                    gpu_id = i
        device = torch.device(f"cuda:{gpu_id}")
        logger.info(f"Using GPU {gpu_id}: {torch.cuda.get_device_name(gpu_id)}")
    else:
        device = torch.device("cpu")
        logger.info("CUDA not available, using CPU")
    return device

def debug_orientation(volume_id, level):
  

    volume_dir = os.path.join(standardized_dir, volume_id)

    # _binary.nii.gz 現為原解析度；0.5mm 用 _binary_sdf（SDF 平滑遮罩）
    # 未旋轉檔：新世代在 crop/ 子資料夾、舊世代在頂層
    sdf = level_file_path(volume_dir, level, 'binary_sdf')
    binary_path = sdf if os.path.exists(sdf) \
        else level_file_path(volume_dir, level, 'binary')

    # azi = azimuth_rotation(binary_path)
    # res = analyze_vertebral_tilt_contour(binary_path, edge_type='superior', show_plot=False, debug=False)
    azi = azimuth_rotation(binary_path, show_plt=True, save_plt=True, output_path=f'{azimuth_rotation_dir}/{level}_{volume_id}.png')
    res = analyze_vertebral_tilt_contour(binary_path, edge_type='superior', show_plot=True, debug=False, save_plt=True, output_path=f'{tilt_contour_dir}/{level}_{volume_id}.png')
    alt = res['superior']['tilt_angle_deg']

    logger.info(binary_path)
    # print(f'Azimuth: {azi}, Alt: {alt}')
    logger.info(f'Alt: {alt}')

def debug_pso(volume_id, level, device=None, side='both', run_id=None):
    # ====== PSO ======
    swarm_size = 100
    max_iter = 100

    # ====== DEVICE ======
    if device is None:
        device = get_device()
  
    # ====== OTHER ======
    spacing = [0.5, 0.5, 0.5]
    CBT = True

    volume_dir = os.path.join(standardized_dir, volume_id)
    # _cortical 現只產出旋轉版（rotated/ 子資料夾）；_roi2 不再存檔，
    # 改用旋轉後 _roi（CT）。
    rotated_dir = os.path.join(volume_dir, 'rotated')
    cortical_path = os.path.join(rotated_dir, f'{level}_cortical.nii.gz')
    binary_path = os.path.join(rotated_dir, f'{level}_binary_sdf.nii.gz')
    roi_path = os.path.join(rotated_dir, f'{level}_roi.nii.gz')

    missing = [p for p in (cortical_path, binary_path, roi_path) if not os.path.exists(p)]
    if missing:
        raise ValueError(f'{volume_id} {level}: missing {missing}')

    cortical_image = sitk.ReadImage(cortical_path)
    binary_image = sitk.ReadImage(binary_path)
    roi_image = sitk.ReadImage(roi_path)

    cortical_array = sitk.GetArrayFromImage(cortical_image)
    binary_array = sitk.GetArrayFromImage(binary_image)
    roi_array = sitk.GetArrayFromImage(roi_image)
    image_shape = binary_array.shape

    # 資料層級的快速失敗（例：…9328.50.4.0653 L5 是 1-voxel 寬的退化體積）
    if binary_array.sum() == 0:
        raise ValueError(f'{volume_id} {level}: empty bone mask')
    if min(image_shape) < 8:
        raise ValueError(f'{volume_id} {level}: degenerate volume shape {image_shape}')

    cortical_tensor = torch.tensor(cortical_array, device=device)
    binary_tensor = torch.tensor(binary_array, device=device)

    grid = create_coordinate_grid(image_shape, device)

    set_global_context(
        cortical=cortical_tensor,
        spine=binary_tensor,
        shape=image_shape,
        spacing_=spacing,
        device_=device,
        grid_=grid,
        use_tip_penalty=False 
    )    
    
    best_l, loss_l, best_r, loss_r, total_time = run_pso_torch_xfr(
        label_str=f'{volume_id} {level}',
        image1_path=cortical_path,
        image2_path=binary_path,
        image3_path=roi_path,
        folder=Output_dir,
        swarm_size=swarm_size,
        max_iter=max_iter,
        spacing=spacing,
        CBT=CBT,
        device=device,
        optimize_size=True,
        grid=grid,
        side=side,
        level=level,
        patient_id=volume_id,
        side_dir=SIDE_RESULT_DIR,
        run_id=run_id,
        )

    # exit()


def list_gpu_ids():
    """透過 nvidia-smi 取得 GPU ID（主流程不初始化 CUDA，避免污染 fork/spawn 子流程）"""
    try:
        out = subprocess.check_output(
            ['nvidia-smi', '--query-gpu=index', '--format=csv,noheader'],
            text=True,
        )
        return [int(line.strip()) for line in out.splitlines() if line.strip()]
    except Exception:
        return []


def gpu_worker(gpu_id, log_path, run_id, task_queue, result_queue):
    """每張 GPU 一個工作流程：先鎖死該 GPU，再從共享隊列領 (volume, level) 任務"""
    # worker 是獨立流程，要自己把輸出 tee 到 log 檔
    setup_tee(log_path)
    _setup_logging()
    # 必須在任何 torch.cuda 呼叫前設定
    os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id)
    torch.cuda.set_device(0)
    device = torch.device('cuda:0')
    logger.info(f'=== [GPU {gpu_id}] worker started ===')

    while True:
        # 注意：multiprocessing.Queue 沒有 task_done()（只有 queue.Queue 有），別加回來
        item = task_queue.get()
        if item is None:
            break
        volume_id, level, side = item
        # tag 沿用 _Tee 的 LEFT/RIGHT 寫法（'L'/'R' 是 optimizer 的 side 值）
        set_task_tag(volume_id, level, 'LEFT' if side == 'L' else 'RIGHT')
        try:
            debug_pso(volume_id, level, device, side=side, run_id=run_id)
            result_queue.put(('task', gpu_id, volume_id, level, side, True, ''))
        except Exception as e:
            logger.error(f'[GPU {gpu_id}] Error in {volume_id} {level} {side}: {e}')
            result_queue.put(('task', gpu_id, volume_id, level, side, False, str(e)))

    result_queue.put(('done', gpu_id))
    logger.info(f'=== [GPU {gpu_id}] worker finished ===')


def _write_cbt_safe(volume_id, run_id):
    """per-volume 收尾：cbt.nii.gz + x-ap/x-lat 投影（錯誤已 log，回傳 ok）"""
    try:
        xfr_cbt_native.write_volume_cbt(volume_id, run_id)
        return True
    except Exception as e:
        logger.error(f'[CBT-NATIVE] {volume_id}: {e}')
        return False


def _run_sequential(tasks, run_id):
    """沒有（或只有一張）GPU 時的回退：單流程串行。
    tasks 依 (volume, level, side) 分組排列：一個 volume 的 (level, side)
    全部跑完後，立刻寫它的 cbt.nii.gz + 投影（不等其他 volume）。"""
    device = get_device()
    current_vid = None
    for volume_id, level, side in tasks:
        if volume_id != current_vid:
            if current_vid is not None:
                _write_cbt_safe(current_vid, run_id)
            current_vid = volume_id
        set_task_tag(volume_id, level, 'LEFT' if side == 'L' else 'RIGHT')
        try:
            debug_pso(volume_id, level, device, side=side, run_id=run_id)
        except Exception as e:
            logger.error(f'Error in {volume_id} {level} {side}: {e}')
    if current_vid is not None:
        _write_cbt_safe(current_vid, run_id)


USAGE = 'Usage: python xfr_debug.py [volume_id] [level]'


def parse_args(argv):
    """volume_id 可用完整 ID 或末段（如 0005）；level 為 L\\d 形式（L1、L2、…，
    不限定 L1~L5；實際執行以各 volume 資料中有的 level 為準，見 available_levels）。
    兩者可省略（＝全部）；level 必須搭配 volume_id 使用。"""
    vid_arg = argv[0] if len(argv) >= 1 else None
    level_arg = argv[1] if len(argv) >= 2 else None
    if len(argv) > 2:
        sys.exit(f'{USAGE}\nToo many arguments')
    if level_arg and not re.fullmatch(r'L[1-9]\d*', level_arg.upper()):
        sys.exit(f'{USAGE}\nInvalid level: {level_arg} (expected L1, L2, ...)')
    if level_arg and not vid_arg:
        sys.exit(f'{USAGE}\nlevel requires volume_id')
    return vid_arg, (level_arg.upper() if level_arg else None)


def main():
    # level = 'L1'

    # volume_id = '1.3.6.1.4.1.9328.50.4.0001'
    # volume_id = '1.3.6.1.4.1.9328.50.4.0003'
    # # volume_id = '1.3.6.1.4.1.9328.50.4.0121'
    # process_volume(volume_id, level)
    # exit()

    # 要處理的 volume 數（並行模式下是「最多嘗試的 volume 數」）
    MAX_SUCCESSFUL_VOLUMES = 100
    # MAX_SUCCESSFUL_VOLUMES = 10

    # log 檔（console 與檔案同時輸出；各 GPU worker 也會 append 進同一個檔）
    os.makedirs(LOG_DIR, exist_ok=True)
    run_id = time.strftime("%Y%m%d_%H%M%S")
    log_path = os.path.join(LOG_DIR, f'xfr_debug_{run_id}.log')
    setup_tee(log_path)
    _setup_logging()
    logger.info(f'Log file: {log_path}')
    logger.info(f'Command: {sys.executable} {" ".join(sys.argv)}')
    logger.info(f'Working directory: {os.getcwd()}')

    volumes = [d for d in sorted(os.listdir(standardized_dir))
               if os.path.isdir(os.path.join(standardized_dir, d))]

    vid_arg, level_arg = parse_args(sys.argv[1:])
    if vid_arg is not None:
        key = vid_arg.lower()
        vols = [v for v in volumes
                if v.lower() == key or v.rsplit('.', 1)[-1] == key]
        if not vols:
            sys.exit(f'Volume not found: {vid_arg}')
        volumes = vols
    volumes = volumes[:MAX_SUCCESSFUL_VOLUMES]

    # 各 volume 的 lumbar level：未指定 level 時，取該 volume 實際有的全部 lumbar level
    # （見 available_levels）；指定 level 時只跑該 level（該 volume 缺檔會直接報錯）
    vol_levels = {}
    for vid in volumes:
        vol_levels[vid] = (level_arg,) if level_arg else available_levels(vid)
    no_levels = [vid for vid in volumes if not level_arg and not vol_levels[vid]]
    if no_levels:
        logger.warning(f'{len(no_levels)} volume(s) have no available lumbar level, '
                       f'skipped: {", ".join(no_levels)}')
    # 任務粒度 = (volume, level, side)：同一 level 的 L/R 是兩個獨立任務，
    # 可被不同 GPU 的 worker 領走並行執行；順序 L1 L -> L1 R -> L2 L -> L2 R -> ...
    tasks = [(vid, level, s) for vid in volumes for level in vol_levels[vid] for s in ('L', 'R')]
    expected = {vid: len(vol_levels[vid]) * 2 for vid in volumes}
    all_levels = sorted({level for lv in vol_levels.values() for level in lv},
                        key=lambda lv: int(lv[1:]))
    logger.info(f'Total {len(volumes)} volumes / {len(tasks)} (volume, level, side) tasks '
                f'(levels: {", ".join(all_levels) if all_levels else "(none)"})')
    if vid_arg or level_arg:
        logger.info(f'Filter: volume_id={vid_arg!r} level={level_arg!r}')

    gpu_ids = list_gpu_ids()

    if len(gpu_ids) <= 1:
        logger.info(f'Only {len(gpu_ids)} GPU(s) available, running sequentially')
        _run_sequential(tasks, run_id)
        return

    logger.info(f'Found {len(gpu_ids)} GPUs: {gpu_ids}, starting {len(gpu_ids)} workers (one per GPU)')

    ctx = mp.get_context('spawn')
    task_queue = ctx.Queue()
    result_queue = ctx.Queue()
    for t in tasks:
        task_queue.put(t)
    for _ in gpu_ids:
        task_queue.put(None)  # 每個 worker 一個結束哨兵

    # per-volume 收尾：一個 volume 排入的 (level, side) 任務全部回報後
    # （成功或失敗），立刻在背景執行緒寫它的 cbt.nii.gz + 投影，不等全部 case
    write_pool = ThreadPoolExecutor(max_workers=4, thread_name_prefix='cbt-write')
    write_futures = set()
    pending_writes = set(volumes)
    remaining = {vid: expected[vid] for vid in volumes}

    def _fire_write(vid):
        pending_writes.discard(vid)
        write_futures.add(write_pool.submit(_write_cbt_safe, vid, run_id))

    def _on_task_result(msg):
        results.append(msg)
        vid = msg[2]
        if vid in remaining:
            remaining[vid] -= 1
            if remaining[vid] <= 0 and vid in pending_writes:
                logger.info(f'{vid}: screw tasks complete, writing cbt + projections now')
                _fire_write(vid)

    procs = [ctx.Process(target=gpu_worker, args=(g, log_path, run_id, task_queue, result_queue), name=f'cbt-gpu-{g}')
              for g in gpu_ids]
    for p in procs:
        p.start()

    start_time = time.time()
    results = []
    finished = 0
    while finished < len(gpu_ids):
        if not any(p.is_alive() for p in procs):
            logger.warning('A worker exited early; reaping remaining tasks...')
            break
        try:
            msg = result_queue.get(timeout=5)
        except queue_module.Empty:
            continue
        if msg[0] == 'done':
            finished += 1
        else:
            _on_task_result(msg)

    # 抽乾剩下排進來的結果
    while True:
        try:
            msg = result_queue.get(timeout=1)
        except queue_module.Empty:
            break
        if msg[0] == 'task':
            _on_task_result(msg)

    for p in procs:
        p.join(timeout=60)

    # 補漏：worker 提早退出、有任務未回報的 volume 仍照舊嘗試寫
    # （無 side 結果時 write_volume_cbt 會自行 skip）
    for vid in volumes:
        if vid in pending_writes:
            _fire_write(vid)

    # 等待所有 per-volume cbt.nii.gz / 投影寫出完成
    write_ok = sum(1 for f in write_futures if f.result())
    write_fail = len(write_futures) - write_ok
    write_pool.shutdown(wait=True)

    total_time = time.time() - start_time
    ok = [r for r in results if r[5]]
    fail = [r for r in results if not r[5]]
    per_volume = {}
    per_volume_n = {}
    for _, _, vid, level, side, success, _ in results:
        per_volume.setdefault(vid, set()).add(success)
        per_volume_n[vid] = per_volume_n.get(vid, 0) + 1
    missing = len(tasks) - len(results)

    # 一個 volume 算「成功」必須它的所有 (level, side) 任務都執行過且全部成功
    # （expected 為該 volume 實際排入的任務數，各 volume 的 level 數可不同）
    n_success_volumes = sum(
        1 for vid in volumes
        if per_volume.get(vid, set()) == {True}
        and per_volume_n.get(vid, 0) == expected.get(vid, 0))

    print('=' * 60)
    logger.info(f'Finished in {total_time / 60:.1f} min | '
                f'tasks {len(results)}/{len(tasks)} (ok {len(ok)} / failed {len(fail)} / not-run {missing})')
    if missing:
        logger.warning(f'{missing} task(s) were never executed (worker crash?)')
    logger.info(f'Successful volumes (all levels OK): {n_success_volumes}/{len(volumes)}')
    if fail:
        logger.info('Failed tasks:')
        for _, g, vid, level, side, _, err in fail:
            logger.error(f'[GPU {g}] {vid} {level} {side}: {err}')

    # cbt.nii.gz + x-ap.jpg / x-lat.jpg 已於各 volume 的螺絲任務完成後立刻
    # 寫出（_fire_write， Output_dir/<run_date>/<volume_id>/；label 1-10 =
    # L1L L1R L2L L2R ... L5L L5R），不再等全部 case 跑完才統一收尾
    logger.info(f'CBT writes: {write_ok}/{len(volumes)} volume(s) ok'
                + (f', {write_fail} failed (見 [CBT-NATIVE] log)' if write_fail else ''))


if __name__ == '__main__':
    main()