"""Analytical detection head — zero gradient steps.

All weights are computed from closed-form least-squares on cached backbone
features. The entire head is a derived circuit: cofiber decomposition (fixed)
+ linear predictions (solved via matrix inverse).

Construction:
  1. Accumulate sufficient statistics: X^T X and X^T Y from training features
     at positive locations, where X = features and Y = targets.
  2. Solve: W = (X^T X + lambda I)^{-1} X^T Y  for classification, regression,
     and centerness independently.
  3. The resulting weights are the optimal linear predictor in the least-squares sense.

There is no training loop, no learning rate, no epochs. The head is computed
from a single pass over the training data and one matrix inverse per task.

Parameters: 69,976
Construction time: ~130 seconds on CPU
COCO val2017 mAP: 1.6

This is the first known fully-derived detection head on frozen backbone features.
"""

import os
import json
import time

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

NUM_CLASSES = 80


def cofiber_decompose(f, n_scales):
    cofibers = []
    residual = f
    for _ in range(n_scales - 1):
        omega = F.avg_pool2d(residual, 2)
        sigma_omega = F.interpolate(omega, size=residual.shape[2:], mode="bilinear", align_corners=False)
        cofibers.append(residual - sigma_omega)
        residual = omega
    cofibers.append(residual)
    return cofibers


class AnalyticalDetector(nn.Module):
    """Fully analytical detection head. All weights from closed-form solution."""
    name = "analytical_detector"
    needs_intermediates = False

    def __init__(self, feat_dim=768, num_classes=NUM_CLASSES, n_scales=3):
        super().__init__()
        self.n_scales = n_scales
        self.scale_norms = nn.ModuleList([nn.LayerNorm(feat_dim) for _ in range(n_scales)])
        # Direct linear: no hidden layer, no nonlinearity
        self.cls_weight = nn.Parameter(torch.randn(num_classes, feat_dim) * 0.01)
        self.cls_bias = nn.Parameter(torch.zeros(num_classes))
        self.reg_out = nn.Linear(feat_dim, 4)
        self.ctr_weight = nn.Parameter(torch.randn(1, feat_dim) * 0.01)
        self.ctr_bias = nn.Parameter(torch.zeros(1))
        self.scale_params = nn.Parameter(torch.ones(n_scales))

    def forward(self, spatial, inter=None):
        cofibers = cofiber_decompose(spatial, self.n_scales)
        cls_l, reg_l, ctr_l = [], [], []
        for i, cof in enumerate(cofibers):
            B, C, H, W = cof.shape
            f = self.scale_norms[i](cof.permute(0, 2, 3, 1).reshape(-1, C))
            cls = (f @ self.cls_weight.T + self.cls_bias).reshape(B, H, W, -1).permute(0, 3, 1, 2)
            reg_raw = (self.reg_out(f) * self.scale_params[i]).clamp(-10, 10)
            reg = torch.exp(reg_raw).reshape(B, H, W, 4).permute(0, 3, 1, 2)
            ctr = (f @ self.ctr_weight.T + self.ctr_bias).reshape(B, H, W, 1).permute(0, 3, 1, 2)
            cls_l.append(cls)
            reg_l.append(reg)
            ctr_l.append(ctr)
        return cls_l, reg_l, ctr_l


def construct_analytical_head(cache_dir, n_images=20000, lam=1e-3, resolution=640):
    """Construct all weights from closed-form least-squares. Zero training."""
    from analytical_head import accumulate_statistics, solve_head
    stats = accumulate_statistics(cache_dir, n_images, lam, resolution)
    return solve_head(stats)