#!/home/xfr/.conda/envs/cbt/bin/python
"""
從 optimizer 存的 output.csv 重新渲染既有螺絲模式四視角圖（X-ray 圖）。

用途：render_bone_figure 改了純顯示層（例如棘突改半透明 + 依深度正確遮擋
螺絲路徑）之後，不需要重跑優化，直接從 CSV 的 best_position 重新出圖。
原圖先備份到 Output/{date}/backup_opaque_spinous/{vol}/，再重渲染同檔名覆蓋。

mask 取法與 xfr_debug / xfr_plot_level 相同（level_file_path：
binary_sdf 優先、缺則 binary；cortical 同），spacing 讀自 mask 檔。

Usage:
  python xfr_rerender_spinous.py                      # 預設 20260912，多 GPU 並行
  python xfr_rerender_spinous.py 20260911             # 指定日期
  python xfr_rerender_spinous.py 20260912 0001        # 只該 volume（全名或末段）
  python xfr_rerender_spinous.py --dry-run 20260912   # 只列出會重新渲染的圖
  python xfr_rerender_spinous.py --cpus 20260912      # 純 CPU 循序
"""

import argparse
import csv
import multiprocessing
import os
import re
import shutil
import sys

import numpy as np
import SimpleITK as sitk

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from imaging.transforms import level_file_path
from visualization.res_bone_figure import render_bone_figure

OUTPUT_BASE = '/mnt/1248/open2/cyrou/Output'
MASK_DIR = '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr'
BACKUP = 'backup_opaque_spinous'

# {level}_{way}_L{d}_{l}_R{d}_{l}_{swarm}_{iter}.png（單側跑法該側可為空）
PNG_RE = re.compile(
    r'^(?P<level>[A-Z]\d+?)_(?P<way>CBT|TPS)'
    r'_L(?P<dl>\d+(?:\.\d+)?|)_(?P<ll>\d+(?:\.\d+)?|)'
    r'_R(?P<dr>\d+(?:\.\d+)?|)_(?P<lr>\d+(?:\.\d+)?|)'
    r'_(?P<sw>\d+|)_(?P<it>\d+|)\.png$')


def _f(s):
    s = (s or '').strip()
    return float(s) if s else None


def _parse_pos(s):
    """CSV '(x, y, z)' -> (z, y, x)（best_position 慣例）。"""
    v = re.findall(r'[-+]?\d*\.?\d+(?:[eE][-+]?\d+)?', s or '')
    if len(v) < 3:
        return None
    x, y, z = (float(t) for t in v[:3])
    return z, y, x


def _pos(row):
    if row is None:
        return None
    p = _parse_pos(row.get('Position_XYZ'))
    if p is None:
        return None
    az, alt = _f(row.get('Raw_Azimuth')), _f(row.get('Raw_Altitude'))
    if az is None or alt is None:
        return None
    return (p[0], p[1], p[2], az, alt)


def read_csv_rows(vol_out_dir):
    path = os.path.join(vol_out_dir, 'output.csv')
    if not os.path.isfile(path):
        return None
    with open(path, newline='') as f:
        return [r for r in csv.DictReader(f) if any(c.strip() for c in r.values())]


def find_row(rows, side, d, l):
    if d is None or l is None:
        return None
    for r in rows:
        if r.get('Side') != side:
            continue
        rd, rl = _f(r.get('Diameter')), _f(r.get('Length'))
        if rd is not None and rl is not None \
                and abs(rd - d) < 1e-6 and abs(rl - l) < 1e-6:
            return r
    return None


def collect_figs(date_dir, vol_filter=None):
    """回傳 (tasks, skipped)：tasks 為可重新渲染的圖（dict），skipped 為 (原因, 路徑)。"""
    vols = sorted(d for d in os.listdir(date_dir)
                  if os.path.isdir(os.path.join(date_dir, d)) and d != BACKUP)
    if vol_filter:
        vols = [v for v in vols
                if v == vol_filter or v.rsplit('.', 1)[-1] == vol_filter]
        if not vols:
            raise SystemExit(f'volume not found: {vol_filter}')
    tasks, skipped = [], []
    for vol in vols:
        vdir = os.path.join(date_dir, vol)
        rows = read_csv_rows(vdir) or []
        for name in sorted(os.listdir(vdir)):
            m = PNG_RE.match(name)
            if not m:
                continue
            png = os.path.join(vdir, name)
            level, way = m['level'], m['way']
            d_l, l_l = _f(m['dl']), _f(m['ll'])
            d_r, l_r = _f(m['dr']), _f(m['lr'])
            row_l = find_row(rows, 'L', d_l, l_l)
            row_r = find_row(rows, 'R', d_r, l_r)
            if (d_l is not None and row_l is None) or (d_r is not None and row_r is None) \
                    or (d_l is None and d_r is None):
                skipped.append((vol, name, 'CSV 無對應 L/R 行'))
                continue
            mask_dir = os.path.join(MASK_DIR, vol)
            sdf = level_file_path(mask_dir, level, 'binary_sdf')
            binary_path = sdf if os.path.exists(sdf) \
                else level_file_path(mask_dir, level, 'binary')
            if not os.path.exists(binary_path):
                skipped.append((vol, name, f'無 bone mask: {binary_path}'))
                continue
            cortical_path = level_file_path(mask_dir, level, 'cortical')
            row_any = row_l or row_r
            tt = _f(row_any.get('Total_Time'))
            tt = None if (tt is None or not np.isfinite(tt)) else tt
            tasks.append({
                'vol': vol, 'name': name, 'png': png, 'level': level, 'way': way,
                'd_l': d_l, 'l_l': l_l, 'd_r': d_r, 'l_r': l_r,
                'pos_l': _pos(row_l), 'pos_r': _pos(row_r),
                'binary': binary_path,
                'cortical': cortical_path if os.path.exists(cortical_path) else None,
                'swarm': int(_f(m['sw']) or 0), 'iter': int(_f(m['it']) or 0),
                'time': tt,
            })
    return tasks, skipped


def _run_batch(batch):
    """一個 worker 固定綁一個 GPU，依序渲染分到的圖（CUDA_VISIBLE_DEVICES
    須在 torch 首次 init CUDA 前設好，故綁定後不再變）。"""
    gpu, g_tasks = batch
    os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu)
    os.environ.setdefault('OMP_NUM_THREADS', '4')
    import matplotlib
    matplotlib.use('Agg')
    return [render_one(t) for t in g_tasks]


def render_one(task):
    import torch
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    img = sitk.ReadImage(task['binary'], sitk.sitkUInt8)
    spacing = list(img.GetSpacing())

    # 備份原圖（同目錄同檔名會重複跑時，備份檔保留第一份）
    backup_dir = os.path.join(os.path.dirname(os.path.dirname(task['png'])), BACKUP, task['vol'])
    bpath = os.path.join(backup_dir, task['name'])
    os.makedirs(backup_dir, exist_ok=True)
    if not os.path.exists(bpath):
        shutil.move(task['png'], bpath)

    try:
        out = render_bone_figure(
            task['vol'], task['level'], task['binary'], task['cortical'],
            base_folder=OUTPUT_BASE, spacing=spacing, way=task['way'],
            best_position_l=task['pos_l'], best_position_r=task['pos_r'],
            diameter_l=task['d_l'], length_l=task['l_l'],
            diameter_r=task['d_r'], length_r=task['l_r'],
            image2_path=None, device=device,
            swarm_size=task['swarm'], max_iter=task['iter'], total_time=task['time'],
            write_csv=False, output_path=task['png'])
    except Exception as e:
        if not os.path.exists(task['png']) and os.path.exists(bpath):
            shutil.move(bpath, task['png'])
        return (task['vol'], task['name'], 'fail', f'{type(e).__name__}: {e}')
    if out is None:
        if os.path.exists(bpath):
            shutil.move(bpath, task['png'])
        return (task['vol'], task['name'], 'fail', 'render 回傳 None（無/空遮罩）')
    if out != task['png']:
        os.replace(out, task['png'])
    return (task['vol'], task['name'], 'ok', f'{out} ({device})')


def main():
    ap = argparse.ArgumentParser(description='從 output.csv 重新渲染螺絲模式四視角圖')
    ap.add_argument('date', nargs='?', default='20260912')
    ap.add_argument('volume', nargs='?', default=None)
    ap.add_argument('--dry-run', action='store_true')
    ap.add_argument('--cpus', action='store_true', help='CPU 循序（不佔 GPU）')
    args = ap.parse_args()

    date_dir = os.path.join(OUTPUT_BASE, args.date)
    if not os.path.isdir(date_dir):
        raise SystemExit(f'no such date dir: {date_dir}')
    tasks, skipped = collect_figs(date_dir, args.volume)
    print(f'{args.date}: {len(tasks)} figure(s) to re-render, {len(skipped)} skip(s)')
    for vol, name, why in skipped:
        print(f'  [skip] {vol}/{name}: {why}')
    if args.dry_run:
        for t in tasks:
            print(f"  [dry] {t['vol']}/{t['name']}  "
                  f"L=({t['d_l']},{t['l_l']}) R=({t['d_r']},{t['l_r']})")
        return

    if not tasks:
        return

    if args.cpus or not _cuda_count():
        os.environ['CUDA_VISIBLE_DEVICES'] = ''
        for i, t in enumerate(tasks, 1):
            r = render_one(t)
            print(f'[{i:3d}/{len(tasks)}] {r[2]:4s} {r[0]}/{r[1]} {r[3]}', flush=True)
    else:
        n = min(_cuda_count(), 4)
        ctx = multiprocessing.get_context('spawn')
        batches = [(i, tasks[i::n]) for i in range(n)]
        with ctx.Pool(n) as pool:
            for results in pool.map(_run_batch, batches):
                for vol, name, status, detail in results:
                    print(f'[{status:4s}] {vol}/{name} {detail}', flush=True)
    print('=' * 60)
    print('Done.')


def _cuda_count():
    try:
        import torch
        return torch.cuda.device_count()
    except Exception:
        return 0


if __name__ == '__main__':
    main()