"""
model.py
--------
Model definitions for Real-World Lung CT Analysis:
  1. load_pretrained_model()  -- existing 2D inference model
  2. build_unet_3d()          -- [NEW] 3D volumetric segmentation model
  3. build_nodule_detector()  -- [NEW] Nodule classification CNN
  4. Custom losses and metrics (Dice, IoU, Focal Loss for class imbalance)
"""

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers


# ===========================================================================
# LOSSES & METRICS
# ===========================================================================

def dice_coefficient(y_true, y_pred, smooth=1.0):
    """Dice Similarity Coefficient (works for 2D and 3D)."""
    y_true_f = keras.backend.flatten(tf.cast(y_true, tf.float32))
    y_pred_f = keras.backend.flatten(tf.cast(y_pred, tf.float32))
    intersection = tf.reduce_sum(y_true_f * y_pred_f)
    return (2.0 * intersection + smooth) / (
        tf.reduce_sum(y_true_f) + tf.reduce_sum(y_pred_f) + smooth
    )

def dice_loss(y_true, y_pred):
    return 1.0 - dice_coefficient(y_true, y_pred)

def iou_score(y_true, y_pred, smooth=1.0):
    """Intersection over Union (Jaccard Index)."""
    y_true_f = keras.backend.flatten(tf.cast(y_true, tf.float32))
    y_pred_f = keras.backend.flatten(tf.cast(y_pred, tf.float32))
    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)

def combined_loss(y_true, y_pred):
    """BCE + Dice Loss for highly accurate segmentation boundaries."""
    bce = keras.losses.binary_crossentropy(
        tf.cast(y_true, tf.float32),
        tf.cast(y_pred, tf.float32)
    )
    dl = dice_loss(y_true, y_pred)
    return 0.5 * bce + 0.5 * dl

def focal_loss(gamma=2.0, alpha=0.25):
    """
    Focal Loss handles severe class imbalance (e.g., millions of healthy lung
    pixels vs a few nodule pixels). It down-weights well-classified examples.
    """
    def focal_loss_fixed(y_true, y_pred):
        pt_1 = tf.where(tf.equal(y_true, 1), y_pred, tf.ones_like(y_pred))
        pt_0 = tf.where(tf.equal(y_true, 0), y_pred, tf.zeros_like(y_pred))
        return -keras.backend.sum(alpha * keras.backend.pow(1. - pt_1, gamma) * keras.backend.log(pt_1 + keras.backend.epsilon())) \
               -keras.backend.sum((1 - alpha) * keras.backend.pow(pt_0, gamma) * keras.backend.log(1. - pt_0 + keras.backend.epsilon()))
    return focal_loss_fixed


# ===========================================================================
# 2D SEGMENTATION
# ===========================================================================

def load_pretrained_model(model_path="LungCT_Seg.keras"):
    """Load the original 2D Keras model."""
    return keras.models.load_model(model_path, compile=False)


# ===========================================================================
# 3D VOLUMETRIC SEGMENTATION  [NEW]
# ===========================================================================

def conv_block_3d(x, filters, kernel_size=3, dropout_rate=0.1):
    """3D Convolutional Block: Conv3D -> BN -> ReLU -> Dropout -> Conv3D"""
    x = layers.Conv3D(filters, kernel_size, padding="same", use_bias=False)(x)
    x = layers.BatchNormalization()(x)
    x = layers.Activation("relu")(x)
    x = layers.Dropout(dropout_rate)(x)

    x = layers.Conv3D(filters, kernel_size, padding="same", use_bias=False)(x)
    x = layers.BatchNormalization()(x)
    x = layers.Activation("relu")(x)
    return x

def build_unet_3d(input_shape=(64, 64, 64, 1), base_filters=16, depth=3):
    """
    3D U-Net for Volumetric CT Segmentation.

    Unlike 2D models that process slice-by-slice, a 3D U-Net processes a cubic
    volume (patch) of the scan. This preserves spatial context across the Z-axis
    (depth), which is crucial for distinguishing spherical nodules from cylindrical
    blood vessels.

    Args:
        input_shape  : tuple, (D, H, W, C). Default 64x64x64 patches.
        base_filters : int, filters in first block. Kept small (16) to fit in GPU RAM.
        depth        : int, number of downsampling levels.

    Returns:
        model : keras.Model (Output shape = input_shape)
    """
    inputs = keras.Input(shape=input_shape, name="ct_volume_input")

    # ── Encoder ──────────────────────────────────────────────────────────────
    encoder_outputs = []
    x = inputs
    for i in range(depth):
        filters = base_filters * (2 ** i)
        x = conv_block_3d(x, filters)
        encoder_outputs.append(x)
        x = layers.MaxPooling3D(pool_size=(2, 2, 2))(x)

    # ── Bottleneck ────────────────────────────────────────────────────────────
    bottleneck_filters = base_filters * (2 ** depth)
    x = conv_block_3d(x, bottleneck_filters)

    # ── Decoder ───────────────────────────────────────────────────────────────
    for i in reversed(range(depth)):
        filters = base_filters * (2 ** i)
        x = layers.UpSampling3D(size=(2, 2, 2))(x)
        x = layers.Concatenate()([x, encoder_outputs[i]])
        x = conv_block_3d(x, filters)

    # ── Output ────────────────────────────────────────────────────────────────
    outputs = layers.Conv3D(1, kernel_size=1, activation="sigmoid", name="seg_mask_3d")(x)

    return keras.Model(inputs, outputs, name="UNet_3D")


# ===========================================================================
# NODULE / ANOMALY DETECTOR  [NEW]
# ===========================================================================

def build_nodule_detector(input_shape=(32, 32, 1)):
    """
    CNN Classifier for identifying lung nodules in localized 2D patches.

    This model is used by the sliding-window detector in predict.py.
    It takes a small 32x32 image patch (cropped from inside the lung mask)
    and predicts whether it contains a nodule.

    Architecture: Standard VGG-style Conv2D + Dense classification head.

    Args:
        input_shape : tuple, expected (32, 32, 1) or (64, 64, 1)

    Returns:
        model : keras.Model (Output = scalar probability [0, 1])
    """
    inputs = keras.Input(shape=input_shape, name="patch_input")

    x = layers.Conv2D(32, 3, padding="same", activation="relu")(inputs)
    x = layers.BatchNormalization()(x)
    x = layers.MaxPooling2D(2)(x)   # 16x16

    x = layers.Conv2D(64, 3, padding="same", activation="relu")(x)
    x = layers.BatchNormalization()(x)
    x = layers.MaxPooling2D(2)(x)   # 8x8

    x = layers.Conv2D(128, 3, padding="same", activation="relu")(x)
    x = layers.BatchNormalization()(x)
    x = layers.MaxPooling2D(2)(x)   # 4x4

    # Global Average Pooling prevents overfitting better than Flatten()
    x = layers.GlobalAveragePooling2D()(x)

    x = layers.Dense(128, activation="relu")(x)
    x = layers.Dropout(0.5)(x)

    # Single output neuron: 1 = Nodule, 0 = Healthy tissue/vessel
    outputs = layers.Dense(1, activation="sigmoid", name="nodule_prob")(x)

    return keras.Model(inputs, outputs, name="Nodule_Detector")