from openrec.legacy.recommenders import VisualPMF
from openrec.legacy.modules.extractions import MultiLayerFC
from openrec.legacy.modules.fusions import Average


class UserVisualPMF(VisualPMF):
    def __init__(
        self,
        batch_size,
        max_user,
        max_item,
        dim_embed,
        dims_user,
        dims_item,
        user_f_source,
        item_f_source,
        test_batch_size=None,
        item_serving_size=None,
        dropout_rate=None,
        l2_reg=None,
        l2_reg_mlp=None,
        opt="SGD",
        sess_config=None,
    ):

        self._dims_user = dims_user
        self._user_f_source = user_f_source

        super(UserVisualPMF, self).__init__(
            batch_size=batch_size,
            max_user=max_user,
            max_item=max_item,
            dim_embed=dim_embed,
            dims=dims_item,
            item_f_source=item_f_source,
            test_batch_size=test_batch_size,
            item_serving_size=item_serving_size,
            dropout_rate=dropout_rate,
            l2_reg=l2_reg,
            l2_reg_mlp=l2_reg_mlp,
            opt=opt,
            sess_config=sess_config,
        )

    def _build_user_inputs(self, train=True):

        super(UserVisualPMF, self)._build_user_inputs(train)
        if train:
            self._add_input(
                name="user_feature",
                dtype="float32",
                shape=[self._batch_size, self._user_f_source.shape[1]],
            )
        else:
            self._add_input(
                name="user_feature",
                dtype="float32",
                shape=[None, self._user_f_source.shape[1]],
                train=False,
            )

    def _input_mappings(self, batch_data, train):

        default_input_map = super(UserVisualPMF, self)._input_mappings(
            batch_data=batch_data, train=train
        )
        default_input_map[
            self._get_input("user_feature", train=train)
        ] = self._user_f_source[batch_data["user_id_input"]]
        return default_input_map

    def _build_user_extractions(self, train=True):

        super(UserVisualPMF, self)._build_user_extractions(train)

        self._add_module(
            "user_f",
            MultiLayerFC(
                in_tensor=self._get_input("user_feature", train=train),
                train=train,
                dims=self._dims_user,
                l2_reg=self._l2_reg_mlp,
                dropout_mid=self._dropout_rate,
                scope="user_MLP",
                reuse=not train,
            ),
            train=train,
        )

    def _build_default_fusions(self, train=True):

        super(UserVisualPMF, self)._build_default_fusions(train)
        self._add_module(
            "user_vec",
            Average(
                scope="user_average",
                reuse=not train,
                module_list=[
                    self._get_module("user_vec", train=train),
                    self._get_module("user_f", train=train),
                ],
                weight=2.0,
            ),
            train=train,
        )
