"""Inference-only compatibility adapter, not a new action architecture.

Uses official SlowOnly backbone and the unmodified official BBoxHeadAVA class.
Torchvision RoIAlign replaces MMCV's CUDA extension with identical requested
alignment/scale/pooling arguments. Training and MMDetection wrappers are excluded.
"""
import ast
import pathlib

import numpy as np
import torch
from torch import nn
from torchvision.ops import roi_align
from mmaction.models.backbones import ResNet3dSlowOnly


def official_head():
    path = pathlib.Path('/opt/mmaction2/mmaction/models/roi_heads/bbox_heads/bbox_head.py')
    tree = ast.parse(path.read_text())
    cls = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == 'BBoxHeadAVA')
    # Training-only MMDetection imports are not needed by the exact forward method.
    body = ast.parse('from __future__ import annotations').body + [cls]
    namespace = {'torch': torch, 'nn': nn, 'np': np}
    exec(compile(ast.fix_missing_locations(ast.Module(body=body,type_ignores=[])),str(path),'exec'),namespace)
    return namespace['BBoxHeadAVA']


class AvaAdapter(nn.Module):
    def __init__(self, checkpoint):
        super().__init__()
        self.backbone = ResNet3dSlowOnly(depth=50,pretrained=None,pretrained2d=False,
            lateral=False,num_stages=4,conv1_kernel=(1,7,7),conv1_stride_t=1,
            pool1_stride_t=1,spatial_strides=(1,2,2,1))
        self.head = official_head()(background_class=True,in_channels=2048,
                                   num_classes=81,multilabel=True,dropout_ratio=.5)
        state = checkpoint['state_dict']
        backbone = {k[len('backbone.'):]:v for k,v in state.items() if k.startswith('backbone.')}
        head = {k[len('roi_head.bbox_head.'):]:v for k,v in state.items() if k.startswith('roi_head.bbox_head.')}
        assert len(state) == len(backbone)+len(head), 'Unaccounted checkpoint keys'
        self.backbone.load_state_dict(backbone,strict=True)
        self.head.load_state_dict(head,strict=True)
        self.checkpoint_key_count = len(state)

    def forward(self, video, boxes):
        feature = self.backbone(video)
        if isinstance(feature, tuple):
            assert len(feature) == 1
            feature = feature[0]
        feature = feature.mean(dim=2)
        roi = roi_align(feature,boxes,output_size=(8,8),spatial_scale=1/16,
                        sampling_ratio=0,aligned=True)
        return self.head(roi.unsqueeze(2))


def check_roi_contract(device):
    grid = (torch.arange(4,device=device)[None,:] + 10*torch.arange(4,device=device)[:,None]).float()[None,None]
    boxes = torch.tensor([[0,1,1,3,3]],dtype=torch.float32,device=device)
    result = roi_align(grid,boxes,(2,2),spatial_scale=1,sampling_ratio=1,aligned=True)
    expected = torch.tensor([[[[11,12],[21,22]]]],device=device,dtype=torch.float32)
    torch.testing.assert_close(result,expected,rtol=0,atol=1e-6)