"""
Raw Custom Object Detection Neural Network
Built from mathematical first principles in pure PyTorch (torch.nn)
Zero third-party model libraries or frameworks.
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F


class ConvBNAct(nn.Module):
    """
    Mathematical Convolutional Block:
    Z = W * X + b
    Z_norm = gamma * (Z - mu) / sqrt(var + eps) + beta
    Out = SiLU(Z_norm) = Z_norm * sigmoid(Z_norm)
    """
    def __init__(self, in_c, out_c, kernel_size=3, stride=1, padding=1):
        super().__init__()
        self.conv = nn.Conv2d(in_c, out_c, kernel_size, stride, padding, bias=False)
        self.bn = nn.BatchNorm2d(out_c)
        self.act = nn.SiLU(inplace=True)

    def forward(self, x):
        return self.act(self.bn(self.conv(x)))


class SqueezeExcitation(nn.Module):
    """
    Mathematical Squeeze-and-Excitation Channel Attention:
    Z_se = Sigmoid( W2 * SiLU( W1 * AdaptiveAvgPool2d(X) ) )
    Out = X * Z_se
    Dynamically recalibrates channel-wise feature responses to amplify defect-specific
    textural frequencies (corrosion roughness, unbolted joint cavities) while suppressing
    clean structural background steelwork.
    """
    def __init__(self, channels, reduction=16):
        super().__init__()
        reduced = max(8, channels // reduction)
        self.fc = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(channels, reduced, kernel_size=1, bias=True),
            nn.SiLU(inplace=True),
            nn.Conv2d(reduced, channels, kernel_size=1, bias=True),
            nn.Sigmoid()
        )

    def forward(self, x):
        return x * self.fc(x)


class ResidualBlock(nn.Module):
    """
    Residual Skip Connection with Squeeze-and-Excitation Channel Attention:
    y = F(x, {W_i}) * SE(F(x)) + x
    Mitigates vanishing gradient and enables deep feature representations.
    """
    def __init__(self, channels, use_se=True):
        super().__init__()
        reduced = max(16, channels // 2)
        self.c1 = ConvBNAct(channels, reduced, kernel_size=1, stride=1, padding=0)
        self.c2 = ConvBNAct(reduced, channels, kernel_size=3, stride=1, padding=1)
        self.se = SqueezeExcitation(channels) if use_se else nn.Identity()

    def forward(self, x):
        return x + self.se(self.c2(self.c1(x)))


class SPPFLayer(nn.Module):
    """
    Spatial Pyramid Pooling - Fast (SPPF):
    Fuses multi-scale receptive field context via sequential MaxPool2d(5, stride=1, padding=2),
    yielding equivalent pooling at 5x5, 9x9, and 13x13 receptive fields.
    Dramatically enhances semantic context for large defects (vegetation canopy,
    extended corrosion fields) with minimal computational latency.
    """
    def __init__(self, in_c, out_c):
        super().__init__()
        c_mid = in_c // 2
        self.c1 = ConvBNAct(in_c, c_mid, kernel_size=1, stride=1, padding=0)
        self.c2 = ConvBNAct(c_mid * 4, out_c, kernel_size=1, stride=1, padding=0)
        self.pool = nn.MaxPool2d(kernel_size=5, stride=1, padding=2)

    def forward(self, x):
        x = self.c1(x)
        p1 = self.pool(x)
        p2 = self.pool(p1)
        p3 = self.pool(p2)
        return self.c2(torch.cat([x, p1, p2, p3], dim=1))


class CustomBackbone(nn.Module):
    """
    Deep Convolutional Feature Extractor with Multi-Scale Pyramid Hierarchy:
    Input:  (B, 3, H, W)
    Stem:   (B, 32, H/2, W/2)    - Stride 2
    Stage1: (B, 64, H/4, W/4)    - Stride 4
    Stage2: (B, 128, H/8, W/8)   - Stride 8  -> Feature Map P3 (high resolution, small objects: bolts)
    Stage3: (B, 256, H/16, W/16) - Stride 16 -> Feature Map P4 (medium objects: nests, insulators)
    Stage4: (B, 512, H/32, W/32) - Stride 32 -> Feature Map P5 (global context & large defects: rust/veg)
    """
    def __init__(self, in_channels=3):
        super().__init__()
        # Stem
        self.stem = ConvBNAct(in_channels, 32, kernel_size=3, stride=2, padding=1)

        # Stage 1: Stride 4
        self.stage1 = nn.Sequential(
            ConvBNAct(32, 64, kernel_size=3, stride=2, padding=1),
            ResidualBlock(64),
        )

        # Stage 2: Stride 8 (P3)
        self.stage2 = nn.Sequential(
            ConvBNAct(64, 128, kernel_size=3, stride=2, padding=1),
            ResidualBlock(128),
            ResidualBlock(128),
        )

        # Stage 3: Stride 16 (P4)
        self.stage3 = nn.Sequential(
            ConvBNAct(128, 256, kernel_size=3, stride=2, padding=1),
            ResidualBlock(256),
            ResidualBlock(256),
        )

        # Stage 4: Stride 32 (P5) with SPPF Receptive Field Aggregation
        self.stage4 = nn.Sequential(
            ConvBNAct(256, 512, kernel_size=3, stride=2, padding=1),
            ResidualBlock(512),
            ResidualBlock(512),
            SPPFLayer(512, 512),
        )

    def forward(self, x):
        x = self.stem(x)
        x = self.stage1(x)
        p3 = self.stage2(x)
        p4 = self.stage3(p3)
        p5 = self.stage4(p4)
        return p3, p4, p5


class FeaturePyramidNeck(nn.Module):
    """
    3-Scale Multi-Scale Feature Aggregator (FPN + PAN dual-path fusion):
    Fuses high-level semantic abstractions from P5 and P4 into fine-grained spatial details in P3,
    then propagates refined spatial patterns back up to P4 and P5.
    """
    def __init__(self, c3=128, c4=256, c5=512, out_c=128):
        super().__init__()
        # Lateral 1x1 convolutions to project all stages to uniform neck channel dimension
        self.lat_p5 = ConvBNAct(c5, out_c, kernel_size=1, stride=1, padding=0)
        self.lat_p4 = ConvBNAct(c4, out_c, kernel_size=1, stride=1, padding=0)
        self.lat_p3 = ConvBNAct(c3, out_c, kernel_size=1, stride=1, padding=0)

        # Top-down FPN smooth convolutions
        self.smooth_p4_td = ConvBNAct(out_c * 2, out_c, kernel_size=3, stride=1, padding=1)
        self.smooth_p3_td = ConvBNAct(out_c * 2, out_c, kernel_size=3, stride=1, padding=1)

        # Bottom-up PAN downsampling and smooth convolutions
        self.down_p3 = ConvBNAct(out_c, out_c, kernel_size=3, stride=2, padding=1)
        self.smooth_p4_pan = ConvBNAct(out_c * 2, out_c, kernel_size=3, stride=1, padding=1)
        self.down_p4 = ConvBNAct(out_c, out_c, kernel_size=3, stride=2, padding=1)
        self.smooth_p5_pan = ConvBNAct(out_c * 2, out_c, kernel_size=3, stride=1, padding=1)

    def forward(self, p3, p4, p5):
        # 1. Lateral projections
        p5_lat = self.lat_p5(p5)
        p4_lat = self.lat_p4(p4)
        p3_lat = self.lat_p3(p3)

        # 2. Top-down pathway (FPN)
        p5_up = F.interpolate(p5_lat, size=p4.shape[2:], mode="nearest")
        p4_td = self.smooth_p4_td(torch.cat([p4_lat, p5_up], dim=1))

        p4_up = F.interpolate(p4_td, size=p3.shape[2:], mode="nearest")
        p3_out = self.smooth_p3_td(torch.cat([p3_lat, p4_up], dim=1))

        # 3. Bottom-up pathway (PAN)
        p3_down = self.down_p3(p3_out)
        p4_out = self.smooth_p4_pan(torch.cat([p4_td, p3_down], dim=1))

        p4_down = self.down_p4(p4_out)
        p5_out = self.smooth_p5_pan(torch.cat([p5_lat, p4_down], dim=1))

        return p3_out, p4_out, p5_out


class DecoupledDetectionHead(nn.Module):
    """
    Mathematical Decoupled Prediction Head:
    Separates Box Regression [tx, ty, tw, th] from Objectness [obj] and Class Probabilities [cls_1..cls_C]
    """
    def __init__(self, in_c, num_classes):
        super().__init__()
        self.num_classes = num_classes

        # Coordinate branch (tx, ty, tw, th)
        self.box_conv = nn.Sequential(
            ConvBNAct(in_c, in_c, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_c, 4, kernel_size=1)
        )

        # Objectness branch
        self.obj_conv = nn.Sequential(
            ConvBNAct(in_c, in_c // 2, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_c // 2, 1, kernel_size=1)
        )

        # Classification branch
        self.cls_conv = nn.Sequential(
            ConvBNAct(in_c, in_c, kernel_size=3, stride=1, padding=1),
            nn.Conv2d(in_c, num_classes, kernel_size=1)
        )

        # Mathematical initialization: initialize objectness bias with negative prior
        # sigmoid(-4.59) ~ 0.01 to prevent initial gradient explosion on mostly empty backgrounds
        nn.init.constant_(self.obj_conv[-1].bias, -4.59)

    def forward(self, x):
        # x: (B, C, H_g, W_g)
        box = self.box_conv(x).permute(0, 2, 3, 1)  # (B, H_g, W_g, 4)
        obj = self.obj_conv(x).permute(0, 2, 3, 1)  # (B, H_g, W_g, 1)
        cls_logits = self.cls_conv(x).permute(0, 2, 3, 1)  # (B, H_g, W_g, num_classes)

        # Concatenate predictions: [tx, ty, tw, th, obj, c1, ..., cK]
        return torch.cat([box, obj, cls_logits], dim=-1)


class RawCustomDetector(nn.Module):
    """
    End-to-end Raw Custom Object Detector
    Combines:
      - Custom Deep Residual Backbone with Squeeze-and-Excitation & SPPF
      - 3-Scale Multi-scale Feature Pyramid Neck (FPN + PAN)
      - Triple-scale Decoupled Detection Heads:
          * Stride 8  (P3): High resolution for tiny defects (missing_bolt, small rust spots)
          * Stride 16 (P4): Medium receptive field (bird_nest, damaged_insulator)
          * Stride 32 (P5): Global receptive field (extensive rust zones, vegetation canopy)
    """
    def __init__(self, num_classes=5, in_channels=3):
        super().__init__()
        self.num_classes = num_classes
        self.strides = [8, 16, 32]

        self.backbone = CustomBackbone(in_channels=in_channels)
        self.neck = FeaturePyramidNeck(c3=128, c4=256, c5=512, out_c=128)
        self.head_p3 = DecoupledDetectionHead(in_c=128, num_classes=num_classes)
        self.head_p4 = DecoupledDetectionHead(in_c=128, num_classes=num_classes)
        self.head_p5 = DecoupledDetectionHead(in_c=128, num_classes=num_classes)

    def forward(self, x):
        p3, p4, p5 = self.backbone(x)
        p3_fused, p4_fused, p5_fused = self.neck(p3, p4, p5)

        pred_p3 = self.head_p3(p3_fused)  # (B, H/8, W/8, 5 + num_classes)
        pred_p4 = self.head_p4(p4_fused)  # (B, H/16, W/16, 5 + num_classes)
        pred_p5 = self.head_p5(p5_fused)  # (B, H/32, W/32, 5 + num_classes)

        return [pred_p3, pred_p4, pred_p5]


def decode_predictions(pred_maps, strides=None, conf_thresh=0.25):
    """
    Mathematical coordinate decoding from raw model logits to bounding boxes [x1, y1, x2, y2, score, class_id].
    bx = (grid_x + sigmoid(tx)) * stride
    by = (grid_y + sigmoid(ty)) * stride
    bw = exp(tw) * stride
    bh = exp(th) * stride
    """
    if strides is None:
        if len(pred_maps) == 3:
            strides = [8, 16, 32]
        else:
            strides = [8, 16]
    batch_detections = []
    batch_size = pred_maps[0].shape[0]

    for b in range(batch_size):
        all_boxes = []
        all_scores = []
        all_classes = []

        for pred_map, stride in zip(pred_maps, strides):
            pred = pred_map[b]  # (H_g, W_g, 5 + C)
            h_g, w_g, _ = pred.shape
            device = pred.device

            tx = pred[..., 0]
            ty = pred[..., 1]
            tw = pred[..., 2]
            th = pred[..., 3]
            obj_logit = pred[..., 4]
            cls_logits = pred[..., 5:]

            # Grid mesh
            gy, gx = torch.meshgrid(
                torch.arange(h_g, device=device, dtype=torch.float32),
                torch.arange(w_g, device=device, dtype=torch.float32),
                indexing="ij"
            )

            # Mathematical center & size transformation
            bx = (gx + torch.sigmoid(tx)) * stride
            by = (gy + torch.sigmoid(ty)) * stride
            bw = torch.exp(torch.clamp(tw, max=5.0)) * stride
            bh = torch.exp(torch.clamp(th, max=5.0)) * stride

            # Convert to [x1, y1, x2, y2]
            x1 = bx - bw / 2.0
            y1 = by - bh / 2.0
            x2 = bx + bw / 2.0
            y2 = by + bh / 2.0

            # Probabilities
            obj_prob = torch.sigmoid(obj_logit)
            cls_prob = torch.softmax(cls_logits, dim=-1)
            scores = obj_prob.unsqueeze(-1) * cls_prob  # (H_g, W_g, num_classes)

            # Flatten [x1, y1, x2, y2, bx, by] (bounding box + exact location point)
            boxes = torch.stack([x1, y1, x2, y2, bx, by], dim=-1).view(-1, 6)
            scores = scores.view(-1, cls_logits.shape[-1])

            max_scores, class_ids = torch.max(scores, dim=-1)
            mask = max_scores >= conf_thresh

            if mask.sum() > 0:
                all_boxes.append(boxes[mask])
                all_scores.append(max_scores[mask])
                all_classes.append(class_ids[mask])

        if len(all_boxes) > 0:
            batch_boxes = torch.cat(all_boxes, dim=0)
            batch_scores = torch.cat(all_scores, dim=0)
            batch_classes = torch.cat(all_classes, dim=0)
            batch_detections.append((batch_boxes, batch_scores, batch_classes))
        else:
            batch_detections.append((
                torch.empty((0, 6), device=pred_maps[0].device),
                torch.empty((0,), device=pred_maps[0].device),
                torch.empty((0,), dtype=torch.long, device=pred_maps[0].device)
            ))

    return batch_detections


def custom_nms(boxes, scores, classes, iou_thresh=0.45):
    """
    Mathematical Non-Maximum Suppression (NMS)
    Removes redundant overlapping bounding boxes using Intersection over Union (IoU).
    Preserves exact location points [x1, y1, x2, y2, bx, by, score, class_id].
    """
    if boxes.shape[0] == 0:
        return torch.empty((0, 8), device=boxes.device)

    keep_indices = []
    # Separate NMS per class to avoid suppressing adjacent different defects
    unique_classes = torch.unique(classes)

    for c in unique_classes:
        c_mask = (classes == c)
        c_indices = torch.nonzero(c_mask).squeeze(-1)
        c_boxes = boxes[c_indices]
        c_scores = scores[c_indices]

        # Sort descending by score
        order = torch.argsort(c_scores, descending=True)
        keep = []

        while order.numel() > 0:
            curr = order[0].item()
            keep.append(c_indices[curr].item())

            if order.numel() == 1:
                break

            curr_box = c_boxes[order[0], :4].unsqueeze(0)
            rest_boxes = c_boxes[order[1:], :4]

            # Mathematical IoU
            x1 = torch.max(curr_box[:, 0], rest_boxes[:, 0])
            y1 = torch.max(curr_box[:, 1], rest_boxes[:, 1])
            x2 = torch.min(curr_box[:, 2], rest_boxes[:, 2])
            y2 = torch.min(curr_box[:, 3], rest_boxes[:, 3])

            inter_w = torch.clamp(x2 - x1, min=0)
            inter_h = torch.clamp(y2 - y1, min=0)
            inter_area = inter_w * inter_h

            area_curr = (curr_box[:, 2] - curr_box[:, 0]) * (curr_box[:, 3] - curr_box[:, 1])
            area_rest = (rest_boxes[:, 2] - rest_boxes[:, 0]) * (rest_boxes[:, 3] - rest_boxes[:, 1])
            union_area = area_curr + area_rest - inter_area + 1e-7

            iou = inter_area / union_area
            remaining = torch.nonzero(iou <= iou_thresh).squeeze(-1)

            if remaining.numel() == 0:
                break
            order = order[remaining + 1]

        keep_indices.extend(keep)

    if len(keep_indices) == 0:
        return torch.empty((0, 8), device=boxes.device)

    keep_idx = torch.tensor(keep_indices, device=boxes.device, dtype=torch.long)
    final_boxes = boxes[keep_idx]  # (N, 6): [x1, y1, x2, y2, bx, by]
    final_scores = scores[keep_idx].unsqueeze(-1)
    final_classes = classes[keep_idx].unsqueeze(-1).float()

    return torch.cat([final_boxes, final_scores, final_classes], dim=-1)