microghost/model.py

"""
MicroGhost-Thermal: Model Module (V2)
========================================
Dual-branch architecture with late gated fusion for multimodal intrusion detection.

V2 Architecture:
- Dual independent GhostNet+MobileNetV2 branches (RGB + Thermal)
- EnergyGate: learned per-location modality weighting at Scale 2
- BiFusion Neck: bidirectional weighted feature pyramid (replaces FPN)
- ReliabilityClassifier: gate-aware classification (Visible vs Camouflaged)
- AuxSegHead: training-only contrast loss head (zero deployment cost)
- 3 anchors/cell for adjacent person detection

V1 architecture preserved for backward compatibility.
"""

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

from config import (
    INPUT_SIZE, INPUT_CHANNELS, NUM_CLASSES, NUM_ANCHORS,
    # V1 constants (backward compat)
    STEM_CHANNELS, RGB_STEM_CHANNELS, THERMAL_STEM_CHANNELS,
    SCALE1_CHANNELS, SCALE2_CHANNELS, SCALE3_CHANNELS,
    FPN_CHANNELS, CLASSIFIER_HIDDEN_DIM, EXPAND_RATIO,
    # V2 constants
    V2_STEM_CHANNELS, V2_SCALE1_CHANNELS, V2_SCALE2_CHANNELS,
    V2_SCALE3_CHANNELS, V2_BIFUSION_CHANNELS, V2_CLASSIFIER_HIDDEN_DIM,
    V2_EXPAND_RATIO,
    ESP32_S3,
)


# ============================================================================
# 1. CORE BUILDING BLOCKS (shared V1 + V2)
# ============================================================================

class GhostModule(nn.Module):
    """
    Ghost Module: Generates feature maps using cheap linear operations.
    Reduces computation by ~2x while maintaining representational capacity.
    """

    def __init__(self, in_channels, out_channels, kernel_size=1,
                 ratio=2, dw_kernel=3, stride=1, relu=True):
        super().__init__()
        self.out_channels = out_channels
        init_channels = math.ceil(out_channels / ratio)
        new_channels = init_channels * (ratio - 1)

        self.primary_conv = nn.Conv2d(in_channels, init_channels, kernel_size, stride,
                                      kernel_size // 2, bias=False)
        self.primary_bn = nn.BatchNorm2d(init_channels)
        
        # RepGhost parallel 1x1 for richer gradients during training
        self.rep_1x1 = nn.Sequential(
            nn.Conv2d(in_channels, init_channels, 1, stride, 0, bias=False),
            nn.BatchNorm2d(init_channels)
        )
        
        self.primary_act = nn.ReLU6(inplace=True) if relu else nn.Identity()

        self.cheap_operation = nn.Sequential(
            nn.Conv2d(init_channels, new_channels, dw_kernel, 1,
                      dw_kernel // 2, groups=init_channels, bias=False),
            nn.BatchNorm2d(new_channels),
            nn.ReLU6(inplace=True) if relu else nn.Identity(),
        )

    def forward(self, x):
        x1 = self.primary_bn(self.primary_conv(x))
        if self.training and hasattr(self, 'rep_1x1'):
            x1 = x1 + self.rep_1x1(x)
        x1 = self.primary_act(x1)
        x2 = self.cheap_operation(x1)
        out = torch.cat([x1, x2], dim=1)
        return out[:, :self.out_channels, :, :]


class GhostBottleneck(nn.Module):
    """
    Ghost Bottleneck: Efficient bottleneck using Ghost modules.
    Structure: Ghost (expansion) → DW Conv → Ghost (projection, linear) → Residual
    """

    def __init__(self, in_channels, mid_channels, out_channels,
                 dw_kernel=3, stride=1):
        super().__init__()
        self.stride = stride

        self.ghost1 = GhostModule(in_channels, mid_channels, relu=True)

        if stride > 1:
            self.conv_dw = nn.Sequential(
                nn.Conv2d(mid_channels, mid_channels, dw_kernel, stride,
                          dw_kernel // 2, groups=mid_channels, bias=False),
                nn.BatchNorm2d(mid_channels),
            )
        else:
            self.conv_dw = nn.Identity()

        self.ghost2 = GhostModule(mid_channels, out_channels, relu=False)

        if in_channels != out_channels or stride > 1:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, in_channels, dw_kernel, stride,
                          dw_kernel // 2, groups=in_channels, bias=False),
                nn.BatchNorm2d(in_channels),
                nn.Conv2d(in_channels, out_channels, 1, 1, 0, bias=False),
                nn.BatchNorm2d(out_channels),
            )
        else:
            self.shortcut = nn.Identity()

    def forward(self, x):
        residual = self.shortcut(x)
        x = self.ghost1(x)
        x = self.conv_dw(x)
        x = self.ghost2(x)
        return x + residual


class InvertedResidual(nn.Module):
    """
    MobileNetV2 Inverted Residual Block with Linear Bottleneck.
    Structure: Expand → Depthwise → Project (Linear, no ReLU at end)
    """

    def __init__(self, in_channels, out_channels, stride=1,
                 expand_ratio=None):
        super().__init__()
        expand_ratio = expand_ratio or EXPAND_RATIO
        self.stride = stride
        self.use_residual = (stride == 1 and in_channels == out_channels)
        hidden_dim = int(in_channels * expand_ratio)

        layers = []

        if expand_ratio != 1:
            layers.extend([
                nn.Conv2d(in_channels, hidden_dim, 1, 1, 0, bias=False),
                nn.BatchNorm2d(hidden_dim),
                nn.ReLU6(inplace=True),
            ])

        layers.extend([
            nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1,
                      groups=hidden_dim, bias=False),
            nn.BatchNorm2d(hidden_dim),
            nn.ReLU6(inplace=True),
        ])

        layers.extend([
            nn.Conv2d(hidden_dim, out_channels, 1, 1, 0, bias=False),
            nn.BatchNorm2d(out_channels),
        ])

        self.conv = nn.Sequential(*layers)

    def forward(self, x):
        if self.use_residual:
            return x + self.conv(x)
        return self.conv(x)


# ============================================================================
# 2. V1 ARCHITECTURE (kept for backward compatibility)
# ============================================================================

class LightweightFPN(nn.Module):
    """V1 Lightweight Feature Pyramid Network for multi-scale fusion."""

    def __init__(self, in_channels_s2, in_channels_s3, out_channels=None):
        super().__init__()
        out_channels = out_channels or FPN_CHANNELS

        self.lateral_s2 = nn.Conv2d(in_channels_s2, out_channels, 1,
                                    bias=False)
        self.lateral_s3 = nn.Conv2d(in_channels_s3, out_channels, 1,
                                    bias=False)

        self.smooth_s2 = nn.Sequential(
            nn.Conv2d(out_channels, out_channels, 3, 1, 1,
                      groups=out_channels, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU6(inplace=True),
            nn.Conv2d(out_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
        )

        self.smooth_s3 = nn.Sequential(
            nn.Conv2d(out_channels, out_channels, 3, 1, 1,
                      groups=out_channels, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU6(inplace=True),
            nn.Conv2d(out_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
        )

    def forward(self, feat_s2, feat_s3):
        lat_s2 = self.lateral_s2(feat_s2)
        lat_s3 = self.lateral_s3(feat_s3)
        upsampled_s3 = F.interpolate(lat_s3, size=lat_s2.shape[2:],
                                     mode='nearest')
        p2 = self.smooth_s2(lat_s2 + upsampled_s3)
        p3 = self.smooth_s3(lat_s3)
        return p2, p3


class SSDLiteHead(nn.Module):
    """SSDLite Detection Head using depthwise-separable convolutions."""

    def __init__(self, in_channels, num_anchors=None):
        super().__init__()
        num_anchors = num_anchors or NUM_ANCHORS

        self.feature = nn.Sequential(
            nn.Conv2d(in_channels, in_channels, 3, 1, 1,
                      groups=in_channels, bias=False),
            nn.BatchNorm2d(in_channels),
            nn.ReLU6(inplace=True),
            nn.Conv2d(in_channels, in_channels, 1, bias=False),
            nn.BatchNorm2d(in_channels),
            nn.ReLU6(inplace=True),
        )

        self.bbox_head = nn.Conv2d(in_channels, num_anchors * 4, 1)
        self.obj_head = nn.Conv2d(in_channels, num_anchors, 1)

    def forward(self, x):
        feat = self.feature(x)
        bbox = self.bbox_head(feat)
        obj = self.obj_head(feat)
        return bbox, obj


class IntrusionClassifier(nn.Module):
    """V1 Multiclass intrusion classifier with objectness-weighted attention."""

    def __init__(self, in_channels, num_classes=None, hidden_dim=None):
        super().__init__()
        num_classes = num_classes or NUM_CLASSES
        hidden_dim = hidden_dim or CLASSIFIER_HIDDEN_DIM

        self.modality_attention = nn.Sequential(
            nn.Linear(in_channels * 2, hidden_dim // 2),
            nn.ReLU6(inplace=True),
            nn.Linear(hidden_dim // 2, in_channels * 2),
            nn.Sigmoid()
        )

        self.classifier = nn.Sequential(
            nn.Linear(in_channels * 2, hidden_dim),
            nn.ReLU6(inplace=True),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, num_classes),
        )

    def forward(self, feat_p2, feat_p3, obj_p2, obj_p3):
        attn_p2 = torch.sigmoid(
            obj_p2.max(dim=1, keepdim=True)[0]
        )
        attn_p3 = torch.sigmoid(
            obj_p3.max(dim=1, keepdim=True)[0]
        )

        p2_w = (feat_p2 * attn_p2).sum(dim=[2, 3]) / \
               (attn_p2.sum(dim=[2, 3]) + 1e-6)
        p3_w = (feat_p3 * attn_p3).sum(dim=[2, 3]) / \
               (attn_p3.sum(dim=[2, 3]) + 1e-6)

        combined = torch.cat([p2_w, p3_w], dim=1)

        attn_weights = self.modality_attention(combined)
        attended_features = combined * attn_weights

        return self.classifier(attended_features)


class MicroGhostThermal(nn.Module):
    """V1 MicroGhost-Thermal: Early fusion, shared backbone, FPN."""

    def __init__(self, num_classes=None, num_anchors=None,
                 input_size=None, classifier_hidden_dim=None):
        super().__init__()
        num_classes = num_classes or NUM_CLASSES
        num_anchors = num_anchors or NUM_ANCHORS
        input_size = input_size or INPUT_SIZE
        classifier_hidden_dim = classifier_hidden_dim or CLASSIFIER_HIDDEN_DIM

        self.input_size = input_size
        if isinstance(input_size, tuple):
            self.input_h, self.input_w = input_size
        else:
            self.input_h, self.input_w = input_size, input_size
        self.num_classes = num_classes
        self.classifier_hidden_dim = classifier_hidden_dim

        # Dual Stem
        self.rgb_stem = nn.Sequential(
            nn.Conv2d(3, 8, 3, 2, 1, bias=False),
            nn.BatchNorm2d(8),
            nn.ReLU6(inplace=True),
            GhostModule(8, RGB_STEM_CHANNELS, kernel_size=1, stride=1),
        )

        self.thermal_stem = nn.Sequential(
            nn.Conv2d(1, 8, 3, 2, 1, bias=False),
            nn.BatchNorm2d(8),
            nn.ReLU6(inplace=True),
            GhostModule(8, THERMAL_STEM_CHANNELS, kernel_size=1, stride=1),
        )

        # Shared backbone
        self.scale1 = nn.Sequential(
            GhostBottleneck(STEM_CHANNELS, STEM_CHANNELS * 2,
                            SCALE1_CHANNELS, stride=2),
            GhostBottleneck(SCALE1_CHANNELS, SCALE1_CHANNELS * 2,
                            SCALE1_CHANNELS, stride=1),
        )

        self.scale2 = nn.Sequential(
            InvertedResidual(SCALE1_CHANNELS, SCALE2_CHANNELS,
                             stride=2, expand_ratio=EXPAND_RATIO),
            InvertedResidual(SCALE2_CHANNELS, SCALE2_CHANNELS,
                             stride=1, expand_ratio=EXPAND_RATIO),
        )

        self.scale3 = nn.Sequential(
            InvertedResidual(SCALE2_CHANNELS, SCALE3_CHANNELS,
                             stride=2, expand_ratio=EXPAND_RATIO),
            InvertedResidual(SCALE3_CHANNELS, SCALE3_CHANNELS,
                             stride=1, expand_ratio=EXPAND_RATIO),
        )

        self.fpn = LightweightFPN(
            in_channels_s2=SCALE2_CHANNELS,
            in_channels_s3=SCALE3_CHANNELS,
            out_channels=FPN_CHANNELS,
        )

        self.head_small = SSDLiteHead(FPN_CHANNELS, num_anchors=num_anchors)
        self.head_large = SSDLiteHead(FPN_CHANNELS, num_anchors=num_anchors)

        self.classifier = IntrusionClassifier(
            in_channels=FPN_CHANNELS,
            num_classes=num_classes,
            hidden_dim=classifier_hidden_dim,
        )

        self._initialize_weights()

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

    def forward(self, x):
        x_rgb = x[:, :3, :, :]
        x_thermal = x[:, 3:, :, :]

        feat_rgb = self.rgb_stem(x_rgb)
        feat_thermal = self.thermal_stem(x_thermal)
        feat_fused = torch.cat([feat_rgb, feat_thermal], dim=1)

        s1 = self.scale1(feat_fused)
        s2 = self.scale2(s1)
        s3 = self.scale3(s2)

        p2, p3 = self.fpn(s2, s3)

        bbox_small, obj_small = self.head_small(p2)
        bbox_large, obj_large = self.head_large(p3)

        label = self.classifier(p2, p3, obj_small, obj_large)

        return {
            'bbox_small': bbox_small,
            'obj_small': obj_small,
            'bbox_large': bbox_large,
            'obj_large': obj_large,
            'label': label,
        }


# ============================================================================
# 3. V2 NEW MODULES
# ============================================================================

class EnergyGate(nn.Module):
    """
    Per-location gating between RGB and Thermal at Scale 2.

    Computes learned energy projections and applies softmax across the two
    branches at each spatial location. This lets the network suppress
    whichever branch has unreliable content (e.g., dark RGB at night,
    hot car bonnet in thermal).

    ~30 parameters. Applied BEFORE BiFusion Neck.
    """

    def __init__(self, channels):
        super().__init__()
        self.proj_rgb = nn.Conv2d(channels, 1, 1, bias=True)
        self.proj_thm = nn.Conv2d(channels, 1, 1, bias=True)
        self.temperature = nn.Parameter(torch.ones(1) * 2.0)

    def forward(self, feat_rgb, feat_thm):
        e_rgb = self.proj_rgb(feat_rgb)          # (B, 1, H, W)
        e_thm = self.proj_thm(feat_thm)          # (B, 1, H, W)
        
        temp = torch.clamp(self.temperature, min=0.5, max=5.0)
        
        weights = torch.softmax(
            torch.stack([e_rgb, e_thm], dim=1) / temp,  # (B, 2, 1, H, W)
            dim=1
        )
        w_rgb = weights[:, 0]                    # (B, 1, H, W)
        w_thm = weights[:, 1]                    # (B, 1, H, W)
        fused = w_rgb * feat_rgb + w_thm * feat_thm
        return fused, w_rgb, w_thm


class BiFusionNeck(nn.Module):
    """
    Bidirectional weighted feature pyramid (replaces LightweightFPN).

    Receives S2 fused features (from EnergyGate) and S3 features from
    both branches separately. Uses learned normalized weights (BiFPN-style)
    for top-down and bottom-up passes.

    Inputs:
        fused_s2:    (B, s2_ch, 16, 20) — gated EnergyGate output
        feat_rgb_s3: (B, s3_ch, 8, 10)
        feat_thm_s3: (B, s3_ch, 8, 10)

    Outputs:
        p2: (B, out_ch, 16, 20) — small/distant target features
        p3: (B, out_ch, 8, 10)  — large/close target features
    """

    def __init__(self, s2_ch=None, s3_ch=None, out_ch=None):
        super().__init__()
        s2_ch = s2_ch or V2_SCALE2_CHANNELS
        s3_ch = s3_ch or V2_SCALE3_CHANNELS
        out_ch = out_ch or V2_BIFUSION_CHANNELS

        # Lateral projections to unified channel count
        self.lat_s2 = nn.Conv2d(s2_ch, out_ch, 1, bias=False)
        self.lat_rgb = nn.Conv2d(s3_ch, out_ch, 1, bias=False)
        self.lat_thm = nn.Conv2d(s3_ch, out_ch, 1, bias=False)

        # Learned BiFPN weights (softmax-normalized)
        # Top-down: P3 = w1*rgb_s3 + w2*thm_s3
        self.w_td = nn.Parameter(torch.zeros(2))

        # Bottom-up P2: w3*s2 + w4*P3_upsampled
        self.w_bu = nn.Parameter(torch.zeros(2))

        # DW-separable refinement convolutions
        self.refine_p3 = self._dw_sep(out_ch, out_ch)
        self.refine_p2 = self._dw_sep(out_ch, out_ch)

    def _dw_sep(self, in_ch, out_ch):
        return nn.Sequential(
            nn.Conv2d(in_ch, in_ch, 3, 1, 1, groups=in_ch, bias=False),
            nn.BatchNorm2d(in_ch),
            nn.ReLU6(inplace=True),
            nn.Conv2d(in_ch, out_ch, 1, bias=False),
            nn.BatchNorm2d(out_ch),
        )

    def forward(self, fused_s2, feat_rgb_s3, feat_thm_s3):
        eps = 1e-4

        # Project to unified channels
        lat_s2 = self.lat_s2(fused_s2)          # (B, CH, 16, 20)
        lat_rgb = self.lat_rgb(feat_rgb_s3)      # (B, CH, 8, 10)
        lat_thm = self.lat_thm(feat_thm_s3)      # (B, CH, 8, 10)

        # ── Top-down: fuse S3 from both branches ──
        w_td = F.softplus(self.w_td) + eps
        w_td = w_td / w_td.sum()
        p3_td = self.refine_p3(
            w_td[0] * lat_rgb + w_td[1] * lat_thm
        )                                        # (B, CH, 8, 10)

        # ── Bottom-up: upsample P3 and merge with S2 ──
        p3_up = F.interpolate(p3_td, size=lat_s2.shape[2:], mode='nearest')
        w_bu = F.softplus(self.w_bu) + eps
        w_bu = w_bu / w_bu.sum()
        p2_out = self.refine_p2(
            w_bu[0] * lat_s2 + w_bu[1] * p3_up
        )                                        # (B, CH, 16, 20)

        return p2_out, p3_td


class ReliabilityClassifier(nn.Module):
    """
    Gate-aware classifier for V2 (replaces IntrusionClassifier).

    Receives the EnergyGate weights as auxiliary input, enabling
    Visible vs Camouflaged classification based on which modality
    was dominant at detection time.

    - w_thm >> w_rgb → thermal dominant → likely Person_Camouflaged
    - w_rgb ≈ w_thm → both agree → Person_Visible
    """

    def __init__(self, in_channels=None, num_classes=None, hidden_dim=None):
        super().__init__()
        in_channels = in_channels or V2_BIFUSION_CHANNELS
        num_classes = num_classes or NUM_CLASSES
        hidden_dim = hidden_dim or V2_CLASSIFIER_HIDDEN_DIM

        # Modality-gate-aware attention (+2 for gate weight scalars)
        self.modality_gate = nn.Sequential(
            nn.Linear(in_channels * 2 + 2, hidden_dim // 2),
            nn.ReLU6(inplace=True),
            nn.Linear(hidden_dim // 2, in_channels * 2),
            nn.Sigmoid(),
        )
        self.classifier = nn.Sequential(
            nn.Linear(in_channels * 2, hidden_dim),
            nn.ReLU6(inplace=True),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, num_classes + 1), # +1 for IoU regression
        )

    def forward(self, feat_p2, feat_p3, obj_p2, obj_p3, w_rgb, w_thm):
        # Objectness-weighted spatial pooling
        attn_p2 = torch.sigmoid(obj_p2.max(dim=1, keepdim=True)[0])
        attn_p3 = torch.sigmoid(obj_p3.max(dim=1, keepdim=True)[0])
        p2_w = (feat_p2 * attn_p2).sum([2, 3]) / (attn_p2.sum([2, 3]) + 1e-6)
        p3_w = (feat_p3 * attn_p3).sum([2, 3]) / (attn_p3.sum([2, 3]) + 1e-6)
        combined = torch.cat([p2_w, p3_w], dim=1)   # (B, 2C)

        # Gate summary: mean weight across spatial dimensions → (B, 2)
        gate_rgb_mean = w_rgb.mean(dim=[1, 2, 3])   # (B,)
        gate_thm_mean = w_thm.mean(dim=[1, 2, 3])   # (B,)
        gate_summary = torch.stack([gate_rgb_mean, gate_thm_mean], dim=1)  # (B, 2)

        # Attend features using gate-aware attention
        gate_input = torch.cat([combined, gate_summary], dim=1)  # (B, 2C+2)
        attn_weights = self.modality_gate(gate_input)  # (B, 2C)
        attended = combined * attn_weights

        return self.classifier(attended)


class AuxSegHead(nn.Module):
    """
    Training-only auxiliary segmentation head for TFDet-style contrast loss.
    Removed at export. Zero deployment cost.

    Produces per-spatial-location person/background logits from p2 features.
    """

    def __init__(self, in_ch=None):
        super().__init__()
        in_ch = in_ch or V2_BIFUSION_CHANNELS
        self.proj = nn.Conv2d(in_ch, 1, 1)

    def forward(self, feat_p2):
        return self.proj(feat_p2)   # (B, 1, H, W)


# ============================================================================
# 4. V2 COMPLETE MODEL: MicroGhostV2
# ============================================================================

class MicroGhostV2(nn.Module):
    """
    MicroGhost-V2: Asynchronous Dual-Branch Architecture.

    Key changes vs V1:
    - Parallel independent RGB + Thermal branches (no shared weights)
    - EnergyGate at Scale 2 for learned modality weighting
    - BiFusion Neck (bidirectional weighted pyramid, replaces FPN)
    - ReliabilityClassifier with gate-weight awareness
    - AuxSegHead for training-only contrast loss
    - 3 anchors per cell
    - Graceful camera failure (either branch can operate independently)

    Architecture:
    ┌────────────────────────────────────────────────────────────┐
    │  RGB (3ch)            Thermal (1ch)                        │
    │     │                      │                               │
    │  RGB Stem(16)          Thm Stem(16)                        │
    │     │                      │                               │
    │  RGB Scale1(24)        Thm Scale1(24)                      │
    │     │                      │                               │
    │  RGB Scale2(32) ──► EnergyGate ◄── Thm Scale2(32)          │
    │                       │ fused(32)                           │
    │  RGB Scale3(48) ──► BiFusion Neck ◄── Thm Scale3(48)       │
    │                    ┌────┴────┐                             │
    │                 p2(48)    p3(48)                            │
    │              SmallHead  LargeHead                          │
    │                    └────┬────┘                             │
    │              ReliabilityClassifier                         │
    └────────────────────────────────────────────────────────────┘
    """

    def __init__(self, num_classes=None, num_anchors=None,
                 input_size=None, classifier_hidden_dim=None,
                 training_mode=True):
        super().__init__()
        num_classes = num_classes or NUM_CLASSES
        num_anchors = num_anchors or NUM_ANCHORS
        input_size = input_size or INPUT_SIZE
        classifier_hidden_dim = classifier_hidden_dim or V2_CLASSIFIER_HIDDEN_DIM

        self.input_size = input_size
        if isinstance(input_size, tuple):
            self.input_h, self.input_w = input_size
        else:
            self.input_h, self.input_w = input_size, input_size
        self.num_classes = num_classes
        self.classifier_hidden_dim = classifier_hidden_dim
        self.training_mode = training_mode

        S = V2_STEM_CHANNELS
        S1 = V2_SCALE1_CHANNELS
        S2 = V2_SCALE2_CHANNELS
        S3 = V2_SCALE3_CHANNELS
        E = V2_EXPAND_RATIO

        # ========== RGB BRANCH (fully independent) ==========
        self.rgb_stem = nn.Sequential(
            nn.Conv2d(3, 8, 3, 2, 1, bias=False),
            nn.BatchNorm2d(8),
            nn.ReLU6(inplace=True),
            GhostModule(8, S, kernel_size=1, stride=1),
        )
        self.rgb_scale1 = nn.Sequential(
            GhostBottleneck(S, S * 2, S1, stride=2),
            GhostBottleneck(S1, S1 * 2, S1, stride=1),
        )
        self.rgb_scale2 = nn.Sequential(
            InvertedResidual(S1, S2, stride=2, expand_ratio=E),
            InvertedResidual(S2, S2, stride=1, expand_ratio=E),
        )
        self.rgb_scale3 = nn.Sequential(
            InvertedResidual(S2, S3, stride=2, expand_ratio=E),
            InvertedResidual(S3, S3, stride=1, expand_ratio=E),
        )

        # ========== THERMAL BRANCH (fully independent) ==========
        self.thm_stem = nn.Sequential(
            nn.Conv2d(1, 8, 3, 2, 1, bias=False),
            nn.BatchNorm2d(8),
            nn.ReLU6(inplace=True),
            GhostModule(8, S, kernel_size=1, stride=1),
        )
        self.thm_scale1 = nn.Sequential(
            GhostBottleneck(S, S * 2, S1, stride=2),
            GhostBottleneck(S1, S1 * 2, S1, stride=1),
        )
        self.thm_scale2 = nn.Sequential(
            InvertedResidual(S1, S2, stride=2, expand_ratio=E),
            InvertedResidual(S2, S2, stride=1, expand_ratio=E),
        )
        self.thm_scale3 = nn.Sequential(
            InvertedResidual(S2, S3, stride=2, expand_ratio=E),
            InvertedResidual(S3, S3, stride=1, expand_ratio=E),
        )

        # ========== ENERGY GATE (at Scale 2 output) ==========
        self.energy_gate = EnergyGate(channels=S2)

        # ========== BIFUSION NECK (replaces FPN) ==========
        self.bifusion_neck = BiFusionNeck(
            s2_ch=S2, s3_ch=S3, out_ch=V2_BIFUSION_CHANNELS,
        )

        # ========== DETECTION HEADS (3 anchors each) ==========
        self.head_small = SSDLiteHead(V2_BIFUSION_CHANNELS, num_anchors=num_anchors)
        self.head_large = SSDLiteHead(V2_BIFUSION_CHANNELS, num_anchors=num_anchors)

        # ========== RELIABILITY CLASSIFIER ==========
        self.classifier = ReliabilityClassifier(
            in_channels=V2_BIFUSION_CHANNELS,
            num_classes=num_classes,
            hidden_dim=classifier_hidden_dim,
        )

        # ========== AUX SEGMENTATION HEAD (training only) ==========
        self.aux_seg_head = AuxSegHead(in_ch=V2_BIFUSION_CHANNELS)

        # Initialize weights
        self._initialize_weights()

    def _initialize_weights(self):
        """Kaiming initialization for better convergence."""
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out',
                                        nonlinearity='relu')
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)

    def forward(self, x):
        """
        Forward pass for V2 dual-branch architecture.

        Args:
            x: (B, 4, H, W) float tensor (channels 0:3 = RGB, channel 3: = Thermal)

        Returns:
            dict with detection outputs + gate weights + aux seg logits
        """
        x_rgb = x[:, :3]       # (B, 3, H, W)
        x_thm = x[:, 3:]       # (B, 1, H, W)

        # Modality masking to prevent BatchNorm shift artifacts on empty inputs
        rgb_present = (x_rgb.abs().mean(dim=[1,2,3], keepdim=True) > 1e-5).float()
        thm_present = (x_thm.abs().mean(dim=[1,2,3], keepdim=True) > 1e-5).float()

        # === RGB Branch (fully independent) ===
        feat_rgb = self.rgb_stem(x_rgb)             # (B, 16, 64, 80)
        feat_rgb = self.rgb_scale1(feat_rgb)        # (B, 24, 32, 40)
        feat_rgb_s2 = self.rgb_scale2(feat_rgb)     # (B, 32, 16, 20)
        feat_rgb_s3 = self.rgb_scale3(feat_rgb_s2)  # (B, 48, 8, 10)
        
        feat_rgb_s2 = feat_rgb_s2 * rgb_present
        feat_rgb_s3 = feat_rgb_s3 * rgb_present

        # === Thermal Branch (fully independent) ===
        feat_thm = self.thm_stem(x_thm)             # (B, 16, 64, 80)
        feat_thm = self.thm_scale1(feat_thm)        # (B, 24, 32, 40)
        feat_thm_s2 = self.thm_scale2(feat_thm)     # (B, 32, 16, 20)
        feat_thm_s3 = self.thm_scale3(feat_thm_s2)  # (B, 48, 8, 10)
        
        feat_thm_s2 = feat_thm_s2 * thm_present
        feat_thm_s3 = feat_thm_s3 * thm_present

        # === Energy Gate (learned modality weighting at S2) ===
        fused_s2, w_rgb, w_thm = self.energy_gate(feat_rgb_s2, feat_thm_s2)

        # === BiFusion Neck (replaces FPN) ===
        p2, p3 = self.bifusion_neck(fused_s2, feat_rgb_s3, feat_thm_s3)

        # === Detection Heads (3 anchors each) ===
        bbox_small, obj_small = self.head_small(p2)
        bbox_large, obj_large = self.head_large(p3)

        # === Reliability Classifier ===
        label = self.classifier(p2, p3, obj_small, obj_large, w_rgb, w_thm)

        result = {
            'bbox_small': bbox_small,
            'obj_small':  obj_small,
            'bbox_large': bbox_large,
            'obj_large':  obj_large,
            'label':      label,
            'w_rgb':      w_rgb,
            'w_thm':      w_thm,
        }

        # Aux seg head (training only — removed at export)
        if self.training_mode and self.training:
            result['aux_seg'] = self.aux_seg_head(p2)

        return result

    def freeze_early_layers(self):
        """Freeze stems and Scale 1 for Phase 3 fine-tuning."""
        for name, param in self.named_parameters():
            if any(prefix in name for prefix in
                   ['rgb_stem', 'thm_stem', 'rgb_scale1', 'thm_scale1']):
                param.requires_grad = False

    def unfreeze_all(self):
        """Unfreeze all parameters for Phase 4 polish."""
        for param in self.parameters():
            param.requires_grad = True


# ============================================================================
# 5. MODEL ANALYSIS UTILITIES
# ============================================================================

def count_parameters(model):
    """Count trainable parameters."""
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


def count_all_parameters(model):
    """Count all parameters (including frozen)."""
    return sum(p.numel() for p in model.parameters())


def estimate_model_size(model):
    """Estimate model size in different quantization formats."""
    param_count = count_all_parameters(model)
    fp32_mb = param_count * 4 / (1024 * 1024)
    int8_kb = param_count * 1 / 1024
    return param_count, fp32_mb, int8_kb


def estimate_peak_sram(model, input_size=None, batch_size=1):
    """Estimate peak SRAM usage during inference on ESP32-S3."""
    input_size = input_size or INPUT_SIZE
    activations = []

    def hook_fn(module, inp, output):
        if isinstance(output, torch.Tensor):
            activations.append(output.numel())

    hooks = []
    for module in model.modules():
        if isinstance(module, (nn.Conv2d, nn.BatchNorm2d, nn.ReLU6)):
            hooks.append(module.register_forward_hook(hook_fn))

    model.eval()
    with torch.no_grad():
        h, w = input_size if isinstance(input_size, tuple) else (input_size, input_size)
        x = torch.randn(batch_size, INPUT_CHANNELS, h, w)
        _ = model(x)

    for hook in hooks:
        hook.remove()

    input_bytes = batch_size * INPUT_CHANNELS * h * w
    max_activation_int8 = max(activations) if activations else 0

    return {
        'input_buffer_kb': input_bytes / 1024,
        'peak_activation_fp32_kb': max(activations) * 4 / 1024 if activations else 0,
        'peak_activation_int8_kb': max_activation_int8 / 1024,
        'total_arena_int8_kb': (input_bytes + max_activation_int8) / 1024,
        'fits_esp32_s3': (input_bytes + max_activation_int8) < ESP32_S3['max_arena_sram_kb'] * 1024,
    }


def print_model_analysis(model):
    """Print comprehensive model analysis for V1 or V2."""
    param_count, fp32_mb, int8_kb = estimate_model_size(model)
    sram = estimate_peak_sram(model)

    is_v2 = isinstance(model, MicroGhostV2)
    model_name = "MicroGhost-V2" if is_v2 else "MicroGhost-V1"

    print(f"\n[OK] {model_name}:")
    print(f"   Parameters:       {param_count:,}")
    print(f"   Trainable:        {count_parameters(model):,}")
    print(f"   Est. Size (FP32): {param_count * 4 / 1024:,.1f} KB")
    print(f"   Est. Size (FP16): {param_count * 2 / 1024:,.1f} KB")
    print(f"   Est. Size (INT8): {int8_kb:,.1f} KB")
    print()
    print(f"  Input buffer:      {sram['input_buffer_kb']:>10.1f} KB")
    print(f"  Peak act (FP32):   {sram['peak_activation_fp32_kb']:>10.1f} KB")
    print(f"  Peak act (INT8):   {sram['peak_activation_int8_kb']:>10.1f} KB")
    print(f"  Total arena INT8:  {sram['total_arena_int8_kb']:>10.1f} KB")
    print(f"  Fits ESP32-S3:     {'OK' if sram['fits_esp32_s3'] else 'FAIL'} "
          f"(limit: {ESP32_S3['max_arena_sram_kb']}KB)")

    # Layer-by-layer breakdown
    print(f"\n  {'Layer':<35} {'Params':>12} {'Size (KB)':>10}")
    print("  " + "-" * 59)

    if is_v2:
        components = [
            ('RGB Stem', model.rgb_stem),
            ('RGB Scale 1 (Ghost)', model.rgb_scale1),
            ('RGB Scale 2 (InvRes)', model.rgb_scale2),
            ('RGB Scale 3 (InvRes)', model.rgb_scale3),
            ('Thm Stem', model.thm_stem),
            ('Thm Scale 1 (Ghost)', model.thm_scale1),
            ('Thm Scale 2 (InvRes)', model.thm_scale2),
            ('Thm Scale 3 (InvRes)', model.thm_scale3),
            ('Energy Gate', model.energy_gate),
            ('BiFusion Neck', model.bifusion_neck),
            ('Head Small', model.head_small),
            ('Head Large', model.head_large),
            ('Classifier', model.classifier),
            ('Aux Seg Head (train only)', model.aux_seg_head),
        ]
    else:
        components = [
            ('RGB Stem', model.rgb_stem),
            ('Thermal Stem', model.thermal_stem),
            ('Scale 1 (Ghost)', model.scale1),
            ('Scale 2 (InvRes)', model.scale2),
            ('Scale 3 (InvRes)', model.scale3),
            ('FPN', model.fpn),
            ('Head Small', model.head_small),
            ('Head Large', model.head_large),
            ('Classifier', model.classifier),
        ]

    total = 0
    for name, module in components:
        params = sum(p.numel() for p in module.parameters())
        kb = params * 4 / 1024
        total += params
        print(f"  {name:<35} {params:>12,} {kb:>8.1f} KB")

    print("  " + "-" * 59)
    print(f"  {'TOTAL':<35} {total:>12,} {total * 4 / 1024:>8.1f} KB")


# ============================================================================
# TEST
# ============================================================================

if __name__ == '__main__':
    print("=" * 60)
    print("  Model Module — V2 Self Test")
    print("=" * 60)

    # Test V2 model
    model = MicroGhostV2()
    print_model_analysis(model)

    # Test forward pass (training mode)
    model.train()
    h, w = INPUT_SIZE if isinstance(INPUT_SIZE, tuple) else (INPUT_SIZE, INPUT_SIZE)
    dummy = torch.randn(2, INPUT_CHANNELS, h, w)
    with torch.no_grad():
        outputs = model(dummy)

    print(f"\n[OK] V2 Forward Pass (training):")
    for key, val in outputs.items():
        print(f"    {key}: {val.shape}")

    # Test forward pass (eval mode — no aux_seg)
    model.eval()
    with torch.no_grad():
        outputs_eval = model(dummy)
    print(f"\n[OK] V2 Forward Pass (eval):")
    for key, val in outputs_eval.items():
        print(f"    {key}: {val.shape}")
    assert 'aux_seg' not in outputs_eval, "aux_seg should not be in eval outputs"

    # Test single-modality (CMM-RXTO: thermal zeroed)
    dummy_rxto = dummy.clone()
    dummy_rxto[:, 3:] = 0.0
    model.eval()
    with torch.no_grad():
        outputs_rxto = model(dummy_rxto)
    print(f"\n[OK] CMM-RXTO (thermal zeroed) — forward pass OK")

    # Test single-modality (CMM-ROTX: RGB zeroed)
    dummy_rotx = dummy.clone()
    dummy_rotx[:, :3] = 0.0
    with torch.no_grad():
        outputs_rotx = model(dummy_rotx)
    print(f"[OK] CMM-ROTX (RGB zeroed) — forward pass OK")

    # Test freeze/unfreeze
    model.freeze_early_layers()
    trainable_after_freeze = count_parameters(model)
    model.unfreeze_all()
    trainable_after_unfreeze = count_parameters(model)
    print(f"\n[OK] Freeze test: {trainable_after_freeze:,} trainable (frozen) -> "
          f"{trainable_after_unfreeze:,} trainable (unfrozen)")

    # Verify gate weights sum to ~1
    w_rgb_mean = outputs_eval['w_rgb'].mean().item()
    w_thm_mean = outputs_eval['w_thm'].mean().item()
    print(f"\n[OK] Gate weights: w_rgb={w_rgb_mean:.4f}, w_thm={w_thm_mean:.4f}, "
          f"sum={w_rgb_mean + w_thm_mean:.4f}")

    print("\n[OK] All V2 model tests passed!")