"""Original, executed mechanism tests for model refactoring.

No pretrained model, image-quality dataset, GPU, or mobile device is used.
Run: python vision_attention_lab.py --out ../evidence/vision_attention.json
Tested with Python 3.13.5 and torch 2.10.0+cpu.
"""
from __future__ import annotations
import argparse
import hashlib
import json
import platform
from pathlib import Path
import torch
from torch import Tensor, nn
from torch.nn import functional as F


class ConvPair(nn.Module):
    """Two dense convolutions with an independently pruneable hidden axis."""
    def __init__(self, hidden: int = 12) -> None:
        super().__init__()
        self.first = nn.Conv2d(3, hidden, 3, padding=1)
        self.second = nn.Conv2d(hidden, 3, 3, padding=1)

    def forward(self, x: Tensor) -> Tensor:
        return self.second(F.relu(self.first(x)))


def check_indices(keep: Tensor, size: int) -> None:
    if keep.ndim != 1 or keep.dtype != torch.long or keep.numel() == 0:
        raise ValueError('Expected nonempty 1D int64 indices')
    if keep.unique().numel() != keep.numel():
        raise ValueError('Duplicate indices')
    if keep.min().item() < 0 or keep.max().item() >= size:
        raise IndexError('Index outside supported range')


@torch.no_grad()
def prune_conv_pair(old: ConvPair, keep: Tensor) -> ConvPair:
    """Only for ConvPair, not an arbitrary CNN or grouped convolution."""
    check_indices(keep, old.first.out_channels)
    keep = keep.to(old.first.weight.device)
    new = ConvPair(len(keep)).to(old.first.weight)
    new.first.weight.copy_(old.first.weight[keep])
    new.first.bias.copy_(old.first.bias[keep])
    new.second.weight.copy_(old.second.weight[:, keep])
    new.second.bias.copy_(old.second.bias)
    return new.train(old.training)


class SplitSelfAttention(nn.Module):
    """Plain attention, no dropout, cache, rotary encoding, or shared KV heads.

    External width stays constant while the number of internal heads changes.
    This is deliberately NOT nn.MultiheadAttention or a pretrained DiT adapter.
    """
    def __init__(self, width: int = 32, heads: int = 4,
                 head_dim: int = 8) -> None:
        super().__init__()
        if min(width, heads, head_dim) < 1:
            raise ValueError('All dimensions must be positive')
        self.width, self.heads, self.head_dim = width, heads, head_dim
        inner = heads * head_dim
        self.q = nn.Linear(width, inner)
        self.k = nn.Linear(width, inner)
        self.v = nn.Linear(width, inner)
        self.out = nn.Linear(inner, width)

    def forward(self, x: Tensor, head_mask: Tensor | None = None) -> Tensor:
        batch, length, _ = x.shape
        def project(layer: nn.Linear) -> Tensor:
            return layer(x).reshape(
                batch, length, self.heads, self.head_dim).transpose(1, 2)
        q, k, v = project(self.q), project(self.k), project(self.v)
        p = ((q @ k.transpose(-1, -2)) / self.head_dim ** 0.5).softmax(-1)
        y = p @ v
        if head_mask is not None:
            if head_mask.shape != (self.heads,):
                raise ValueError('Head mask must contain one value per head')
            y = y * head_mask.to(y).reshape(1, self.heads, 1, 1)
        y = y.transpose(1, 2).contiguous().reshape(batch, length, -1)
        return self.out(y)


@torch.no_grad()
def prune_heads(old: SplitSelfAttention, keep: Tensor) -> SplitSelfAttention:
    check_indices(keep, old.heads)
    keep = keep.to(old.q.weight.device)
    offsets = torch.arange(old.head_dim, device=keep.device)
    columns = (keep[:, None] * old.head_dim + offsets).reshape(-1)
    new = SplitSelfAttention(old.width, len(keep), old.head_dim).to(old.q.weight)
    for name in ('q', 'k', 'v'):
        a, b = getattr(old, name), getattr(new, name)
        b.weight.copy_(a.weight[columns])
        b.bias.copy_(a.bias[columns])
    new.out.weight.copy_(old.out.weight[:, columns])
    new.out.bias.copy_(old.out.bias)
    return new.train(old.training)


@torch.inference_mode()
def tiled_single_conv(conv: nn.Conv2d, x: Tensor, tile: int,
                      tile_norm: nn.Module | None = None) -> Tensor:
    """Halo tiling for ONE stride-one, dilation-one convolution.

    Expects padding=0, square odd kernel, and groups=1. Applies zero padding
    only at the complete image boundary. This test keeps full input/output
    resident and does NOT demonstrate bounded total process memory.
    tile_norm is used only to demonstrate a deliberately invalid extension.
    """
    if (conv.stride != (1, 1) or conv.dilation != (1, 1)
            or conv.padding != (0, 0) or conv.groups != 1):
        raise ValueError('Unsupported convolution contract')
    k1, k2 = conv.kernel_size
    if k1 != k2 or k1 % 2 != 1 or tile < 1:
        raise ValueError('Requires an odd square kernel and positive tile')
    if x.ndim != 4 or x.shape[1] != conv.in_channels:
        raise ValueError('Expected [batch, in_channels, height, width]')
    radius = k1 // 2
    pad = F.pad(x, (radius,) * 4)
    height, width = x.shape[-2:]
    output = x.new_empty(x.shape[0], conv.out_channels, height, width)
    for y0 in range(0, height, tile):
        y1 = min(height, y0 + tile)
        for x0 in range(0, width, tile):
            x1 = min(width, x0 + tile)
            patch = pad[..., y0:y1 + 2*radius, x0:x1 + 2*radius]
            result = conv(patch)
            if tile_norm is not None:
                result = tile_norm(result)
            output[..., y0:y1, x0:x1] = result
    return output


def parameters(model: nn.Module) -> int:
    return sum(p.numel() for p in model.parameters())


def run() -> dict:
    torch.set_num_threads(1)
    torch.manual_seed(20260924)
    tests = []
    with torch.inference_mode():
        pair = ConvPair().eval()
        image = torch.randn(2, 3, 37, 41)
        keep = torch.tensor([0, 2, 4, 6, 8, 10])
        small = prune_conv_pair(pair, keep)
        hidden = F.relu(pair.first(image))
        mask = torch.zeros(1, 12, 1, 1)
        mask[:, keep] = 1
        expected = pair.second(hidden * mask)
        actual = small(image)
        torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6)
        tests.append({'name': 'conv_channels_masked_equivalence', 'status': 'passed',
                      'max_abs_error': (actual-expected).abs().max().item(),
                      'original_parameters': parameters(pair),
                      'pruned_parameters': parameters(small)})

        attention = SplitSelfAttention().eval()
        seq = torch.randn(2, 19, 32)
        keep = torch.tensor([0, 2])
        smaller_attention = prune_heads(attention, keep)
        head_mask = torch.tensor([1., 0., 1., 0.])
        expected = attention(seq, head_mask)
        actual = smaller_attention(seq)
        torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6)
        tests.append({'name': 'attention_heads_masked_equivalence', 'status': 'passed',
                      'max_abs_error': (actual-expected).abs().max().item(),
                      'original_parameters': parameters(attention),
                      'pruned_parameters': parameters(smaller_attention),
                      'external_width': 32, 'original_inner_width': 32,
                      'pruned_inner_width': 16})

        conv = nn.Conv2d(3, 5, 5, padding=0).eval()
        x = torch.randn(1, 3, 37, 41)
        full = conv(F.pad(x, (2, 2, 2, 2)))
        for tile in (7, 16, 64):
            tiled = tiled_single_conv(conv, x, tile)
            torch.testing.assert_close(full, tiled, rtol=1e-5, atol=1e-6)
            tests.append({'name': f'single_conv_halo_tile_{tile}', 'status': 'passed',
                          'max_abs_error': (full-tiled).abs().max().item()})

        norm = nn.GroupNorm(1, 5).eval()
        full_norm = norm(full)
        wrong = tiled_single_conv(conv, x, 7, tile_norm=norm)
        error = (full_norm-wrong).abs().max().item()
        assert error > 1e-3, 'Expected a detectable normalization mismatch'
        tests.append({'name': 'per_tile_groupnorm_is_not_equivalent',
                      'status': 'expected_mismatch_confirmed',
                      'max_abs_error': error})

        ffn = nn.Sequential(nn.Linear(32, 96), nn.GELU(), nn.Linear(96, 32)).eval()
        whole = ffn(seq)
        pieces = torch.cat([ffn(z) for z in seq.split(5, dim=1)], dim=1)
        torch.testing.assert_close(whole, pieces, rtol=1e-5, atol=1e-6)
        tests.append({'name': 'tokenwise_ffn_chunking', 'status': 'passed',
                      'max_abs_error': (whole-pieces).abs().max().item()})

        # Correctly reject head-index errors instead of silently changing mapping.
        for invalid in (torch.tensor([0, 0]), torch.tensor([4])):
            try:
                prune_heads(attention, invalid)
            except (ValueError, IndexError):
                pass
            else:
                raise AssertionError('Invalid head indices accepted')
        tests.append({'name': 'invalid_head_indices_rejected', 'status': 'passed'})
    return {'experiment': 'Vision and attention mechanism tests 1.0',
            'environment': {'python': platform.python_version(), 'torch': torch.__version__,
                            'platform': platform.platform(), 'threads': 1},
            'source_sha256': hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
            'tests': tests,
            'limitations': ['Random inputs and randomly initialized modules',
                           'No learned task quality or compression benchmark',
                           'No diffusion checkpoint, video dataset, GPU, or Android execution',
                           'No peak RAM measurement; output buffers are retained for testing']}


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--out', type=Path)
    args = parser.parse_args()
    report = run()
    text = json.dumps(report, indent=2) + '\n'
    if args.out:
        if args.out.exists():
            raise FileExistsError('Use a new output path to preserve prior evidence')
        args.out.parent.mkdir(parents=True, exist_ok=True)
        args.out.write_text(text)
    print(text)