# examples/resnet.py
"""Offline reduced original post-activation ResNet mechanics.

Paper: https://arxiv.org/html/1512.03385v1, §§3.2–3.4 and Figure 5.
Bottleneck stride convention: official ResNet-50-deploy.prototxt at
https://github.com/KaimingHe/deep-residual-networks/blob/a7026cb6d478e131b765b898c312e25f9f6dc031/prototxt/ResNet-50-deploy.prototxt
The first 1x1 convolution downsamples, unlike torchvision ResNet-v1.5.
Random reduced widths and synthetic inputs do not reproduce ImageNet results.
PyTorch BN running statistics/SGD are not an exact Caffe training conversion.
Import does not instantiate models, seed RNGs, or execute experiments.
"""

import argparse
import hashlib
import json
from pathlib import Path

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


UPSTREAM_REVISION = "a7026cb6d478e131b765b898c312e25f9f6dc031"


def validate_config(in_channels, out_channels, stride):
    """Require positive widths and the original stride-one/two transitions."""
    if any(type(c) is not int or c < 1 for c in (in_channels, out_channels)):
        raise ValueError("channels must be positive integers")
    if type(stride) is not int or stride not in (1, 2):
        raise ValueError("stride must be 1 or 2")


def validate_input(x, channels):
    """Require finite, nonempty NCHW floating-point activations."""
    if x.ndim != 4 or min(x.shape) < 1 or x.shape[1] != channels:
        raise ValueError("input must be nonempty NCHW with matching channels")
    if not x.is_floating_point() or not torch.isfinite(x).all():
        raise ValueError("input must contain finite floating-point values")


# region: basic_branch
def basic_block(in_channels, channels, stride=1):
    """Build the original two-convolution residual branch and shortcut."""
    validate_config(in_channels, channels, stride)
    branch = nn.Sequential(
        nn.Conv2d(
            in_channels, channels, 3, stride=stride, padding=1, bias=False
        ),
        nn.BatchNorm2d(channels),
        nn.ReLU(),
        nn.Conv2d(channels, channels, 3, padding=1, bias=False),
        nn.BatchNorm2d(channels),
        # No ReLU here: the second activation follows the addition.
    )
    return nn.ModuleDict(
        {
            "branch": branch,
            "shortcut": make_shortcut(in_channels, channels, stride),
        }
    )


# endregion: basic_branch


# region: identity_shortcut
def identity_shortcut(x, residual):
    """Return x only if all dimensions match; forbid implicit broadcasting."""
    if x.shape != residual.shape:
        raise ValueError("identity shortcut requires equal NCHW shapes")
    return x


# endregion: identity_shortcut


# region: projection_shortcut
def make_shortcut(in_channels, out_channels, stride):
    """Select option B: project only at channel or spatial transitions."""
    validate_config(in_channels, out_channels, stride)
    if in_channels == out_channels and stride == 1:
        return nn.Identity()
    return nn.Sequential(
        nn.Conv2d(in_channels, out_channels, 1, stride=stride, bias=False),
        nn.BatchNorm2d(out_channels),
    )


# endregion: projection_shortcut


# region: post_add_relu
def post_add_relu(residual, shortcut):
    """Add exactly matching tensors, then apply the original final ReLU."""
    if residual.shape != shortcut.shape:
        raise ValueError("residual addition forbids broadcasting")
    return F.relu(residual + shortcut)


def block_forward(block, x):
    """Evaluate both routes before the post-add nonlinearity."""
    validate_input(x, block["branch"][0].in_channels)
    residual = block["branch"](x)
    shortcut = block["shortcut"](x)
    if isinstance(block["shortcut"], nn.Identity):
        shortcut = identity_shortcut(x, residual)
    return post_add_relu(residual, shortcut)


# endregion: post_add_relu


# region: bottleneck
def bottleneck_block(in_channels, width, stride=1):
    """Use 4x expansion and the official Caffe first-1x1 stride placement."""
    validate_config(in_channels, width, stride)
    out_channels = 4 * width
    branch = nn.Sequential(
        nn.Conv2d(in_channels, width, 1, stride=stride, bias=False),
        nn.BatchNorm2d(width),
        nn.ReLU(),
        nn.Conv2d(width, width, 3, stride=1, padding=1, bias=False),
        nn.BatchNorm2d(width),
        nn.ReLU(),
        nn.Conv2d(width, out_channels, 1, bias=False),
        nn.BatchNorm2d(out_channels),
    )
    return nn.ModuleDict(
        {
            "branch": branch,
            "shortcut": make_shortcut(in_channels, out_channels, stride),
        }
    )


# endregion: bottleneck


# region: readout
def small_stack(kind="basic", width=4, classes=5):
    """Build a two-block mechanics stack, not a full ResNet-18/50."""
    validate_config(3, width, 1)
    if type(classes) is not int or classes < 2:
        raise ValueError("classes must be an integer >= 2")
    if kind not in ("basic", "bottleneck"):
        raise ValueError("kind must be basic or bottleneck")
    factory = basic_block if kind == "basic" else bottleneck_block
    expansion = 1 if kind == "basic" else 4
    stack = nn.ModuleDict(
        {
            "stem": nn.Sequential(
                nn.Conv2d(3, width, 3, padding=1, bias=False),
                nn.BatchNorm2d(width),
                nn.ReLU(),
            ),
            "blocks": nn.ModuleList(
                [
                    factory(width, width),
                    factory(expansion * width, 2 * width, stride=2),
                ]
            ),
            "classifier": nn.Linear(expansion * 2 * width, classes),
        }
    )
    for module in stack.modules():
        if isinstance(module, nn.Conv2d):
            nn.init.kaiming_normal_(
                module.weight, mode="fan_out", nonlinearity="relu"
            )
    return stack


def stack_forward(stack, images):
    """Global-average-pool the spatial axes and return raw class logits."""
    validate_input(images, 3)
    x = stack["stem"](images)
    for block in stack["blocks"]:
        x = block_forward(block, x)
    pooled = x.mean(dim=(2, 3))  # [B, C, H, W] -> [B, C]
    return stack["classifier"](pooled)


# endregion: readout


def jacobian_observation():
    """Measure a nontrivial full-block Jacobian away from all ReLU kinks."""
    block = basic_block(1, 1).double().eval()
    with torch.no_grad():
        for layer in block["branch"]:
            if isinstance(layer, nn.Conv2d):
                layer.weight.fill_(0.05)
            elif isinstance(layer, nn.BatchNorm2d):
                layer.running_var.fill_(1.0)
    x = torch.tensor([1.0, 1.5, 2.0, 2.5], dtype=torch.double).reshape(
        1, 1, 2, 2
    )
    fn = lambda value: block_forward(block, value)
    analytic = torch.autograd.functional.jacobian(fn, x).reshape(4, 4)
    columns = []
    epsilon = 1e-6
    for index in range(x.numel()):
        perturbation = torch.zeros_like(x)
        perturbation.flatten()[index] = epsilon
        columns.append(
            (
                (fn(x + perturbation) - fn(x - perturbation)) / (2 * epsilon)
            ).flatten()
        )
    numeric = torch.stack(columns, dim=1)
    return float((analytic - numeric).abs().max().detach())


def smoke(seed=7):
    """Run seeded reduced CPU forward/backward/update and factual checks."""
    torch.manual_seed(seed)
    images = torch.randn(2, 3, 9, 11)
    labels = torch.tensor([1, 4])
    observations = {}
    for kind in ("basic", "bottleneck"):
        stack = small_stack(kind)
        optimizer = torch.optim.SGD(stack.parameters(), lr=0.01, momentum=0.9)
        before = stack["stem"][0].weight.detach().clone()
        logits = stack_forward(stack, images)
        loss = F.cross_entropy(logits, labels)
        loss.backward()
        gradients = [p.grad for p in stack.parameters() if p.grad is not None]
        if not torch.isfinite(loss) or not all(
            torch.isfinite(g).all() for g in gradients
        ):
            raise RuntimeError("nonfinite loss or gradients")
        gradient_norm = torch.stack([g.norm() for g in gradients]).norm()
        optimizer.step()
        stack.eval()
        with torch.no_grad():
            first = stack_forward(stack, images)
            second = stack_forward(stack, images)
            block = stack["blocks"][1]
            channels = block["branch"][0].in_channels
            transition = block_forward(block, torch.randn(2, channels, 9, 11))
        observations[kind] = {
            "logits_shape": list(logits.shape),
            "transition_shape": list(transition.shape),
            "loss": float(loss.detach()),
            "gradient_norm": float(gradient_norm),
            "stem_update_norm": float(
                (stack["stem"][0].weight - before).norm().detach()
            ),
            "bn_batches_tracked": int(stack["stem"][1].num_batches_tracked),
            "eval_repeat_max_error": float((first - second).abs().max()),
        }
    zero = basic_block(1, 1).eval()
    with torch.no_grad():
        zero["branch"][3].weight.zero_()
    signed = torch.tensor([-2.0, -0.5, 0.5, 2.0]).reshape(1, 1, 2, 2)
    with torch.no_grad():
        zero_result = block_forward(zero, signed)
    signed_grad = torch.autograd.functional.jacobian(
        lambda x: block_forward(zero, x), signed
    )
    return {
        "source_sha256": hashlib.sha256(
            Path(__file__).read_bytes()
        ).hexdigest(),
        "upstream_revision": UPSTREAM_REVISION,
        "seed": seed,
        "device": "cpu",
        "reductions": {
            "width": 4,
            "blocks": 2,
            "classes": 5,
            "synthetic_shape": list(images.shape),
        },
        "variants": observations,
        "zero_residual_input": signed.flatten().tolist(),
        "zero_residual_output": zero_result.flatten().tolist(),
        "zero_residual_jacobian_diagonal": signed_grad.reshape(4, 4)
        .diag()
        .tolist(),
        "full_block_finite_difference_max_error": jacobian_observation(),
        "boundary": "random reduced mechanics; no pretrained semantics or benchmark reproduction",
    }


def main():
    """Print strict finite JSON only when the CLI smoke is explicitly invoked."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--smoke", action="store_true")
    parser.add_argument("--seed", type=int, default=7)
    args = parser.parse_args()
    if not args.smoke:
        parser.error("pass --smoke to run the offline reduced reference")
    torch.set_num_threads(1)
    print(json.dumps(smoke(args.seed), indent=2, allow_nan=False))


if __name__ == "__main__":
    main()