import os
import numpy as np
import SimpleITK as sitk
from imaging.resample import resample_img
from imaging.affine import standardize_affine
from imaging.segmentation import seg_bone
import json
import glob
from config.constant import LABEL_MAP
from imaging.nifti_io import sitk_to_nibabel, nibabel_to_sitk
from imaging.orientation import anterior_y_side
from imaging.transforms import (build_volume_meta, img_geom, margined_box,
                                save_transform, std_flip_axes_for_direction)


def flip_y_sitk(img):
    """index 系 y 軸（array axis 1，前後方向）翻轉：
    只翻數據、spacing/origin/direction 等幾何不變，即整顆體積的前後
    朝向在 index 系翻轉（y 大側 <-> y 小側）。供 supine / prone 個案
    統一前後慣例（前側 = y 大側）用。"""
    arr = np.flip(sitk.GetArrayFromImage(img), axis=1)
    out = sitk.GetImageFromArray(arr)
    out.CopyInformation(img)
    return out



def load_progress(PROGRESS_FILE):
    if os.path.exists(PROGRESS_FILE):
        with open(PROGRESS_FILE, "r") as f:
            return json.load(f)
    return {}

def save_progress(progress, PROGRESS_FILE):
    with open(PROGRESS_FILE, "w") as f:
        json.dump(progress, f, indent=2)
        
def process_single_image(image_path, label_path, output_dir_base=None, max_z_spacing=None, allowed_levels=None, min_levels=None, metadata_cache=None, crop_subdir=False):
    """metadata_cache：可選，需提供 get(name) -> dict|None 與
    put(name, dict)。dict 可含 spacing=[x,y,z]、labels=[label id]。
    兩者都在 db 裡時整支跳過判定不需讀影像 / label 檔。
    crop_subdir=True：各 level 的未旋轉輸出（_binary / _smd /
    _smd_resampled / _binary_sdf / _binary_nn / _roi）改写到
    <vol>/crop/ 子資料夾（False＝舊佈局、<vol>/ 頂層）。"""

    file_name = os.path.basename(image_path)
    name = file_name.replace(".nii.gz", "")

    # pixel spacing / labels 優先用 metadata db，避免每次 run 都讀檔
    meta = metadata_cache.get(name) if metadata_cache is not None else None
    image = None
    label = None
    spacing = meta.get("spacing") if meta is not None else None
    existing_labels = meta.get("labels") if meta is not None else None

    if spacing is None:
        image = sitk.ReadImage(image_path)
        spacing = [float(v) for v in image.GetSpacing()]
        if metadata_cache is not None:
            metadata_cache.put(name, {"spacing": spacing})

    # z spacing 過大（低解析度掃描）跳過整支 pipeline
    if max_z_spacing is not None and spacing[2] > max_z_spacing:
        z_spacing = spacing[2]
        print(f"z spacing {z_spacing:.2f} mm > {max_z_spacing} mm, "
              f"skipping pipeline for {name}")
        return {
            "processed_labels": [],
            "missing_labels": [],
            "skipped": True,
            "skip_reason": f"z spacing {z_spacing:.2f} mm > {max_z_spacing} mm"
        }

    if existing_labels is None:
        label = sitk.ReadImage(label_path)
        # 取得現有 label
        lssif = sitk.LabelShapeStatisticsImageFilter()
        lssif.Execute(label)
        existing_labels = [int(v) for v in lssif.GetLabels()]  # 例如 [1,2,3,20,21]
        if metadata_cache is not None:
            metadata_cache.put(name, {"labels": existing_labels})
    else:
        print(f"Metadata db hit for {name} (spacing={spacing}, labels={existing_labels})")

    print(f"Existing labels in {os.path.basename(label_path)}: {existing_labels}")

    allowed_label_list = [n for n in existing_labels
                          if n in LABEL_MAP and (allowed_levels is None
                                                 or LABEL_MAP[n] in allowed_levels)]

    # 沒有任何符合 allowed levels（如 lumbar）的 label：不建立輸出資料夾、
    # 不做重取樣，整檔跳過（不計入 max_images）
    if not allowed_label_list:
        print(f"No label matches allowed levels {allowed_levels} in {name}, "
              f"skipping (no output folder)")
        return {
            "processed_labels": [],
            "missing_labels": [],
            "skipped": True,
            "skip_reason": f"no label in allowed levels {allowed_levels}"
        }

    # 符合 allowed levels 的 label 少於 min_levels（如 lumbar < 2 層）：
    # 整支 pipeline 跳過，不建立輸出資料夾（不計入 max_images）
    if min_levels is not None and len(allowed_label_list) < min_levels:
        have = ", ".join(f"{n} ({LABEL_MAP[n]})" for n in allowed_label_list)
        print(f"Only {len(allowed_label_list)} allowed-level label(s) [{have}] in {name} "
              f"< {min_levels}, skipping pipeline (no output folder)")
        return {
            "processed_labels": [],
            "missing_labels": [],
            "skipped": True,
            "skip_reason": (f"only {len(allowed_label_list)} label(s) in allowed "
                            f"levels < {min_levels}")
        }

    # 進入 pipeline：若上面的 spacing / labels 來自 metadata db，
    # 影像與 label 檔還沒讀，這裡補讀
    if image is None:
        image = sitk.ReadImage(image_path)
    if label is None:
        label = sitk.ReadImage(label_path)

    # LabelStatisticsImageFilter computes statistics (e.g., mean, minimum, maximum, median) of pixel values in an image, segmented by labels in a corresponding label image.
    lsif = sitk.LabelStatisticsImageFilter()
    lsif.Execute(image, label)

    # Assume to have some sitk image (itk_image) and label (itk_label)
    resampled_sitk_img = resample_img(image, out_spacing=[0.5, 0.5, 0.5], is_label=False)
    resampled_sitk_lbl = resample_img(label, out_spacing=[0.5, 0.5, 0.5], is_label=True)

    # 前後（AP）方向判定：本流程最終輸出慣例是前側 = y 大側（後 = y 小側）。
    # 但 prone（伏位）掃描個案經 standardize_affine 後前側落在 y 小側
    # （例：CTSpine1K colon 0003、0075、0460...，全 dataset 約 1%），
    # 不修正時上終板 / 棘突 / 椎體分割與 rotated/ 對齊全部反掉。
    #
    # 判定（對每個 allowed level 的 0.5mm label 個別做、再票決——不能
    # 直接 union 多層：腰椎前凸（lordosis）下各層椎體在 AP 投影會散開，
    # 椎管「空隙」被其它層的骨填掉）：
#   1) anterior_y_side 回傳工作體積（0.5mm 重取樣、standardize_affine
    #      之前）grid 中椎體塊所在的 y 端（y_min / y_max / None）；
    #   2) standardize_affine（nibabel 端）會在輸出 affine y 分量 < 0 時
    #      再翻一次 y。注意 nibabel affine 與 SimpleITK direction 的 y 分
    #      量符號相反（NIfTI RAS <-> SITK LPS），所以
    #      standardize_affine 會翻 y  <=>  direction[4] > 0；
    #      最終 y 端 = pre_side（若會翻 y 則 y_min<->y_max 互換）；
    #   3) 最終前側會落 y 小側時，現在先對 CT 與 label 做 y 翻轉
    #      （純 index 翻轉、幾何不變），與 standardize_affine 的翻轉
    #      組成淨效果，使所有輸出與其它個案同方向。
    # 無法判定 / 投票平手時維持原方向（寧可不翻、不誤翻）。
    # 判定結果存入 metadata db（ap_flip），重跑免重算。
    ap_flip = (meta or {}).get("ap_flip")
    if ap_flip is None:
        # standardize_affine 是否會翻轉 y（nibabel affine[1,1] < 0，
        # 等价於 sitk direction[4] > 0，兩者符號相反）
        std_flips_y = resampled_sitk_img.GetDirection()[4] > 0
        arr_lbl = sitk.GetArrayFromImage(resampled_sitk_lbl)
        votes = []
        for n in allowed_label_list:
            side = anterior_y_side(arr_lbl == n)
            if side is not None:
                votes.append(side)
        n_min = votes.count("y_min")
        n_max = votes.count("y_max")
        if n_min > n_max:
            pre_side = "y_min"
        elif n_max > n_min:
            pre_side = "y_max"
        else:
            pre_side = None
        if pre_side is None:
            ap_flip = False
            print(f"AP orientation undetermined for {name} (votes={votes}); "
                  f"proceeding with default orientation (anterior = large y)")
        else:
            final_side = (pre_side if not std_flips_y
                          else ("y_min" if pre_side == "y_max" else "y_max"))
            ap_flip = final_side == "y_min"
            print(f"AP orientation for {name}: pre={pre_side} "
                  f"(std_flips_y={std_flips_y}) -> final={final_side} "
                  f"[votes={votes}], flip={ap_flip}")
            if metadata_cache is not None:
                metadata_cache.put(name, {"ap_flip": bool(ap_flip)})
    if ap_flip:
        # label（原解析度 raw label）也要翻：seg_bone 的主遮罩鏈
        #（_binary / SMD / _binary_sdf）是用 original_label（= label）
        # 算的，不是用 0.5mm resampled label；漏翻時翻轉不生效。
        label = flip_y_sitk(label)
        resampled_sitk_img = flip_y_sitk(resampled_sitk_img)
        resampled_sitk_lbl = flip_y_sitk(resampled_sitk_lbl)
        print(f"AP orientation corrected for {name}: CT and label flipped "
              f"along y; outputs unified to anterior = large y")

    # 建立每個檔案的輸出資料夾
    file_name = os.path.basename(image_path)
    name = file_name.replace(".nii.gz", "")
    output_dir = os.path.join(output_dir_base, name)
    os.makedirs(output_dir, exist_ok=True)
    # 各 level 的未旋轉輸出寫入位置：crop_subdir=True 時在 <vol>/crop/
    # 子資料夾（旋轉版在 <vol>/rotated/，xfr_preprocess 產出）
    if crop_subdir:
        seg_dir = os.path.join(output_dir, 'crop')
        os.makedirs(seg_dir, exist_ok=True)
    else:
        seg_dir = output_dir
    
    # 存現有 label 到 txt
    txt_path = os.path.join(output_dir, f"{name}_labels.txt")
    with open(txt_path, "w") as f:
        for lab in existing_labels:
            f.write(f"{lab}\t{LABEL_MAP.get(lab, 'Unknown')}\n")
    
    # 遍歷現有 label 做分割；個別 label 失敗（label_map 缺該 label、
    # 或 seg_bone 報錯）只跳過該 label，不中斷整檔處理
    processed = []
    skipped = []
    level_entries = {}      # level name -> transform.json 的 level 條目
    std_flips = None        # standardize_affine 實際翻的軸（整卷一致；記錄用）
    for n in existing_labels:
        if n not in LABEL_MAP:
            print(f"Label {n} not found in label_map, skipping this label (file continues).")
            skipped.append(n)
            continue
        if allowed_levels is not None and LABEL_MAP[n] not in allowed_levels:
            print(f"Label {n} ({LABEL_MAP[n]}) not in allowed levels, "
                  f"skipping this label (file continues).")
            skipped.append(n)
            continue
        try:
            res = seg_bone(n, name, resampled_sitk_img, resampled_sitk_lbl, seg_dir,
                           label_map=LABEL_MAP, original_label=label)
            if res is None:
                print(f"Label {n}: empty after largest-CC extraction, skipping this label.")
                skipped.append(n)
                continue
            (roi_path, binary_path, roi2_path, cortical_path, binary_nn_path,
             binary_linear_path, smd_path, resampled_path, binary_sdf_path,
             binary_erode_path, seg_meta) = res
            for path in [roi_path, binary_path, roi2_path, cortical_path,
                         binary_nn_path, binary_linear_path,
                         smd_path, resampled_path, binary_sdf_path,
                         binary_erode_path]:
                if path is not None:
                    flips = standardize_affine(path, seg_dir)
                    if flips:
                        if std_flips is None:
                            std_flips = sorted(flips)
                        elif sorted(flips) != std_flips:
                            print(f"WARNING: standardize_affine 翻軸不一致 "
                                  f"({name}): {std_flips} vs {sorted(flips)}")
            # 該 level 的裁切 box（x,y,z 序；imaging/transforms.py 定義）
            boxes = {}
            if seg_meta['has_smd']:
                boxes['smd_resampled'] = seg_meta['bbox2']
                boxes['binary_sdf'] = seg_meta['bbox2']
                boxes['roi'] = seg_meta['bbox2']
            else:
                boxes['roi'] = seg_meta['bbox2']
            if seg_meta['nn_bbox'] is not None:
                boxes['binary_nn'] = seg_meta['nn_bbox']
            if seg_meta['bbox_orig'] is not None:
                boxes['binary'] = seg_meta['bbox_orig']
                boxes['smd'] = margined_box(seg_meta['bbox_orig'],
                                            image.GetSize(), margin=4)
            level_entries[LABEL_MAP[n]] = {
                'label': int(n),
                'std_flip_axes': list(std_flips) if std_flips is not None else [],
                'boxes': boxes,
            }
            processed.append(n)
        except RuntimeError as e:
            print(f"Label {n} could not be processed, skipping. Error: {e}")
            skipped.append(n)

    # transform.json：每 level 記錄 原始 CT <-> 標準化 grid 的完整座標鏈
    # （box / flip 軸 / ap_flip / 原 CT 幾何），rotated/ 的 R/center/start 由
    # _write_rotated_level 補寫（xfr_preprocess）；供分割後 mask 映回原始座標
    if level_entries:
        try:
            tmeta = build_volume_meta(name, image_path, img_geom(image),
                                      img_geom(resampled_sitk_img), ap_flip,
                                      level_entries)
            save_transform(output_dir, tmeta)
            print(f"Transform metadata: {os.path.join(output_dir, 'transform.json')} "
                  f"({len(level_entries)} level(s), std_flips={tmeta['levels'][list(level_entries)[0]]['std_flip_axes']})")
            try:
                expected = std_flip_axes_for_direction(image.GetDirection())
                if std_flips is not None and list(std_flips) != expected:
                    print(f"WARNING: {name} std_flips {std_flips} != 預期 "
                          f"{expected}（依原 CT direction 推斷）")
                elif std_flips is None and expected:
                    print(f"WARNING: {name} 預期 std_flips {expected} 但未翻任何軸")
            except ValueError:
                print(f"WARNING: {name} 原 CT direction 非對角 ±1；"
                      f"transform.json 的反向映射將無法使用該卷")
        except Exception as e:
            print(f"WARNING: 寫 transform.json 失敗（{name}）：{e}")

    return {
        "processed_labels": processed,
        "missing_labels": skipped
    }

def process_dataset(image_dir, label_dir, output_dir, labels_to_process=None, max_images=None, post_process=None, max_z_spacing=None, allowed_levels=None, min_levels=None, metadata_cache=None, crop_subdir=False):
    image_files = sorted(glob.glob(os.path.join(image_dir, "*.nii.gz")))
    total_files = len(image_files)
    print(f"Total files: {total_files}")

    PROGRESS_FILE = os.path.join(output_dir, "progress.json")

    progress = load_progress(PROGRESS_FILE)
    all_file_summary = []
    # max_images 統計「實際進入 pipeline 的檔數」：
    # z spacing 超標被跳過的檔不計入配額，繼續掃描後續檔案
    processed_count = 0

    for idx, image_path in enumerate(image_files, 1):
        if max_images is not None and processed_count >= max_images:
            print(f"Reached max_images={max_images}, stopping early.")
            break

        file_name = os.path.basename(image_path)
        name = file_name.replace(".nii.gz", "")
        label_path = os.path.join(label_dir, file_name.replace(".nii.gz", "_seg.nii.gz"))

        file_summary = {
            "file_name": file_name,
            "current_labels": [],
            "missing_labels": []
        }

        if progress.get(name, {}).get("finished", False):
            print(f"[{idx}/{total_files}] Already finished: {file_name}")
            file_summary["current_labels"] = progress[name].get("processed_labels", [])
            file_summary["missing_labels"] = progress[name].get("missing_labels", [])
            processed_count += 1
            all_file_summary.append(file_summary)
            continue

        if not os.path.exists(label_path):
            print(f"[{idx}/{total_files}] Warning: label not found for {file_name}")
            file_summary["note"] = "Label file not found"
            processed_count += 1
            all_file_summary.append(file_summary)
            continue

        try:
            result = process_single_image(image_path, label_path, output_dir_base=output_dir,
                                           max_z_spacing=max_z_spacing,
                                           allowed_levels=allowed_levels,
                                           min_levels=min_levels,
                                           metadata_cache=metadata_cache,
                                           crop_subdir=crop_subdir)
            # print(result)
            # exit()
        except Exception as e:
            print(f"[{idx}/{total_files}] Error processing {file_name}: {e}")
            file_summary["note"] = f"Error: {e}"
            processed_count += 1
            all_file_summary.append(file_summary)
            continue

        if result.get("skipped"):
            # z spacing 超標：不計入 max_images 配額，繼續掃描後續檔案
            print(f"[{idx}/{total_files}] Skipped ({result['skip_reason']}): {file_name}")
            file_summary["note"] = result["skip_reason"]
            all_file_summary.append(file_summary)
            continue

        processed_count += 1
        file_summary["current_labels"] = result["processed_labels"]
        file_summary["missing_labels"] = result["missing_labels"]
        all_file_summary.append(file_summary)

        progress[name] = {
            "finished": True,
            "processed_labels": result["processed_labels"],
            "missing_labels": result["missing_labels"]
        }
        save_progress(progress, PROGRESS_FILE)

        count_msg = f" | processed {processed_count}/{max_images}" if max_images is not None else ""
        print(f"[{idx}/{total_files}] Finished: {file_name} | "
              f"Missing labels: {result['missing_labels'] or 'None'}{count_msg}")

        if post_process is not None:
            try:
                post_process(os.path.join(output_dir, name), result["processed_labels"])
            except Exception as e:
                print(f"[{idx}/{total_files}] post_process error for {name}: {e}")

    # --- Summary ---
    summary_path = os.path.join(output_dir, "all_files_label_summary.txt")
    os.makedirs(os.path.dirname(summary_path), exist_ok=True)
    print(f"\nWriting summary to {summary_path}...")

    with open(summary_path, "w") as f:
        f.write("--- CTSpine1K Dataset Label Summary ---\n")
        f.write(f"Total files processed: {total_files}\n\n")

        for summary in all_file_summary:
            f.write("================================================\n")
            f.write(f"File: {summary['file_name']}\n")

            processed_labels_str = ", ".join([str(l) for l in summary['current_labels']])
            f.write(f"Labels processed: {processed_labels_str}\n")

            if summary['missing_labels']:
                missing_str = ", ".join([f"{l} ({LABEL_MAP.get(l, 'Unknown')})" for l in summary['missing_labels']])
                f.write(f"🚨 Missing Labels: {missing_str}\n")
            else:
                f.write("✅ Missing Labels: None\n")

            if "note" in summary:
                f.write(f"Note: {summary['note']}\n")
            f.write("================================================\n\n")

    print("All done! Summary file created.")