"""Train a CNN on generated vertical/horizontal bars. Requires torch; no downloads."""
from __future__ import annotations
import torch
from torch import nn
from torch.nn import functional as F


def bar_images(count: int, *, seed: int):
    """Pre: count>0. Post: images (N,1,16,16), labels (N,) from a reproducible generator."""
    if count <= 0:
        raise ValueError("count must be positive")
    generator = torch.Generator().manual_seed(seed)
    images = torch.randn(count, 1, 16, 16, generator=generator)*.05
    labels = torch.arange(count) % 2
    positions = torch.randint(3, 12, (count,), generator=generator)
    # Invariant: preceding images contain the bar matching their class. Bound: count-i.
    for i, position in enumerate(positions):
        if labels[i] == 0:
            images[i, 0, :, position:position+2] += 1
        else:
            images[i, 0, position:position+2, :] += 1
    return images, labels


class BarCNN(nn.Module):
    """Map (N,1,16,16) to two raw class logits per image."""
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(nn.Conv2d(1, 8, 3, padding=1), nn.ReLU(),
                                      nn.Conv2d(8, 16, 3, padding=1), nn.ReLU(),
                                      nn.AdaptiveAvgPool2d(1))
        self.head = nn.Linear(16, 2)

    def forward(self, images):
        return self.head(self.features(images).flatten(1))


def train(*, epochs=80):
    """Pre: epochs>0. Post: model trained on seed-7 data; evaluated on distinct seed-9 data."""
    if epochs <= 0:
        raise ValueError("epochs must be positive")
    torch.set_num_threads(1)
    torch.manual_seed(7)
    images, labels = bar_images(256, seed=7)
    test_images, test_labels = bar_images(64, seed=9)
    model = BarCNN()
    optimizer = torch.optim.Adam(model.parameters(), lr=.01)
    history = []
    # Invariant: history records all completed full-batch updates. Bound: epochs-step.
    for _ in range(epochs):
        model.train()
        optimizer.zero_grad(set_to_none=True)
        loss = F.cross_entropy(model(images), labels)
        loss.backward()
        optimizer.step()
        history.append(loss.item())
    model.eval()
    with torch.no_grad():
        accuracy = (model(test_images).argmax(1) == test_labels).float().mean().item()
    return model, history, accuracy


def demo():
    model, history, accuracy = train()
    print("Training images: 256; independent test images: 64")
    print("Initial/final cross entropy:", history[0], history[-1])
    print("Test accuracy:", accuracy)
    print("Parameters:", sum(p.numel() for p in model.parameters()))
    print("This demonstrates the training pipeline, not performance on natural photos.")


def self_test():
    images, labels = bar_images(12, seed=5)
    check(images.shape, torch.Size([12, 1, 16, 16]))
    check(labels.dtype, torch.int64)
    check(BarCNN()(images).shape, torch.Size([12, 2]))
    model, history, accuracy = train()
    check(history[-1] < history[0], True)
    check(accuracy >= .9, True)
    # Validate one weight gradient against a central finite difference.
    layer = nn.Conv2d(1, 1, 2, bias=False).double()
    data = torch.arange(9, dtype=torch.float64).reshape(1, 1, 3, 3)
    layer(data).square().sum().backward()
    analytic = layer.weight.grad[0, 0, 0, 0].item()
    with torch.no_grad():
        original = layer.weight[0, 0, 0, 0].item()
        layer.weight[0, 0, 0, 0] = original + 1e-6
        plus = layer(data).square().sum().item()
        layer.weight[0, 0, 0, 0] = original - 1e-6
        minus = layer(data).square().sum().item()
        layer.weight[0, 0, 0, 0] = original
    check(abs(analytic-(plus-minus)/2e-6) < 1e-4, True)
    return 6

def check(actual, expected):
    """Require an expected result; this check also runs under python -O."""
    if actual != expected:
        raise AssertionError(f"expected {expected!r}, got {actual!r}")


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--self-test", action="store_true")
    args = parser.parse_args()
    if args.self_test:
        print(f"{self_test()} checks passed")
    else:
        demo()