#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
# Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights Reserved.
#  MIT License  (https://opensource.org/licenses/MIT)

import types
import torch
import torch.nn.functional as F


def export_rebuild_model(model, **kwargs):
    model.device = kwargs.get("device")

    # store original forward since self.extract_features is calling it
    model._original_forward = model.forward

    model.forward = types.MethodType(export_forward, model)
    model.export_dummy_inputs = types.MethodType(export_dummy_inputs, model)
    model.export_input_names = types.MethodType(export_input_names, model)
    model.export_output_names = types.MethodType(export_output_names, model)
    model.export_dynamic_axes = types.MethodType(export_dynamic_axes, model)
    model.export_name = types.MethodType(export_name, model)

    return model


def export_forward(self, x: torch.Tensor):
    with torch.no_grad():
        if self.cfg.normalize:
            mean = torch.mean(x, dim=1, keepdim=True)
            var = torch.var(x, dim=1, keepdim=True, unbiased=False)
            x = (x - mean) / torch.sqrt(var + 1e-5)
            x = x.view(x.shape[0], -1)

        # Call the original forward directly just like extract_features
        # Cannot directly use self.extract_features since it is being replaced by export_forward
        res = self._original_forward(
            source=x, padding_mask=None, mask=False, features_only=True, remove_extra_tokens=True
        )

        x = res["x"]

        return x


def export_dummy_inputs(self):
    return (torch.randn(1, 16000),)


def export_input_names(self):
    return ["input"]


def export_output_names(self):
    return ["output"]


def export_dynamic_axes(self):
    return {
        "input": {
            0: "batch_size",
            1: "sequence_length",
        },
        "output": {0: "batch_size", 1: "sequence_length"},
    }


def export_name(self):
    return "emotion2vec"