import torch
from torchvision import models, transforms
from torchvision.io import read_image
from torch import nn

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

backbone_spatial = models.mobilenet_v3_large(weights='DEFAULT').features.to(device)
backbone_motion  = models.mobilenet_v3_small(weights='DEFAULT').features.to(device)

backbone_spatial.eval()
backbone_motion.eval()

pool = nn.AdaptiveAvgPool2d(1)

norm = transforms.Normalize([0.485, 0.456, 0.406],
                            [0.229, 0.224, 0.225])
resize = transforms.Resize((224, 224))

def extract_features(frame_paths):
    imgs = []
    for p in frame_paths:
        img = read_image(p).float() / 255.0
        imgs.append(norm(resize(img)))

    imgs = torch.stack(imgs).to(device)

    with torch.no_grad():
        spatial = imgs[::3]
        spatial_feat = pool(backbone_spatial(spatial)).flatten(1)

        diffs = imgs[1:] - imgs[:-1]
        diffs = torch.where(torch.abs(diffs) > 0.08, diffs, torch.zeros_like(diffs))
        motion_feat = pool(backbone_motion(diffs)).flatten(1)

        if spatial_feat.shape[0] < motion_feat.shape[0]:
            pad = motion_feat.shape[0] - spatial_feat.shape[0]
            spatial_feat = torch.cat([spatial_feat,
                                      spatial_feat[-1:].repeat(pad, 1)])
        else:
            spatial_feat = spatial_feat[:motion_feat.shape[0]]

        features = torch.cat([spatial_feat, motion_feat], dim=1)

    return features.cpu()