# examples/efficientnet.py
"""Educational EfficientNet B0/B4, compatible with torchvision 0.27.0.

Run --smoke without downloads, or --predict --image PHOTO for released weights.
Architecture reference: torchvision at 78839c2b06c83c6cfb5c4da692ffb331bbd4c4cc.
"""

import argparse
import json
import math
from pathlib import Path

from PIL import Image
import torch
from torch import nn
from torchvision.models import EfficientNet_B0_Weights, EfficientNet_B4_Weights
from torchvision.ops import StochasticDepth
from torchvision.transforms import functional as F

PRESETS = {"b0": (1.0, 1.0, 224, 0.2), "b4": (1.4, 1.8, 380, 0.4)}
# (expansion, kernel, first stride, input channels, output channels, repeats)
BASE_STAGES = [
    (1, 3, 1, 32, 16, 1),
    (6, 3, 2, 16, 24, 2),
    (6, 5, 2, 24, 40, 2),
    (6, 3, 2, 40, 80, 3),
    (6, 5, 1, 80, 112, 3),
    (6, 5, 2, 112, 192, 4),
    (6, 3, 1, 192, 320, 1),
]


# region: scaling
def round_channels(channels, width_multiplier, divisor=8):
    """Round width while avoiding a reduction greater than ten percent."""
    scaled = channels * width_multiplier
    if scaled <= 0 or divisor <= 0:
        raise ValueError("Channels and divisor must be positive")
    rounded = max(divisor, int(scaled + divisor / 2) // divisor * divisor)
    if rounded < 0.9 * scaled:
        rounded += divisor
    return rounded


def round_repeats(repeats, depth_multiplier):
    """Round each stage's depth upward."""
    if repeats <= 0 or depth_multiplier <= 0:
        raise ValueError("Repeats and depth multiplier must be positive")
    return math.ceil(repeats * depth_multiplier)


# endregion: scaling


def conv_norm(cin, cout, kernel=1, stride=1, groups=1, activation=True):
    """Convolution, batch normalization, and optional SiLU."""
    layers = [
        nn.Conv2d(
            cin,
            cout,
            kernel,
            stride,
            (kernel - 1) // 2,
            groups=groups,
            bias=False,
        ),
        nn.BatchNorm2d(cout),
    ]
    if activation:
        layers.append(nn.SiLU(inplace=True))
    return nn.Sequential(*layers)


# region: se
class SqueezeExcitation(nn.Module):
    """Gate expanded channels using an input-width squeeze bottleneck."""

    def __init__(self, expanded, input_channels):
        super().__init__()
        squeeze = max(1, input_channels // 4)
        self.avgpool = nn.AdaptiveAvgPool2d(1)
        self.fc1 = nn.Conv2d(expanded, squeeze, 1)
        self.fc2 = nn.Conv2d(squeeze, expanded, 1)
        self.activation = nn.SiLU(inplace=True)
        self.scale_activation = nn.Sigmoid()

    def forward(self, x):
        pooled = self.avgpool(x)
        gates = self.scale_activation(
            self.fc2(self.activation(self.fc1(pooled)))
        )
        return x * gates


# endregion: se


class MBConv(nn.Module):
    """An inverted bottleneck with the pinned port's state-key layout."""

    def __init__(self, cin, cout, expansion, kernel, stride, drop_probability):
        super().__init__()
        # region: spatial
        expanded = round_channels(cin, expansion)
        layers = []
        if expanded != cin:
            layers.append(conv_norm(cin, expanded, kernel=1))
        layers.append(
            conv_norm(
                expanded,
                expanded,
                kernel=kernel,
                stride=stride,
                groups=expanded,
            )
        )
        layers.append(SqueezeExcitation(expanded, cin))
        # endregion: spatial
        # region: residual
        layers.append(conv_norm(expanded, cout, activation=False))
        self.block = nn.Sequential(*layers)
        self.use_res_connect = stride == 1 and cin == cout
        self.stochastic_depth = StochasticDepth(drop_probability, "row")

    def forward(self, x):
        branch = self.block(x)
        if self.use_res_connect:
            branch = self.stochastic_depth(branch) + x
        return branch

    # endregion: residual


class EfficientNet(nn.Module):
    """Seven scaled MBConv stages, stem, and ImageNet classification head."""

    def __init__(self, model_id):
        super().__init__()
        if model_id not in PRESETS:
            raise ValueError("Supported models are b0 and b4")
        # region: stages
        width, depth, _, dropout = PRESETS[model_id]
        stem_channels = round_channels(32, width)
        layers = [conv_norm(3, stem_channels, kernel=3, stride=2)]
        total = sum(round_repeats(s[5], depth) for s in BASE_STAGES)
        block_index = 0
        for expansion, kernel, stride, cin, cout, repeats in BASE_STAGES:
            cin, cout = round_channels(cin, width), round_channels(cout, width)
            blocks = []
            for repeat in range(round_repeats(repeats, depth)):
                blocks.append(
                    MBConv(
                        cin if repeat == 0 else cout,
                        cout,
                        expansion,
                        kernel,
                        stride if repeat == 0 else 1,
                        0.2 * block_index / total,
                    )
                )
                block_index += 1
            layers.append(nn.Sequential(*blocks))
        head_channels = 4 * round_channels(320, width)
        layers.append(conv_norm(round_channels(320, width), head_channels))
        self.features = nn.Sequential(*layers)
        self.avgpool = nn.AdaptiveAvgPool2d(1)
        self.classifier = nn.Sequential(
            nn.Dropout(dropout, inplace=True), nn.Linear(head_channels, 1000)
        )
        # endregion: stages

    # region: head
    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        return self.classifier(x)  # Logits; softmax belongs to the display.

    # endregion: head


def build_efficientnet(model_id):
    """Build an untrained, CPU model without network access."""
    return EfficientNet(model_id)


def get_weights(model_id):
    """Select an explicit pretrained enum instead of a mutable DEFAULT."""
    if model_id == "b0":
        return EfficientNet_B0_Weights.IMAGENET1K_V1
    if model_id == "b4":
        return EfficientNet_B4_Weights.IMAGENET1K_V1
    raise ValueError("Supported models are b0 and b4")


# region: inference
def load_pretrained(model_id):
    """Strictly load the released port; weights are cached outside this repo."""
    weights = get_weights(model_id)
    model = build_efficientnet(model_id)
    model.load_state_dict(
        weights.get_state_dict(progress=False, check_hash=True), strict=True
    )
    return model.cpu().float().eval(), weights


def prepare_image(image, weights_enum):
    """Keep the exact PIL crop and normalized tensor together."""
    transform = weights_enum.transforms()
    image = image.convert("RGB")
    resized = F.resize(
        image,
        transform.resize_size,
        interpolation=transform.interpolation,
        antialias=transform.antialias,
    )
    crop = F.center_crop(resized, transform.crop_size)
    tensor = F.convert_image_dtype(F.pil_to_tensor(crop), torch.float32)
    batch = F.normalize(tensor, transform.mean, transform.std).unsqueeze(0)
    size = transform.crop_size[0]
    left = int(round((resized.width - size) / 2))
    top = int(round((resized.height - size) / 2))
    record = {
        "resizeSize": transform.resize_size[0],
        "cropSize": size,
        "interpolation": "bicubic",
        "antialias": transform.antialias,
        "mean": transform.mean,
        "std": transform.std,
        "resizedShape": [resized.height, resized.width],
        "cropBox": [left, top, left + size, top + size],
    }
    return crop, batch, record


# endregion: inference


def main():
    """Run a mechanics check or an explicit local-photo prediction."""
    parser = argparse.ArgumentParser(description=__doc__)
    mode = parser.add_mutually_exclusive_group(required=True)
    mode.add_argument("--smoke", action="store_true")
    mode.add_argument("--predict", action="store_true")
    parser.add_argument("--model", choices=PRESETS, default="b0")
    parser.add_argument("--image", type=Path)
    args = parser.parse_args()
    torch.set_num_threads(2)
    torch.manual_seed(0)
    if args.smoke:
        model = build_efficientnet(args.model).train()
        loss = nn.functional.cross_entropy(
            model(torch.randn(2, 3, 64, 64)), torch.tensor([0, 1])
        )
        loss.backward()
        if not torch.isfinite(loss) or not all(
            p.grad is not None and torch.isfinite(p.grad).all()
            for p in model.parameters()
        ):
            raise RuntimeError("Nonfinite loss or missing gradients")
        print(
            json.dumps(
                {
                    "mode": "random-weight mechanics check",
                    "loss": loss.item(),
                    "finiteGradients": True,
                }
            )
        )
        return
    if args.image is None:
        parser.error("--predict requires --image")
    try:
        with Image.open(args.image) as source:
            image = source.convert("RGB")
    except (OSError, ValueError) as error:
        parser.error(f"Cannot read image: {error}")
    model, weights = load_pretrained(args.model)
    _, batch, _ = prepare_image(image, weights)
    with torch.inference_mode():
        probabilities = model(batch).softmax(dim=1)[0].tolist()
    indices = sorted(range(1000), key=lambda i: (-probabilities[i], i))[:5]
    print(
        json.dumps(
            [
                {
                    "class": weights.meta["categories"][i],
                    "probability": probabilities[i],
                }
                for i in indices
            ],
            indent=2,
        )
    )


if __name__ == "__main__":
    main()