"""BRN: Bbox Refinement Network
===============================
A lightweight CNN that takes a high-resolution ROI crop from the ORIGINAL IMAGE
and predicts sub-pixel (Δcx, Δcy, Δw, Δh) corrections to a bounding box.

Unlike YOLO which sees the world through stride-8 features, BRN sees RAW PIXELS
at 1:1 resolution. This gives it precision that YOLO physically cannot achieve.

Architecture: EfficientNet-B0-lite inspired, ~500K params.
Training: Generate infinite samples by adding noise to GT boxes.
Inference: YOLO detects → BRN refines → final boxes.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as T
import numpy as np
from PIL import Image


class BRN(nn.Module):
    """Lightweight Bbox Refinement Network.

    Input:  [B, 3, H, W] ROI crop (default 96x96)
    Output: [B, 4] (Δcx, Δcy, Δw, Δh) in normalized ROI coordinates [-0.5, 0.5]
    """

    def __init__(self, input_size=96):
        super().__init__()

        # Stem: 96 → 48
        self.stem = nn.Sequential(
            nn.Conv2d(3, 32, 3, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(32),
            nn.SiLU(inplace=True),
        )
        # Block1: 48 → 24
        self.block1 = nn.Sequential(
            nn.Conv2d(32, 48, 3, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(48),
            nn.SiLU(inplace=True),
            nn.Conv2d(48, 48, 3, padding=1, bias=False),
            nn.BatchNorm2d(48),
            nn.SiLU(inplace=True),
        )
        # Block2: 24 → 12
        self.block2 = nn.Sequential(
            nn.Conv2d(48, 64, 3, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.SiLU(inplace=True),
            nn.Conv2d(64, 64, 3, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.SiLU(inplace=True),
        )
        # Block3: 12 → 6
        self.block3 = nn.Sequential(
            nn.Conv2d(64, 96, 3, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(96),
            nn.SiLU(inplace=True),
            nn.Conv2d(96, 96, 3, padding=1, bias=False),
            nn.BatchNorm2d(96),
            nn.SiLU(inplace=True),
        )
        # Head
        self.head = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(96, 64),
            nn.SiLU(inplace=True),
            nn.Dropout(0.1),
            nn.Linear(64, 4),
        )

        self._init_weights()

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.zeros_(m.bias)

    def forward(self, x):
        x = self.stem(x)
        x = self.block1(x)
        x = self.block2(x)
        x = self.block3(x)
        x = self.head(x)
        return x


class BRNRefiner:
    """BRN-based bounding box refiner.

    Usage:
        refiner = BRNRefiner('path/to/brn.pt', input_size=96)
        refined_boxes = refiner.refine(image, yolo_boxes)
    """

    def __init__(self, model_path, input_size=96, device='cuda'):
        self.input_size = input_size
        self.device = device

        self.model = BRN(input_size=input_size).to(device)
        if model_path:
            self.model.load_state_dict(torch.load(model_path, map_location=device))
        self.model.eval()

        self.transform = T.Compose([
            T.ToTensor(),
            T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])

    @torch.no_grad()
    def refine(self, image, boxes):
        """Refine bounding boxes using BRN.

        Args:
            image: PIL Image or numpy array
            boxes: [N, 4] in (x1, y1, x2, y2) pixel coordinates
        Returns:
            refined: [N, 4] refined boxes
        """
        if len(boxes) == 0:
            return boxes

        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        iw, ih = image.size
        refined = []

        for box in boxes:
            x1, y1, x2, y2 = box
            cx, cy = (x1 + x2) / 2, (y1 + y2) / 2
            w, h = max(x2 - x1, 1), max(y2 - y1, 1)

            # Expand ROI by 50% to give context
            roi_size = max(w, h) * 1.5
            roi_size = min(roi_size, min(iw, ih))  # Don't exceed image
            half = roi_size / 2

            # ROI bounds
            rx1 = max(0, cx - half)
            ry1 = max(0, cy - half)
            rx2 = min(iw, cx + half)
            ry2 = min(ih, cy + half)

            # Crop and resize
            roi = image.crop((rx1, ry1, rx2, ry2))
            roi_tensor = self.transform(roi.resize((self.input_size, self.input_size), Image.BICUBIC))
            roi_tensor = roi_tensor.unsqueeze(0).to(self.device)

            # Predict correction
            delta = self.model(roi_tensor)[0].cpu().numpy()  # [Δcx, Δcy, Δw, Δh]

            # Decode: correction is in ROI-normalized space
            roi_w, roi_h = rx2 - rx1, ry2 - ry1
            dcx = delta[0] * roi_w
            dcy = delta[1] * roi_h
            dw = delta[2] * roi_w
            dh = delta[3] * roi_h

            new_cx = cx + dcx
            new_cy = cy + dcy
            new_w = max(1, w + dw)
            new_h = max(1, h + dh)

            new_x1 = max(0, new_cx - new_w / 2)
            new_y1 = max(0, new_cy - new_h / 2)
            new_x2 = min(iw, new_cx + new_w / 2)
            new_y2 = min(ih, new_cy + new_h / 2)

            refined.append([new_x1, new_y1, new_x2, new_y2])

        return np.array(refined)


# ── Training data generation ────────────────────────────────

def generate_training_sample(image, gt_boxes, noise_level=0.05):
    """Generate a BRN training sample from one image.

    For each GT box, add random noise and create a (noisy_box, GT_box) pair.

    Args:
        image: PIL Image
        gt_boxes: [N, 4] in (cx, cy, w, h) normalized [0,1]
        noise_level: std of Gaussian noise relative to box size

    Returns:
        rois: [M, 3, S, S] tensor of ROI crops
        deltas: [M, 4] target corrections to recover GT from noisy box
    """
    iw, ih = image.size
    rois = []
    deltas = []

    for box in gt_boxes:
        cx, cy, w, h = box
        cx_px, cy_px = cx * iw, cy * ih
        w_px, h_px = w * iw, h * ih

        # Add Gaussian noise to all 4 params
        std_c = noise_level * min(w_px, h_px)  # Center noise proportional to size
        std_s = noise_level * max(w_px, h_px)  # Size noise

        noisy_cx = cx_px + np.random.randn() * std_c
        noisy_cy = cy_px + np.random.randn() * std_c
        noisy_w = w_px * (1 + np.random.randn() * noise_level)
        noisy_h = h_px * (1 + np.random.randn() * noise_level)
        noisy_w = max(4, noisy_w)
        noisy_h = max(4, noisy_h)

        # Compute target delta: what correction needed from noisy → GT
        delta_cx = (cx_px - noisy_cx) / max(w_px, h_px)  # Normalize by box size
        delta_cy = (cy_px - noisy_cy) / max(w_px, h_px)
        delta_w = (w_px - noisy_w) / max(w_px, h_px)
        delta_h = (h_px - noisy_h) / max(w_px, h_px)

        # Expand ROI around noisy box
        half = max(noisy_w, noisy_h) * 0.75
        half = max(half, 16)  # Minimum ROI size

        rx1 = max(0, int(noisy_cx - half))
        ry1 = max(0, int(noisy_cy - half))
        rx2 = min(iw, int(noisy_cx + half))
        ry2 = min(ih, int(noisy_cy + half))

        if rx2 - rx1 < 8 or ry2 - ry1 < 8:
            continue

        try:
            roi = image.crop((rx1, ry1, rx2, ry2))
            rois.append(roi)
            deltas.append([delta_cx, delta_cy, delta_w, delta_h])
        except Exception:
            continue

    return rois, deltas