import torch
import torch.nn.functional as F
import numpy as np
from torch import nn

from tagging.modules.frontend import FrontEnd
from tagging.modules.mlp import Projection, MLP


class TaggingModel(nn.Module):
    def __init__(
        self,
        frontend_name="musicfm_concat",
        model_path=None,
        is_flash=True,
    ):
        super().__init__()

        # prepare encoders
        self.frontend = FrontEnd(frontend_name, layer_ix=12, is_flash=is_flash)

        # get projection layers
        self.projection = MLP([1024, 512, 407])

        # load model
        if model_path:
            S = torch.load(model_path)["state_dict"]
            SS = {k[6:]: v for k, v in S.items()}
            self.load_state_dict(SS, strict=True)

    def forward(self, wav):
        out = self.frontend(wav)
        out = self.projection(out)

        return out