#!/usr/bin/env python3
# compare_local_contrast.py

import os
import argparse
import numpy as np
import matplotlib.pyplot as plt

# Required: tifffile (TIFF読み込みの安定版)
try:
    import tifffile as tiff
except ImportError as e:
    raise SystemExit("tifffile が必要です。pip install tifffile で入れてください。") from e

# Optional: SciPy (高速な局所平均/分散)
try:
    from scipy.ndimage import uniform_filter
    HAVE_SCIPY = True
except ImportError:
    HAVE_SCIPY = False

# Optional: Pillow (PNG/JPG等の読み込み)
try:
    from PIL import Image
    HAVE_PIL = True
except ImportError:
    HAVE_PIL = False

# Optional: FITS
try:
    from astropy.io import fits
    HAVE_ASTROPY = True
except ImportError:
    HAVE_ASTROPY = False

# Optional: 距離指標（JS divergence）
def js_divergence(p, q, eps=1e-12):
    p = np.asarray(p, dtype=np.float64)
    q = np.asarray(q, dtype=np.float64)
    p = p / (p.sum() + eps)
    q = q / (q.sum() + eps)
    m = 0.5 * (p + q)
    # KL with eps
    def kl(a, b):
        a = np.clip(a, eps, None)
        b = np.clip(b, eps, None)
        return np.sum(a * np.log(a / b))
    return 0.5 * kl(p, m) + 0.5 * kl(q, m)

def ks_statistic(x, y):
    """Two-sample KS statistic (no p-value)"""
    x = np.sort(np.asarray(x))
    y = np.sort(np.asarray(y))
    n = x.size
    m = y.size
    # Merge points
    allv = np.sort(np.concatenate([x, y]))
    cx = np.searchsorted(x, allv, side="right") / n
    cy = np.searchsorted(y, allv, side="right") / m
    return np.max(np.abs(cx - cy))

def read_image(path):
    ext = os.path.splitext(path)[1].lower()
    if ext in [".tif", ".tiff"]:
        img = tiff.imread(path)
        return img
    if ext in [".fits", ".fit", ".fts"]:
        if not HAVE_ASTROPY:
            raise RuntimeError("FITSを読むには astropy が必要です。pip install astropy")
        return fits.getdata(path)
    if ext in [".png", ".jpg", ".jpeg"]:
        if not HAVE_PIL:
            raise RuntimeError("PNG/JPGを読むには pillow が必要です。pip install pillow")
        return np.array(Image.open(path))
    raise RuntimeError(f"未対応の拡張子です: {ext}")

def to_luminance(img):
    """RGBなら輝度へ。単一チャンネルならそのまま。"""
    a = np.asarray(img)
    if a.ndim == 2:
        return a
    if a.ndim == 3 and a.shape[2] >= 3:
        # sRGBっぽい係数（ざっくり）
        r, g, b = a[..., 0], a[..., 1], a[..., 2]
        return 0.2126 * r + 0.7152 * g + 0.0722 * b
    # それ以外は平均で潰す
    return a.mean(axis=-1)

def robust_normalize(x, low_pct=0.1, high_pct=99.9):
    x = np.asarray(x, dtype=np.float32)
    lo, hi = np.percentile(x, [low_pct, high_pct])
    if hi <= lo:
        return np.zeros_like(x, dtype=np.float32)
    y = np.clip(x, lo, hi)
    y = (y - lo) / (hi - lo)
    return y.astype(np.float32)

def local_mean_std(img01, window):
    """img01: 0..1 float32"""
    if HAVE_SCIPY:
        # uniform_filter は O(1) で速い
        mu = uniform_filter(img01, size=window, mode="reflect")
        mu2 = uniform_filter(img01 * img01, size=window, mode="reflect")
        var = np.maximum(mu2 - mu * mu, 0.0)
        sigma = np.sqrt(var)
        return mu, sigma
    # SciPyなし版（遅いけど動く）：積分画像で実装
    # 端処理は簡易に reflect 近似（pad）
    pad = window // 2
    x = np.pad(img01, pad, mode="reflect")
    # integral image
    ii = x.cumsum(0).cumsum(1)
    ii2 = (x * x).cumsum(0).cumsum(1)

    def box_sum(ii_):
        w = window
        A = ii_[w:, w:]
        B = ii_[:-w, w:]
        C = ii_[w:, :-w]
        D = ii_[:-w, :-w]
        return A - B - C + D

    s = box_sum(ii)
    s2 = box_sum(ii2)
    area = float(window * window)
    mu = s / area
    var = np.maximum(s2 / area - mu * mu, 0.0)
    sigma = np.sqrt(var)
    return mu.astype(np.float32), sigma.astype(np.float32)

def compute_local_contrast(img01, window, eps=1e-6):
    mu, sigma = local_mean_std(img01, window)
    # あなたの定義: local_std / local_mean
    lc = sigma / np.maximum(mu, eps)
    return lc.astype(np.float32)

def summarize(arr):
    arr = np.asarray(arr)
    return {
        "mean": float(np.mean(arr)),
        "median": float(np.median(arr)),
        "p90": float(np.percentile(arr, 90)),
        "p95": float(np.percentile(arr, 95)),
        "p99": float(np.percentile(arr, 99)),
        "max": float(np.max(arr)),
    }

def savefig(path):
    plt.tight_layout()
    plt.savefig(path, dpi=160)
    plt.close()

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--images", nargs="+", required=True, help="画像パスを並べて指定")
    ap.add_argument("--labels", nargs="*", default=None, help="画像ラベル（省略可）")
    ap.add_argument("--window", type=int, default=64, help="ローカル統計の窓サイズ(px)")
    ap.add_argument("--bins", type=int, default=300, help="ヒストグラムbin数")
    ap.add_argument("--crop", type=str, default=None,
                    help="x0,y0,w,h 例: 1000,2000,3000,2000 (比較を同一領域に固定したい場合)")
    ap.add_argument("--ref", type=int, default=0, help="差分の基準画像インデックス（0始まり）")
    ap.add_argument("--outdir", type=str, default="contrast_report", help="出力フォルダ")
    args = ap.parse_args()

    os.makedirs(args.outdir, exist_ok=True)

    paths = args.images
    labels = args.labels
    if labels is None or len(labels) == 0:
        labels = [os.path.splitext(os.path.basename(p))[0] for p in paths]
    if len(labels) != len(paths):
        raise SystemExit("labels の数は images と同じにしてください。")

    crop = None
    if args.crop:
        x0, y0, w, h = map(int, args.crop.split(","))
        crop = (x0, y0, w, h)

    imgs01 = []
    for p in paths:
        raw = read_image(p)
        lum = to_luminance(raw)
        if crop:
            x0, y0, w, h = crop
            lum = lum[y0:y0+h, x0:x0+w]
        img01 = robust_normalize(lum)
        imgs01.append(img01)

    # local contrast arrays/maps
    lcs = []
    summaries = []
    for img01 in imgs01:
        lc = compute_local_contrast(img01, window=args.window)
        lcs.append(lc)
        summaries.append(summarize(lc.ravel()))

    # ===== 1) Overlay histogram (same bins) =====
    all_vals = np.concatenate([lc.ravel() for lc in lcs])
    # 外れ値で潰れないように上側を少し切る（必要なら調整）
    vmax = np.percentile(all_vals, 99.9)
    bins = np.linspace(0.0, vmax, args.bins)

    plt.figure()
    for lc, lab in zip(lcs, labels):
        hist, _ = np.histogram(np.clip(lc.ravel(), 0, vmax), bins=bins)
        plt.plot((bins[:-1] + bins[1:]) / 2, hist + 1, label=lab)  # +1 for log
    plt.yscale("log")
    plt.title(f"Local contrast distribution (window={args.window}px) overlay")
    plt.xlabel("local_std / local_mean (clipped)")
    plt.ylabel("Count (log)")
    plt.legend()
    savefig(os.path.join(args.outdir, "01_hist_overlay.png"))

    # ===== 2) ECDF (差が見えやすい) =====
    plt.figure()
    for lc, lab in zip(lcs, labels):
        v = np.clip(lc.ravel(), 0, vmax)
        v = np.sort(v)
        y = np.linspace(0, 1, v.size, endpoint=False)
        plt.plot(v, y, label=lab)
    plt.title(f"ECDF of local contrast (window={args.window}px)")
    plt.xlabel("local_std / local_mean (clipped)")
    plt.ylabel("Cumulative probability")
    plt.legend()
    savefig(os.path.join(args.outdir, "02_ecdf.png"))

    # ===== 3) Summary table (text) =====
    txt_path = os.path.join(args.outdir, "03_summary.txt")
    with open(txt_path, "w", encoding="utf-8") as f:
        f.write(f"window={args.window}px\n")
        f.write("label\tmean\tmedian\tp90\tp95\tp99\tmax\n")
        for lab, s in zip(labels, summaries):
            f.write(f"{lab}\t{s['mean']:.6f}\t{s['median']:.6f}\t{s['p90']:.6f}\t"
                    f"{s['p95']:.6f}\t{s['p99']:.6f}\t{s['max']:.6f}\n")

    # ===== 4) Distance matrices: KS / JS =====
    n = len(lcs)
    ks = np.zeros((n, n), dtype=np.float64)
    js = np.zeros((n, n), dtype=np.float64)

    # JSはヒストグラムで近似
    hists = []
    for lc in lcs:
        hist, _ = np.histogram(np.clip(lc.ravel(), 0, vmax), bins=bins)
        hists.append(hist.astype(np.float64))

    for i in range(n):
        for j in range(n):
            if i == j:
                continue
            ks[i, j] = ks_statistic(np.clip(lcs[i].ravel(), 0, vmax),
                                    np.clip(lcs[j].ravel(), 0, vmax))
            js[i, j] = js_divergence(hists[i], hists[j])

    def plot_matrix(M, title, fname):
        plt.figure()
        plt.imshow(M, interpolation="nearest")
        plt.title(title)
        plt.colorbar()
        plt.xticks(range(n), labels, rotation=45, ha="right")
        plt.yticks(range(n), labels)
        savefig(os.path.join(args.outdir, fname))

    plot_matrix(ks, "KS distance (local contrast distributions)", "04_ks_matrix.png")
    plot_matrix(js, "JS divergence (hist approx)", "05_js_matrix.png")

    # ===== 5) Local contrast maps montage (same color scale) =====
    # ここも比較が見えるように同一スケールに固定
    map_vmax = np.percentile(np.concatenate([lc.ravel() for lc in lcs]), 99.5)
    cols = min(3, n)
    rows = int(np.ceil(n / cols))
    plt.figure(figsize=(5 * cols, 4 * rows))
    for idx, (lc, lab) in enumerate(zip(lcs, labels)):
        ax = plt.subplot(rows, cols, idx + 1)
        im = ax.imshow(np.clip(lc, 0, map_vmax), vmin=0, vmax=map_vmax)
        ax.set_title(lab)
        ax.set_axis_off()
    plt.colorbar(im, ax=plt.gcf().axes, fraction=0.015, pad=0.01)
    plt.suptitle("Local contrast maps (same scale)", y=1.02)
    savefig(os.path.join(args.outdir, "06_lc_maps_montage.png"))

    # ===== 6) Difference maps vs reference =====
    ref = args.ref
    ref_lc = lcs[ref]
    # 差分は見やすいように対称レンジに
    diffs = []
    for i in range(n):
        if i == ref:
            diffs.append(None)
        else:
            diffs.append(lcs[i] - ref_lc)

    # 差分の表示スケール
    diff_vals = np.concatenate([d.ravel() for d in diffs if d is not None])
    dlim = np.percentile(np.abs(diff_vals), 99.5)

    m = n - 1
    cols = min(3, m) if m > 0 else 1
    rows = int(np.ceil(m / cols)) if m > 0 else 1
    plt.figure(figsize=(5 * cols, 4 * rows))
    k = 1
    for i in range(n):
        if i == ref:
            continue
        ax = plt.subplot(rows, cols, k)
        k += 1
        im = ax.imshow(np.clip(diffs[i], -dlim, dlim), vmin=-dlim, vmax=dlim)
        ax.set_title(f"{labels[i]} - {labels[ref]}")
        ax.set_axis_off()
    if m > 0:
        plt.colorbar(im, ax=plt.gcf().axes, fraction=0.015, pad=0.01)
        plt.suptitle("Local contrast difference maps (vs reference)", y=1.02)
        savefig(os.path.join(args.outdir, "07_lc_diff_vs_ref.png"))

    print("Done. Outputs saved to:", args.outdir)
    print("Summary:", txt_path)

if __name__ == "__main__":
    main()