"""Check MapTR query/output shapes on CPU, not full-model inference.

Source: hustvl/MapTR@a6872d8d9670bde17b4b01560f1221f88b443d55
Decoder outputs and prediction heads below are deliberately simplified.
"""

import torch
from torch import nn


@torch.no_grad()
def check_shapes(name, instances, points, dim, layers, batch=2, classes=3):
    torch.manual_seed(0)
    queries = instances * points

    # Distinct values let us check pairing order, not only tensor sizes.
    instance = torch.arange(instances, dtype=torch.float32)[:, None].repeat(1, 2 * dim) * 10
    point = torch.arange(points, dtype=torch.float32)[:, None].repeat(1, 2 * dim)
    combined = instance[:, None, :] + point[None, :, :]
    query_embed = combined.flatten(0, 1)
    assert combined.shape == (instances, points, 2 * dim)
    assert query_embed.shape == (queries, 2 * dim)
    for i in range(instances):
        for j in range(points):
            torch.testing.assert_close(query_embed[i * points + j], instance[i] + point[j])

    query_pos, query = query_embed.split(dim, dim=1)
    query_pos = query_pos.unsqueeze(0).expand(batch, -1, -1)
    query = query.unsqueeze(0).expand(batch, -1, -1)
    # Zero initialization keeps reference points at 0.5 for this demonstration.
    reference_head = nn.Linear(dim, 2)
    nn.init.zeros_(reference_head.weight)
    nn.init.zeros_(reference_head.bias)
    reference = reference_head(query_pos).sigmoid()
    assert query_pos.shape == query.shape == (batch, queries, dim)
    assert reference.shape == (batch, queries, 2)

    query_seq = query.permute(1, 0, 2)
    query_pos_seq = query_pos.permute(1, 0, 2)
    assert query_seq.shape == query_pos_seq.shape == (queries, batch, dim)

    # Placeholder only: NO self-attention, BEV sampling or reference refinement.
    hs_seq = query_seq.unsqueeze(0).repeat(layers, 1, 1, 1)
    assert hs_seq.shape == (layers, queries, batch, dim)
    hs = hs_seq.permute(0, 2, 1, 3)
    grouped = hs.reshape(layers, batch, instances, points, dim)
    cls_scores = nn.Linear(dim, classes)(grouped.mean(dim=3))
    offsets = nn.Linear(dim, 2)(hs)
    # Reuse the initial reference for every placeholder layer; not real refinement.
    coords = (offsets + torch.logit(reference, eps=1e-5).unsqueeze(0)).sigmoid()
    coords = coords.reshape(layers, batch, instances, points, 2)
    assert cls_scores.shape == (layers, batch, instances, classes)
    assert coords.shape == (layers, batch, instances, points, 2)
    assert torch.isfinite(coords).all()
    assert ((coords >= 0) & (coords <= 1)).all()

    print(f'[{name}] I={instances}, P={points}, D={dim}, B={batch}, L={layers}, C={classes}')
    if queries <= 6:
        print('flat order (instance, point):', [divmod(q, points) for q in range(queries)])
        print('first channel:', query_embed[:, 0].tolist())
    for label, tensor in (
        ('combined', combined), ('query_embed', query_embed),
        ('query_pos / query (batch first)', query), ('reference', reference),
        ('decoder query (sequence first)', query_seq), ('hs (placeholder)', hs_seq),
        ('all_cls_scores (toy head)', cls_scores), ('all_pts_preds (toy head)', coords),
    ):
        print(f'{label}: {tuple(tensor.shape)}')
    print('all assertions passed\n')


if __name__ == '__main__':
    check_shapes('small example', instances=2, points=3, dim=4, layers=2)
    check_shapes('tiny config dimensions', instances=50, points=20, dim=256, layers=6)