import os
import re
import sys
import time
import queue as queue_module
import subprocess
import multiprocessing as mp

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

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

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'

LEVELS = ('L1', 'L2', 'L3', 'L4', 'L5')

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


# 目前任務標記（每個 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):
    _TASK_TAG['vid'] = volume_id
    _TASK_TAG['level'] = level
    _TASK_TAG['side'] = None


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 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}")
        print(f"Using GPU {gpu_id}: {torch.cuda.get_device_name(gpu_id)}")
    else:
        device = torch.device("cpu")
        print("CUDA not available, using CPU")
    return device

def debug_orientation(volume_id, level):
  

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

    cortical_path = os.path.join(volume_dir, f'{level}_cortical.nii.gz')
    binary_path = os.path.join(volume_dir, f'{level}_binary.nii.gz')
    roi_path = os.path.join(volume_dir, f'{level}_roi2.nii.gz')

    # 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']

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

def debug_pso(volume_id, level, device=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_path = os.path.join(volume_dir, f'{level}_cortical.nii.gz')
    binary_path = os.path.join(volume_dir, f'{level}_binary.nii.gz')
    roi_path = os.path.join(volume_dir, f'{level}_roi2.nii.gz')

    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,
        )

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

    while True:
        # 注意：multiprocessing.Queue 沒有 task_done()（只有 queue.Queue 有），別加回來
        item = task_queue.get()
        if item is None:
            break
        volume_id, level = item
        set_task_tag(volume_id, level)
        try:
            debug_pso(volume_id, level, device)
            result_queue.put(('task', gpu_id, volume_id, level, True, ''))
        except Exception as e:
            print(f'[GPU {gpu_id}] Error in {volume_id} {level}: {e}', flush=True)
            result_queue.put(('task', gpu_id, volume_id, level, False, str(e)))

    result_queue.put(('done', gpu_id))
    print(f'=== [GPU {gpu_id}] worker finished ===', flush=True)


def _run_sequential(tasks):
    """沒有（或只有一張）GPU 時的回退：單流程串行"""
    device = get_device()
    for volume_id, level in tasks:
        set_task_tag(volume_id, level)
        try:
            debug_pso(volume_id, level, device)
        except Exception as e:
            print(f'Error in {volume_id} {level}: {e}')


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 = 1

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

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

    tasks = [(vid, level) for vid in volumes for level in LEVELS]
    print(f'Total {len(volumes)} volumes / {len(tasks)} (volume, level) tasks', flush=True)

    gpu_ids = list_gpu_ids()

    if len(gpu_ids) <= 1:
        print(f'Only {len(gpu_ids)} GPU(s) available, running sequentially', flush=True)
        _run_sequential(tasks)
        return

    print(f'Found {len(gpu_ids)} GPUs: {gpu_ids}, starting {len(gpu_ids)} workers (one per GPU)', flush=True)

    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 一個結束哨兵

    procs = [ctx.Process(target=gpu_worker, args=(g, log_path, 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):
            print('Warning: a worker exited early; reaping remaining tasks...', flush=True)
            break
        try:
            msg = result_queue.get(timeout=5)
        except queue_module.Empty:
            continue
        if msg[0] == 'done':
            finished += 1
        else:
            results.append(msg)

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

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

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

    # 一個 volume 算「成功」必須它的(level)全部執行過且全部成功
    n_success_volumes = sum(1 for vid in volumes
                            if per_volume.get(vid, set()) == {True})

    print('=' * 60)
    print(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:
        print(f'Warning: {missing} task(s) were never executed (worker crash?)')
    print(f'Successful volumes (all levels OK): {n_success_volumes}/{len(volumes)}')
    if fail:
        print('Failed tasks:')
        for _, g, vid, level, _, err in fail:
            print(f'  [GPU {g}] {vid} {level}: {err}')


if __name__ == '__main__':
    main()