import torch
from .hourglass import FANAU


class Model:
    def __init__(self, npts=12, corenet='pretrained_models/disfa_adaptation_f0.pth', use_cuda=True):
        self.FAN = FANAU(num_modules=1, n_points=npts)
        self.FAN.load_state_dict(torch.load(corenet, map_location='cpu')['state_dict'])
        self.FAN.eval()
        if use_cuda:
            self.FAN.cuda()

    def __call__(self, x):
        H = self.FAN(x)
        H = H if H.__class__.__name__ == 'Tensor' else H[-1]
        return H

    def _forward_FAN(self, images):
        with torch.no_grad():
            self.FAN.eval()
            H = self.FAN(images)
        return H

    def forward_FAN(self, images):
        H = self.FAN(images)
        return H


class AUdetector:
    def __init__(self, au_model_path='models/pretrained/au_detector/disfa_adaptation_f0.pth', use_cuda=True):
        self.naus = 12
        self.AUdetector = Model(npts=self.naus, corenet=au_model_path, use_cuda=use_cuda)
        self.use_cuda = use_cuda

    def detect_AU(self, img):
        img_normalized = (img - img.min()) / (img.max() - img.min())
        if self.use_cuda:
            img_normalized = img_normalized.cuda()

        if img_normalized.ndim == 3:
            img_normalized = img_normalized.unsqueeze(0)

        heatmaps = self.AUdetector.forward_FAN(img_normalized)
        intensities = torch.nn.MaxPool2d((64, 64))(heatmaps).squeeze(2).squeeze(2)

        return intensities