"""
Utility functions for preprocessing, TTA inference, postprocessing,
visualization, and Kaggle RLE formatting.
"""

import cv2
import numpy as np
import torch
from PIL import Image

IMG_H = 128
IMG_W = 800
NORM_MEAN = 0.449
NORM_STD = 0.226

CLASS_NAMES = ["defect_a", "defect_b", "defect_c", "defect_d"]
CLASS_LABELS = {
    "defect_a": "Defect A (Pitted / Scratch)",
    "defect_b": "Defect B (Minor defect)",
    "defect_c": "Defect C (Inclusion / Slag)",
    "defect_d": "Defect D (Surface patch / Scuff)"
}

# Color palette for classes (R, G, B)
CLASS_COLORS = [
    (255, 60, 60),    # Defect A: Red
    (50, 220, 50),    # Defect B: Green
    (30, 150, 255),   # Defect C: Cyan-Blue
    (255, 180, 20)    # Defect D: Orange-Gold
]

# Calibrated optimal thresholds & min component areas from notebook
DEFAULT_THRESHOLDS = [0.35, 0.22, 0.42, 0.42]
DEFAULT_MIN_AREAS = [25, 25, 220, 350]


def rle_encode(mask):
    """
    1-indexed, column-major (Fortran) Run-Length Encoding.
    Matches the official Kaggle competition format.
    """
    mask = np.asarray(mask)
    if mask.ndim != 2:
        raise ValueError(f"rle_encode expects a 2D mask, got shape {mask.shape}")
    pixels = np.concatenate([[0], (mask > 0).astype(np.uint8).flatten(order="F"), [0]])
    runs = np.where(pixels[1:] != pixels[:-1])[0] + 1
    runs[1::2] -= runs[::2]
    return " ".join(map(str, runs))


def rle_decode(rle, shape=(IMG_H, IMG_W)):
    """Decodes 1-indexed column-major RLE string back to 2D numpy mask."""
    total = int(shape[0] * shape[1])
    flat = np.zeros(total, dtype=np.uint8)
    if not isinstance(rle, str) or not rle.strip():
        return flat.reshape(shape, order="F")
    nums = np.asarray(rle.split(), dtype=np.int64)
    if nums.size % 2:
        return flat.reshape(shape, order="F")
    starts = nums[0::2] - 1
    lengths = nums[1::2]
    for start, length in zip(starts, lengths):
        end = min(int(start + length), total)
        if 0 <= start < total:
            flat[int(start):end] = 1
    return flat.reshape(shape, order="F")


def filter_small_components(mask, min_area):
    """Removes 8-connected components smaller than min_area pixels."""
    if min_area <= 1 or not mask.any():
        return mask
    mask_u8 = mask.astype(np.uint8)
    count, labels, stats, _ = cv2.connectedComponentsWithStats(mask_u8, connectivity=8)
    if count <= 1:
        return np.zeros_like(mask, dtype=bool)
    keep = np.flatnonzero(stats[:, cv2.CC_STAT_AREA] >= min_area)
    keep = keep[keep != 0]  # exclude background (0)
    return np.isin(labels, keep)


def preprocess_image(input_img):
    """
    Prepares input image for model:
    - Converts to Grayscale
    - Resizes to (IMG_H=128, IMG_W=800)
    - Normalizes with dataset mean & std
    Returns:
        tensor: shape (1, 1, 128, 800)
        display_rgb: numpy uint8 (128, 800, 3) for overlay rendering
    """
    if isinstance(input_img, str):
        gray = cv2.imread(input_img, cv2.IMREAD_GRAYSCALE)
    elif isinstance(input_img, np.ndarray):
        if input_img.ndim == 3:
            gray = cv2.cvtColor(input_img, cv2.COLOR_RGB2GRAY)
        else:
            gray = input_img.copy()
    elif isinstance(input_img, Image.Image):
        gray = np.array(input_img.convert("L"))
    else:
        raise ValueError("Unsupported image input format.")

    if gray.shape != (IMG_H, IMG_W):
        gray = cv2.resize(gray, (IMG_W, IMG_H), interpolation=cv2.INTER_AREA)

    norm_img = gray.astype(np.float32) / 255.0
    norm_img = (norm_img - NORM_MEAN) / NORM_STD
    tensor = torch.from_numpy(norm_img).unsqueeze(0).unsqueeze(0)  # (1, 1, 128, 800)

    display_rgb = cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB)
    return tensor, display_rgb


@torch.no_grad()
def predict_with_tta(model, tensor_x, tta=True, device="cpu"):
    """
    Inference with optional 4-view Test-Time Augmentation (flips).
    Returns numpy probability array of shape (4, 128, 800) with values in [0, 1].
    """
    model.eval()
    tensor_x = tensor_x.to(device)
    probs = torch.sigmoid(model(tensor_x).float())
    if tta:
        # 4-view flip TTA: original, hflip, vflip, hvflip
        for dims in ([3], [2], [2, 3]):
            flipped_in = torch.flip(tensor_x, dims)
            flipped_pred = torch.sigmoid(model(flipped_in).float())
            probs = probs + torch.flip(flipped_pred, dims)
        probs = probs / 4.0
    return probs.squeeze(0).cpu().numpy()


def apply_postprocess(probs_np, thresholds=None, min_areas=None):
    """
    Applies calibrated probability thresholds and minimum connected component areas.
    """
    if thresholds is None:
        thresholds = DEFAULT_THRESHOLDS
    if min_areas is None:
        min_areas = DEFAULT_MIN_AREAS

    n_cls = probs_np.shape[0]
    masks = np.zeros_like(probs_np, dtype=bool)

    for c in range(n_cls):
        th = float(thresholds[c])
        area = int(min_areas[c])
        raw_mask = probs_np[c] >= th
        masks[c] = filter_small_components(raw_mask, area)

    return masks


def create_segmentation_overlay(base_rgb, masks, alpha=0.55):
    """
    Overlays colored semi-transparent defect masks and clear contour borders onto base image.
    """
    overlay = base_rgb.copy().astype(np.float32)

    for c, mask in enumerate(masks):
        if not mask.any():
            continue
        color = np.array(CLASS_COLORS[c], dtype=np.float32)
        # Apply semi-transparent fill
        idx = mask > 0
        overlay[idx] = (1 - alpha) * overlay[idx] + alpha * color

    overlay = np.clip(overlay, 0, 255).astype(np.uint8)

    # Draw sharp contour borders on top for precision visibility
    for c, mask in enumerate(masks):
        if not mask.any():
            continue
        mask_u8 = mask.astype(np.uint8)
        contours, _ = cv2.findContours(mask_u8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        color = CLASS_COLORS[c]
        cv2.drawContours(overlay, contours, -1, color, 1)

    return overlay


def compute_defect_diagnostics(masks, display_shape=(IMG_H, IMG_W)):
    """
    Computes per-class analytics:
    - Detection boolean
    - Component count
    - Total pixel area
    - Coverage percentage
    - RLE encoding string
    """
    total_canvas_pixels = float(display_shape[0] * display_shape[1])
    diagnostics = []
    has_any_defect = False

    for c, name in enumerate(CLASS_NAMES):
        mask = masks[c]
        is_detected = bool(mask.any())
        if is_detected:
            has_any_defect = True
            count, _, stats, _ = cv2.connectedComponentsWithStats(mask.astype(np.uint8), connectivity=8)
            comp_count = max(0, count - 1)
            area_px = int(mask.sum())
            area_pct = round(100.0 * (area_px / total_canvas_pixels), 3)
            rle_str = rle_encode(mask)
        else:
            comp_count = 0
            area_px = 0
            area_pct = 0.0
            rle_str = ""

        diagnostics.append({
            "Class Name": name,
            "Label": CLASS_LABELS[name],
            "Detected": "YES" if is_detected else "NO",
            "Components": comp_count,
            "Area (px)": area_px,
            "Surface %": f"{area_pct}%",
            "RLE Encoded": (rle_str[:60] + "...") if len(rle_str) > 60 else (rle_str or "EMPTY")
        })

    overall_status = "DEFECT DETECTED" if has_any_defect else "CLEAN / NO DEFECT"
    return diagnostics, overall_status