"""Camera-Adaptive Feature Modulation for fixed-camera detection setups.

Since this dataset uses 4 fixed cameras (EastLeft, EastRight, WestLeft, WestRight),
each camera has consistent perspective/viewpoint. CameraFiLM adds lightweight
camera-conditioned feature modulation to help the detection head compensate
for camera-specific geometric biases.

Cost: 4×C×2 extra params per modulated layer (~negligible).
"""
import torch
import torch.nn as nn
from pathlib import Path


CAMERA_NAMES = ['EastLeft', 'EastRight', 'WestLeft', 'WestRight']


def extract_camera_ids(im_files):
    """Extract camera index from image file paths. Returns list of ints."""
    ids = []
    for f in im_files:
        fname = Path(f).stem
        found = False
        for i, cam in enumerate(CAMERA_NAMES):
            if cam in fname:
                ids.append(i)
                found = True
                break
        if not found:
            ids.append(0)  # default fallback
    return ids


class CameraFiLM(nn.Module):
    """Feature-wise Linear Modulation conditioned on camera identity."""

    def __init__(self, channels, num_cameras=4):
        super().__init__()
        self.scale = nn.Parameter(torch.zeros(num_cameras, channels))
        self.bias = nn.Parameter(torch.zeros(num_cameras, channels))

    def forward(self, x, camera_ids):
        """x: [B, C, H, W], camera_ids: LongTensor [B] or None"""
        if camera_ids is None:
            return x
        s = self.scale[camera_ids]  # [B, C]
        b = self.bias[camera_ids]   # [B, C]
        return x * (1.0 + s[:, :, None, None]) + b[:, :, None, None]


def _get_film_channels(detect_head):
    """Get input channel counts for each detection layer."""
    ch = []
    for i in range(detect_head.nl):
        # cv2[0] is Conv(input_ch, c2). Get input channels.
        ch.append(detect_head.cv2[i][0].conv.in_channels)
    return ch


def inject_camera_film(detect_head):
    """Inject CameraFiLM into the detection head's forward_head.

    The FiLM is applied to the FPN features BEFORE detection convolutions,
    modulating per-camera feature statistics. Camera IDs are read from
    detect_head.camera_ids (set before forward).

    Args:
        detect_head: model.model.model[-1] — the Detect module
    """
    assert hasattr(detect_head, 'nl'), 'Must be a Detect-like module'

    channels = _get_film_channels(detect_head)

    films = nn.ModuleList([
        CameraFiLM(c, len(CAMERA_NAMES))
        for c in channels
    ])
    detect_head.add_module('camera_films', films)
    detect_head.camera_ids = None

    _orig_forward_head = detect_head.forward_head

    def forward_head_with_film(x, **kwargs):
        cids = detect_head.camera_ids
        for i in range(detect_head.nl):
            x[i] = detect_head.camera_films[i](x[i], cids)
        return _orig_forward_head(x, **kwargs)

    detect_head.forward_head = forward_head_with_film
    return detect_head


def inject_camera_to_model(model):
    """Inject camera-adaptive FiLM and patch model.loss for camera-id flow.

    The flow: batch['im_file'] -> extract_camera_ids() -> detect.camera_ids
    -> forward_head_with_film() -> CameraFiLM -> detection convolutions
    """
    detect = model.model.model[-1]  # DetectionModel -> Sequential -> Detect
    inject_camera_film(detect)

    _orig_loss = model.loss

    def loss_with_camera(batch, preds=None):
        if 'im_file' in batch:
            cids = extract_camera_ids(batch['im_file'])
            detect.camera_ids = torch.tensor(cids, device=model.device)
        else:
            detect.camera_ids = None
        return _orig_loss(batch, preds)

    import types
    model.loss = types.MethodType(loss_with_camera, model)
    return model