src/Wiener_Filter.py

import numpy as np
from scipy.fft import fft2, ifft2, fftfreq, fftshift, rfftn, irfftn
from scipy.fft import rfftn, irfftn
from numpy.random import gamma as sample_gamma



def _pad_to_shape(ir, shape):
    """Zero-pad IR đến cùng kích thước ảnh và circular shift tâm về (0,...)."""
    out = np.zeros(shape, dtype=ir.dtype)
    insert_slices = tuple(slice(0, s) for s in ir.shape)
    out[insert_slices] = ir
    for ax, k in enumerate(ir.shape):
        out = np.roll(out, - (k // 2), axis=ax)
    return out


def _laplacian_ir(ndim=2):
    """Sinh Laplacian kernel kích thước 3^N."""
    shape = (3,) * ndim
    ir = np.zeros(shape, dtype=float)
    center = (shape[0] // 2,) * ndim
    ir[center] = -2 * ndim
    for ax in range(ndim):
        idx_plus = list(center)
        idx_minus = list(center)
        idx_plus[ax] = 2
        idx_minus[ax] = 0
        ir[tuple(idx_plus)] = 1.0
        ir[tuple(idx_minus)] = 1.0
    return ir


def gaussian_psf(size=5, sigma=1.5, ndim=2):
    """Sinh Gaussian PSF chuẩn hóa (sum=1)."""
    # tạo grid
    coords = [np.arange(size) - size // 2 for _ in range(ndim)]
    grids = np.meshgrid(*coords, indexing='ij')
    dist2 = sum((g ** 2 for g in grids))
    psf = np.exp(-dist2 / (2 * sigma ** 2))
    psf /= psf.sum()
    return psf


def wiener_base_real(image, psf=None, balance=0.01, reg=None,
                     clip=True, eps=1e-12, psf_size=5, psf_sigma=1.5):
    """
    Wiener deconvolution cho ảnh thực

    Parameters
    ----------
    image : ndarray
        Ảnh đầu vào (real-valued).
    psf : ndarray or None
        Point Spread Function (impulse response). Nếu None → Gaussian.
    balance : float
        Hệ số điều chỉnh λ giữa fidelity và regularization.
    reg : ndarray or None
        Toán tử regularization (mặc định Laplacian).
    clip : bool
        Nếu True, giới hạn output về [-1,1].
    eps : float
        Giá trị nhỏ tránh chia cho 0.
    psf_size : int
        Kích thước kernel Gaussian mặc định.
    psf_sigma : float
        Sigma mặc định cho PSF Gaussian.

    Returns
    -------
    ndarray : ảnh khôi phục (cùng shape với input)
    """
    image = np.asarray(image, dtype=np.float64)
    ndim = image.ndim
    shape = image.shape

    # PSF mặc định = Gaussian kernel
    if psf is None:
        psf = gaussian_psf(psf_size, psf_sigma, ndim)

    # Regularization mặc định = Laplacian
    if reg is None:
        reg = _laplacian_ir(ndim)

    # Fourier transform
    H = rfftn(_pad_to_shape(psf, shape))
    D = rfftn(_pad_to_shape(reg, shape))
    Y = rfftn(image)

    # Wiener filter: conj(H) / (|H|^2 + λ|D|^2)
    denom = np.abs(H)**2 + balance * np.abs(D)**2
    denom = np.where(denom < eps, eps, denom)
    W = np.conj(H) / denom

    # Áp dụng và inverse FFT
    X = irfftn(W * Y, s=shape)

    if clip:
        X = np.clip(X, -1.0, 1.0)

    return X

def wiener_unsupervised_mcmc(
    image, 
    psf=None, 
    reg=None,
    n_iter=100,
    burn_in=20,
    clip=True, 
    eps=1e-12, 
    psf_size=15, 
    psf_sigma=5
):
    """
    Wiener deconvolution không giám sát sử dụng MCMC (Gibbs Sampler)
    để ước tính các siêu tham số.

    Hàm này thay thế việc truyền vào 'balance' cố định bằng cách
    tự động lấy mẫu các tham số precision (gamma_epsilon, gamma_1)
    dựa trên Mục 5.B của bài báo.

    Parameters
    ----------
    image : ndarray
        Ảnh đầu vào (real-valued).
    psf : ndarray or None
        Point Spread Function. Nếu None -> Gaussian.
    reg : ndarray or None
        Toán tử regularization (mặc định Laplacian).
    n_iter : int
        Tổng số lần lặp MCMC.
    burn_in : int
        Số lần lặp ban đầu để loại bỏ (giai đoạn 'cháy').
    clip : bool
        Nếu True, giới hạn output về [-1, 1].
    eps : float
        Giá trị nhỏ tránh chia cho 0.
    
    Returns
    -------
    ndarray : ảnh khôi phục (cùng shape với input)
    """
    
    # 1. CHUẨN HÓA ẢNH (QUAN TRỌNG NHẤT)
    # Đưa ảnh về khoảng [0, 1] để tránh residual quá lớn làm sập gamma_e
    img_min = image.min()
    img_max = image.max()
    if img_max - img_min > eps:
        image_norm = (image - img_min) / (img_max - img_min)
    else:
        image_norm = image.copy()

    image_norm = np.asarray(image_norm, dtype=np.float64)
    ndim = image_norm.ndim
    shape = image_norm.shape
    N = image_norm.size

    # --- Khởi tạo ---
    # PSF mặc định = Gaussian kernel
    if psf is None:
        psf = gaussian_psf(psf_size, psf_sigma, ndim)

    # Regularization mặc định = Laplacian [cite: 348]
    if reg is None:
        reg = _laplacian_ir(ndim)

    # Chuyển đổi sang miền Fourier
    H = rfftn(_pad_to_shape(psf, shape))
    D = rfftn(_pad_to_shape(reg, shape))
    Y = rfftn(image_norm)

    # Các thành phần tiền tính toán
    H_conj = np.conj(H)
    H_abs_sq = np.abs(H)**2
    D_abs_sq = np.abs(D)**2

    # Khởi tạo các siêu tham số (precision)
    gamma_epsilon = 1.0 / (np.var(image_norm) * 0.01 + eps) # Giả định nhiễu thấp ban đầu
    gamma_1 = 1.0 / np.var(image_norm)

    alpha0_eps, beta0_eps = 1e-2, 1e-2
    alpha0_1, beta0_1 = 1e-2, 1e-2
    GAMMA_MIN, GAMMA_MAX = 0.031, 1e2

    # Khởi tạo ảnh phục hồi (X_fft)
    # Bắt đầu bằng chính ảnh mờ
    X_fft = Y.copy()

    # Biến để lưu trữ tổng các mẫu (để lấy trung bình)
    X_samples_sum = np.zeros_like(Y, dtype=complex)
    samples_count = 0
    
    print(f"Bắt đầu MCMC với {n_iter} lần lặp (burn-in: {burn_in})...")

    # --- Vòng lặp MCMC (Gibbs Sampler) ---
    for k in range(n_iter):
        
        # 1. LẤY MẪU ẢNH (Mục 5.A)
        # ---------------------------
        # Tính toán bộ lọc Wiener-Hunt dựa trên các tham số gamma hiện tại
        # Đây chính là phần lõi của hàm wiener_base_real cũ
        
        # balance = gamma_1 / gamma_epsilon
        # W = H_conj / (H_abs_sq + balance * D_abs_sq)
        
        # Thay vì tính 'balance', ta dùng công thức gốc từ Mục 5.A:
        # mu = gamma_e * (gamma_e * |H|^2 + gamma_1 * |D|^2)^-1 * H_conj * Y
        # X_fft = mu 
        # (Đây là ước tính MAP/Mean, không phải mẫu đầy đủ, nhưng phổ biến)

        inv_sigma_fft = gamma_epsilon * H_abs_sq + gamma_1 * D_abs_sq
        sigma_fft = 1.0 / np.where(inv_sigma_fft < eps, eps, inv_sigma_fft)

        mu_fft = gamma_epsilon * sigma_fft * H_conj * Y

        # Thay vì X_fft = mu, ta cần X ~ N(mu, sigma)
        # Trong miền tần số: X = mu + sqrt(sigma) * noise
        # Noise phức chuẩn tắc
        noise_real = np.random.standard_normal(X_fft.shape)
        noise_imag = np.random.standard_normal(X_fft.shape)
        complex_noise = (noise_real + 1j * noise_imag)
        
        # Tính trung bình hậu nghiệm mu
        # X_fft = gamma_epsilon * sigma_fft * H_conj * Y
        X_fft = mu_fft + np.sqrt(sigma_fft) * complex_noise * 0.5 
        
        
        # 2. LẤY MẪU CÁC THAM SỐ PRECISION (Mục 5.B)
        # ---------------------------------------------
        # Sử dụng Jeffreys' prior (alpha=0, beta_inv=0)
        
        # A. Lấy mẫu gamma_epsilon (Noise precision)
        # alpha_eps = N / 2.0
        # residual_sq_sum = np.sum(np.abs(Y - H * X_fft)**2)
        # beta_eps = 2.0 / np.where(residual_sq_sum < eps, eps, residual_sq_sum)
        # gamma_epsilon = sample_gamma(alpha_eps, beta_eps)
        alpha_eps = alpha0_eps + N / 2.0
        residual_sq_sum = np.sum(np.abs(Y - H * X_fft)**2)
        beta_eps = beta0_eps + residual_sq_sum / 2.0
        gamma_epsilon = np.clip(sample_gamma(alpha_eps, 1.0/beta_eps), GAMMA_MIN, GAMMA_MAX)
        
        # B. Lấy mẫu gamma_1 (Image/Regularization precision)
        # alpha_1 = (N - 1) / 2.0 # (Bỏ qua 1 bậc tự do của gamma_0)
        # reg_sq_sum = np.sum(np.abs(D * X_fft)**2)
        # beta_1 = 2.0 / np.where(reg_sq_sum < eps, eps, reg_sq_sum)
        # gamma_1 = sample_gamma(alpha_1, beta_1)
        alpha_1 = alpha0_1 + (N - 1) / 2.0
        reg_sq_sum = np.sum(np.abs(D * X_fft)**2)
        beta_1 = beta0_1 + reg_sq_sum / 2.0
        gamma_1 = np.clip(sample_gamma(alpha_1, 1.0/beta_1), GAMMA_MIN, GAMMA_MAX)

        # 3. LƯU TRỮ MẪU (SAU KHI BURN-IN)
        # ---------------------------------
        if k >= burn_in:
            X_samples_sum += mu_fft
            samples_count += 1
            
        if (k + 1) % 10 == 0:
            balance = gamma_1 / (gamma_epsilon + eps)
            print(f" Iter {k+1}: ge={gamma_epsilon:.3f}, g1={gamma_1:.3f}, bal={balance:.4f}, Res={residual_sq_sum:.4f}")

    # --- HOÀN TẤT ---
    # Tính toán ước tính cuối cùng bằng trung bình các mẫu
    # Đây là ước tính trung bình hậu nghiệm (posterior mean)
    if samples_count == 0:
        print("Cảnh báo: Không có mẫu nào được thu thập (burn_in >= n_iter)")
        X_mean_fft = X_fft # Trả về mẫu cuối cùng
    else:
        X_mean_fft = X_samples_sum / samples_count
    
    X_final = irfftn(X_mean_fft, s=shape)
    
    # 4. KHÔI PHỤC SCALE GỐC (DENORMALIZE)
    if img_max - img_min > eps:
        X_final = X_final * (img_max - img_min) + img_min
        
    if clip:
        # Clip theo giá trị gốc của ảnh đầu vào (ví dụ 0-255) hoặc -1,1 tuỳ input
        # Ở đây clip theo min/max ban đầu của ảnh
        X_final = np.clip(X_final, image.min(), image.max())

    return X_final