"""
Script doc lap de test model phan doan (segmentation) khoi u nao
(DeepLabV3+ ResNet101V2) da train (.keras).
Khong can toan bo pipeline training - chi can model + 1 file de test.

Input: file .h5 co dataset 'image' (va 'mask' neu muon so sanh voi ground truth).

Cach dung:
    # Chi du doan, luu anh overlay ra file
    python test_segmentation.py --model best_model.keras --input slice.h5 --output result.png

    # Du doan + tinh metric neu file h5 co san mask ground truth
    python test_segmentation.py --model best_model.keras --input slice.h5 --output result.png
"""

import argparse
import numpy as np
import cv2
import h5py
import tensorflow as tf
import matplotlib.pyplot as plt
from scipy.ndimage import label, binary_closing, binary_dilation

IMG_SIZE = (256, 256)
NUM_SLICES = 1
PIXEL_THRESHOLD = 0.01


# ============================================================
# LOSS / METRIC FUNCTIONS - BAT BUOC PHAI CO DE load_model() HOAT DONG
# (copy y nguyen tu script training goc)
# ============================================================
@tf.function
def dice_coef(y_true, y_pred, smooth=1e-6):
    y_true = tf.cast(y_true, tf.float32)
    y_pred = tf.cast(y_pred, tf.float32)
    if len(y_true.shape) == 3:
        y_true = tf.expand_dims(y_true, axis=-1)
    if len(y_pred.shape) == 3:
        y_pred = tf.expand_dims(y_pred, axis=-1)
    y_true_f = tf.reshape(y_true, [-1])
    y_pred_f = tf.reshape(y_pred, [-1])
    intersection = tf.reduce_sum(y_true_f * y_pred_f)
    union = tf.reduce_sum(y_true_f) + tf.reduce_sum(y_pred_f)
    return (2.0 * intersection + smooth) / (union + smooth)


@tf.function
def generalized_dice_coef(y_true, y_pred, smooth=1e-6):
    y_true = tf.cast(y_true, tf.float32)
    y_pred = tf.cast(y_pred, tf.float32)
    if len(y_true.shape) == 3:
        y_true = tf.expand_dims(y_true, axis=-1)
    if len(y_pred.shape) == 3:
        y_pred = tf.expand_dims(y_pred, axis=-1)
    w_fg = 1.0 / (tf.reduce_sum(y_true) ** 2 + smooth)
    w_bg = 1.0 / (tf.reduce_sum(1.0 - y_true) ** 2 + smooth)
    intersection_fg = tf.reduce_sum(y_true * y_pred)
    union_fg = tf.reduce_sum(y_true) + tf.reduce_sum(y_pred)
    intersection_bg = tf.reduce_sum((1.0 - y_true) * (1.0 - y_pred))
    union_bg = tf.reduce_sum(1.0 - y_true) + tf.reduce_sum(1.0 - y_pred)
    numerator = w_fg * intersection_fg + w_bg * intersection_bg
    denominator = w_fg * union_fg + w_bg * union_bg
    return 2.0 * (numerator + smooth) / (denominator + smooth)


@tf.function
def weighted_dice_coef(y_true, y_pred, weight_fg=0.75, weight_bg=0.25, smooth=1e-6):
    y_true = tf.cast(y_true, tf.float32)
    y_pred = tf.cast(y_pred, tf.float32)
    if len(y_true.shape) == 3:
        y_true = tf.expand_dims(y_true, axis=-1)
    if len(y_pred.shape) == 3:
        y_pred = tf.expand_dims(y_pred, axis=-1)
    intersection_fg = tf.reduce_sum(y_true * y_pred)
    union_fg = tf.reduce_sum(y_true) + tf.reduce_sum(y_pred)
    dice_fg = (2.0 * intersection_fg + smooth) / (union_fg + smooth)
    intersection_bg = tf.reduce_sum((1.0 - y_true) * (1.0 - y_pred))
    union_bg = tf.reduce_sum(1.0 - y_true) + tf.reduce_sum(1.0 - y_pred)
    dice_bg = (2.0 * intersection_bg + smooth) / (union_bg + smooth)
    return weight_fg * dice_fg + weight_bg * dice_bg


@tf.function
def iou(y_true, y_pred, smooth=1e-6):
    y_true = tf.cast(y_true, tf.float32)
    y_pred = tf.cast(y_pred, tf.float32)
    if len(y_true.shape) == 3:
        y_true = tf.expand_dims(y_true, axis=-1)
    if len(y_pred.shape) == 3:
        y_pred = tf.expand_dims(y_pred, axis=-1)
    y_true_f = tf.reshape(y_true, [-1])
    y_pred_f = tf.reshape(y_pred, [-1])
    intersection = tf.reduce_sum(y_true_f * y_pred_f)
    union = tf.reduce_sum(y_true_f) + tf.reduce_sum(y_pred_f) - intersection
    return (intersection + smooth) / (union + smooth)


@tf.function
def boundary_iou(y_true, y_pred, smooth=1e-6):
    y_true = tf.cast(y_true, tf.float32)
    y_pred = tf.cast(y_pred, tf.float32)
    if len(y_true.shape) == 3:
        y_true = tf.expand_dims(y_true, axis=-1)
    if len(y_pred.shape) == 3:
        y_pred = tf.expand_dims(y_pred, axis=-1)

    def erosion(x):
        return -tf.nn.max_pool2d(-x, ksize=3, strides=1, padding='SAME')

    y_true_eroded = erosion(y_true)
    y_pred_eroded = erosion(y_pred)
    y_true_boundary = tf.abs(y_true - y_true_eroded)
    y_pred_boundary = tf.abs(y_pred - y_pred_eroded)
    intersection = tf.reduce_sum(y_true_boundary * y_pred_boundary)
    union = tf.reduce_sum(y_true_boundary) + tf.reduce_sum(y_pred_boundary) - intersection
    return (intersection + smooth) / (union + smooth)


@tf.function
def sensitivity(y_true, y_pred, smooth=1e-6):
    y_true = tf.cast(y_true, tf.float32)
    y_pred = tf.cast(y_pred > 0.5, tf.float32)
    true_pos = tf.reduce_sum(y_true * y_pred)
    false_neg = tf.reduce_sum(y_true * (1.0 - y_pred))
    return (true_pos + smooth) / (true_pos + false_neg + smooth)


@tf.function
def specificity(y_true, y_pred, smooth=1e-6):
    y_true = tf.cast(y_true, tf.float32)
    y_pred = tf.cast(y_pred > 0.5, tf.float32)
    true_neg = tf.reduce_sum((1.0 - y_true) * (1.0 - y_pred))
    false_pos = tf.reduce_sum((1.0 - y_true) * y_pred)
    return (true_neg + smooth) / (true_neg + false_pos + smooth)


@tf.function
def focal_loss(y_true, y_pred, alpha=0.75, gamma=2.0):
    y_true = tf.cast(y_true, tf.float32)
    y_pred = tf.cast(y_pred, tf.float32)
    y_pred = tf.clip_by_value(y_pred, 1e-7, 1.0 - 1e-7)
    pt_1 = tf.where(tf.equal(y_true, 1), y_pred, tf.ones_like(y_pred))
    focal_pos = -alpha * tf.pow(1.0 - pt_1, gamma) * tf.math.log(pt_1)
    pt_0 = tf.where(tf.equal(y_true, 0), y_pred, tf.zeros_like(y_pred))
    focal_neg = -(1 - alpha) * tf.pow(pt_0, gamma) * tf.math.log(1.0 - pt_0)
    return tf.reduce_mean(focal_pos + focal_neg)


@tf.function
def tversky_loss(y_true, y_pred, alpha=0.8, beta=0.2, smooth=1e-6):
    y_true = tf.cast(y_true, tf.float32)
    y_pred = tf.cast(y_pred, tf.float32)
    if len(y_true.shape) == 3:
        y_true = tf.expand_dims(y_true, axis=-1)
    if len(y_pred.shape) == 3:
        y_pred = tf.expand_dims(y_pred, axis=-1)
    y_true_f = tf.reshape(y_true, [-1])
    y_pred_f = tf.reshape(y_pred, [-1])
    true_pos = tf.reduce_sum(y_true_f * y_pred_f)
    false_neg = tf.reduce_sum(y_true_f * (1 - y_pred_f))
    false_pos = tf.reduce_sum((1 - y_true_f) * y_pred_f)
    tversky_index = (true_pos + smooth) / (true_pos + alpha * false_neg + beta * false_pos + smooth)
    return 1.0 - tversky_index


@tf.function
def focal_tversky_loss(y_true, y_pred, alpha=0.8, beta=0.2, gamma=0.75, smooth=1e-6):
    tversky = tversky_loss(y_true, y_pred, alpha, beta, smooth)
    return tf.pow(tversky, gamma)


@tf.function
def generalized_dice_loss(y_true, y_pred, smooth=1e-6):
    return 1.0 - generalized_dice_coef(y_true, y_pred, smooth)


@tf.function
def boundary_loss(y_true, y_pred, smooth=1e-6):
    y_true = tf.cast(y_true, tf.float32)
    y_pred = tf.cast(y_pred, tf.float32)
    if len(y_true.shape) == 3:
        y_true = tf.expand_dims(y_true, axis=-1)
    if len(y_pred.shape) == 3:
        y_pred = tf.expand_dims(y_pred, axis=-1)
    sobel_x = tf.reshape(tf.constant([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=tf.float32), [3, 3, 1, 1])
    sobel_y = tf.reshape(tf.constant([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=tf.float32), [3, 3, 1, 1])
    edges_true_x = tf.nn.conv2d(y_true, sobel_x, strides=[1, 1, 1, 1], padding='SAME')
    edges_true_y = tf.nn.conv2d(y_true, sobel_y, strides=[1, 1, 1, 1], padding='SAME')
    edges_true = tf.sqrt(edges_true_x ** 2 + edges_true_y ** 2 + smooth)
    edges_pred_x = tf.nn.conv2d(y_pred, sobel_x, strides=[1, 1, 1, 1], padding='SAME')
    edges_pred_y = tf.nn.conv2d(y_pred, sobel_y, strides=[1, 1, 1, 1], padding='SAME')
    edges_pred = tf.sqrt(edges_pred_x ** 2 + edges_pred_y ** 2 + smooth)
    return tf.reduce_mean(tf.square(edges_true - edges_pred))


@tf.function
def hybrid_loss_optimized(y_true, y_pred, w_ft=0.5, w_dice=0.35, w_boundary=0.1, w_focal=0.05):
    ft_loss = focal_tversky_loss(y_true, y_pred, alpha=0.8, beta=0.2, gamma=0.75)
    d_loss = generalized_dice_loss(y_true, y_pred)
    b_loss = boundary_loss(y_true, y_pred)
    f_loss = focal_loss(y_true, y_pred, alpha=0.75, gamma=2.0)
    return w_ft * ft_loss + w_dice * d_loss + w_boundary * b_loss + w_focal * f_loss


CUSTOM_OBJECTS = {
    'dice_coef': dice_coef,
    'generalized_dice_coef': generalized_dice_coef,
    'weighted_dice_coef': weighted_dice_coef,
    'iou': iou,
    'boundary_iou': boundary_iou,
    'sensitivity': sensitivity,
    'specificity': specificity,
    'focal_loss': focal_loss,
    'tversky_loss': tversky_loss,
    'focal_tversky_loss': focal_tversky_loss,
    'generalized_dice_loss': generalized_dice_loss,
    'boundary_loss': boundary_loss,
    'hybrid_loss_optimized': hybrid_loss_optimized,
}


# ============================================================
# TIEN XU LY - GIONG HET LOGIC TRONG ImprovedNiftiGenerator BAN GOC
# ============================================================
def load_image_and_mask_from_h5(path):
    """Doc anh (va mask neu co) tu file .h5."""
    with h5py.File(path, 'r') as f:
        if 'image' not in f:
            raise ValueError(f"File h5 khong co dataset 'image': {path}")
        img = f['image'][:].astype(np.float32)
        if img.ndim == 3:
            img = img[..., 0]

        mask = None
        if 'mask' in f:
            mask = f['mask'][:].astype(np.float32)
            if mask.ndim == 3:
                mask = mask[..., 0]

        return img, mask


def preprocess_image(img, mask=None):
    """Resize + normalize percentile (2, 98) - dung y nhu _load_slice va
    _normalize_batch trong ImprovedNiftiGenerator ban goc."""
    if img.shape != IMG_SIZE:
        img = cv2.resize(img, IMG_SIZE, interpolation=cv2.INTER_LINEAR)
    if mask is not None and mask.shape != IMG_SIZE:
        mask = cv2.resize(mask, IMG_SIZE, interpolation=cv2.INTER_NEAREST)
    if mask is not None:
        mask = (mask > 0).astype(np.float32)

    p2, p98 = np.percentile(img, (2, 98))
    img_norm = np.clip(img, p2, p98)
    img_norm = (img_norm - p2) / (p98 - p2 + 1e-8)

    X = img_norm[np.newaxis, ..., np.newaxis].astype(np.float32)  # (1, 256, 256, 1)
    return X, img, mask


def advanced_post_process(pred, threshold=0.25, min_size=100):
    """Hau xu ly du doan de loai bo noise va lam min (copy tu ban goc)."""
    binary_pred = (pred > threshold).astype(np.uint8)

    kernel_close = np.ones((3, 3), np.uint8)
    binary_pred = binary_closing(binary_pred, structure=kernel_close).astype(np.uint8)

    labeled, num_features = label(binary_pred)
    if num_features > 0:
        sizes = np.bincount(labeled.flat)[1:]
        for i, size in enumerate(sizes, 1):
            if size < min_size:
                binary_pred[labeled == i] = 0

    kernel_dilate = np.ones((2, 2), np.uint8)
    binary_pred = binary_dilation(binary_pred, structure=kernel_dilate).astype(np.uint8)

    return binary_pred


# ============================================================
# CHAY INFERENCE
# ============================================================
def run_inference(model_path, input_path, output_path=None, threshold=0.25, min_size=100):
    print(f"Dang load model: {model_path}")
    model = tf.keras.models.load_model(model_path, custom_objects=CUSTOM_OBJECTS)

    print(f"Dang doc va tien xu ly file: {input_path}")
    raw_img, raw_mask = load_image_and_mask_from_h5(input_path)
    X, img_resized, mask_resized = preprocess_image(raw_img, raw_mask)

    print("Dang du doan...")
    pred_raw = model.predict(X, verbose=0)[0, ..., 0]
    pred_processed = advanced_post_process(pred_raw, threshold=threshold, min_size=min_size)

    tumor_pixels = int(pred_processed.sum())
    total_pixels = pred_processed.size
    print("\n" + "=" * 50)
    print("KET QUA DU DOAN (SEGMENTATION)")
    print("=" * 50)
    print(f"So pixel duoc du doan la tumor: {tumor_pixels} / {total_pixels} "
          f"({100 * tumor_pixels / total_pixels:.2f}%)")
    print(f"Xac suat trung binh (raw) vung tumor: {pred_raw.mean():.4f}")
    print(f"Xac suat max (raw): {pred_raw.max():.4f}")

    # Neu file h5 co san ground truth mask -> tinh metric so sanh
    if mask_resized is not None:
        y_true = tf.constant(mask_resized[..., np.newaxis], dtype=tf.float32)
        y_pred = tf.constant(pred_processed[..., np.newaxis].astype(np.float32), dtype=tf.float32)

        dice = dice_coef(y_true, y_pred).numpy()
        iou_score = iou(y_true, y_pred).numpy()
        sens = sensitivity(y_true, y_pred).numpy()
        spec = specificity(y_true, y_pred).numpy()

        true_pos = np.sum((mask_resized == 1) & (pred_processed == 1))
        false_pos = np.sum((mask_resized == 0) & (pred_processed == 1))
        precision = true_pos / (true_pos + false_pos + 1e-7)
        f1 = 2 * precision * sens / (precision + sens + 1e-7)

        print("\nSo sanh voi ground truth mask co san trong file h5:")
        print(f"  Dice coefficient : {dice:.4f}")
        print(f"  IoU              : {iou_score:.4f}")
        print(f"  Sensitivity      : {sens:.4f}")
        print(f"  Specificity      : {spec:.4f}")
        print(f"  Precision        : {precision:.4f}")
        print(f"  F1 Score         : {f1:.4f}")
    else:
        print("\n(File h5 khong co dataset 'mask' nen khong tinh duoc metric so sanh)")

    print("=" * 50)

    # Xuat anh minh hoa
    if output_path:
        n_panels = 4 if mask_resized is not None else 3
        fig, axes = plt.subplots(1, n_panels, figsize=(4 * n_panels, 4))

        axes[0].imshow(img_resized, cmap='gray')
        axes[0].set_title('Anh dau vao')
        axes[0].axis('off')

        idx = 1
        if mask_resized is not None:
            axes[idx].imshow(mask_resized, cmap='jet', alpha=0.8)
            axes[idx].set_title('Ground Truth')
            axes[idx].axis('off')
            idx += 1

        axes[idx].imshow(pred_processed, cmap='jet', alpha=0.8)
        axes[idx].set_title('Du doan')
        axes[idx].axis('off')
        idx += 1

        axes[idx].imshow(img_resized, cmap='gray')
        axes[idx].imshow(pred_processed, cmap='jet', alpha=0.5)
        axes[idx].set_title('Overlay')
        axes[idx].axis('off')

        plt.tight_layout()
        plt.savefig(output_path, dpi=150, bbox_inches='tight')
        plt.close()
        print(f"\nDa luu anh minh hoa tai: {output_path}")

    return pred_raw, pred_processed


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Test model segmentation khoi u nao voi 1 file .h5")
    parser.add_argument("--model", required=True, help="Duong dan file model .keras")
    parser.add_argument("--input", required=True, help="Duong dan file .h5 can test (co dataset 'image', 'mask' neu co)")
    parser.add_argument("--output", default=None, help="Duong dan file .png de luu anh minh hoa (tuy chon)")
    parser.add_argument("--threshold", type=float, default=0.25, help="Nguong nhi phan hoa du doan (mac dinh 0.25)")
    parser.add_argument("--min_size", type=int, default=100, help="Kich thuoc vung nho nhat giu lai, tinh bang pixel (mac dinh 100)")
    args = parser.parse_args()

    run_inference(args.model, args.input, args.output, args.threshold, args.min_size)