import copy
import logging
from typing import Optional, Tuple, Union

import humanfriendly
import torch
import os

from s3prl.upstream.hubert.convert import load_and_convert_fairseq_ckpt

os.environ["CUDA_VISIBLE_DEVICES"] = "6"



class S3prlFrontend(torch.nn.Module):
    """Speech Pretrained Representation frontend structure for ASR."""

    def __init__(
        self,
        fs: Union[int, str] = 16000,
        frontend_conf: Optional[dict] = None,
        download_dir: str = None,
        multilayer_feature: bool = False,
        layer: int = -1,
    ):
        try:
            import s3prl
            from s3prl.nn import Featurizer, S3PRLUpstream
            from s3prl.upstream.hubert.convert import load_and_convert_fairseq_ckpt
        except Exception as e:
            print("Error: S3PRL is not properly installed.")
            print("Please install S3PRL: cd ${MAIN_ROOT}/tools && make s3prl.done")
            raise e

        super().__init__()

        if isinstance(fs, str):
            fs = humanfriendly.parse_size(fs)
        if fs != 16000:
            logging.warning(
                "All the upstream models in S3PRL now only support 16 kHz audio."
            )

        if download_dir is not None:
            s3prl.util.download.set_dir(download_dir)

        assert frontend_conf.get("upstream", None) in S3PRLUpstream.available_names()
        upstream = S3PRLUpstream(
            frontend_conf.get("upstream"),
            path_or_url=frontend_conf.get("upstream_ckpt", None),
            normalize=frontend_conf.get("normalize", False),
            extra_conf=frontend_conf.get("extra_conf", None),
        )
        # upstream = hubert_local(
        #     ckpt=frontend_conf.get("upstream_ckpt"),
        #     normalize=frontend_conf.get("normalize", False),
        # )

        if getattr(upstream.upstream, "model", None):
            if getattr(upstream.upstream.model, "feature_grad_mult", None) is not None:
                upstream.upstream.model.feature_grad_mult = 1.0
        upstream.eval()

        if layer != -1:
            layer_selections = [layer]
            assert (
                not multilayer_feature
            ), "multilayer feature will be deactivated, when specific layer used"
        else:
            layer_selections = None
        featurizer = Featurizer(upstream, layer_selections=layer_selections)

        self.multilayer_feature = multilayer_feature
        self.layer = layer
        self.upstream, self.featurizer = upstream, featurizer
        self.pretrained_params = copy.deepcopy(self.upstream.state_dict())
        self.frontend_type = "s3prl"
        self.hop_length = self.featurizer.downsample_rate
        self.tile_factor = frontend_conf.get("tile_factor", 1)

    def _tile_representations(self, feature):
        """Tile up the representations by `tile_factor`.

        Input - sequence of representations
                shape: (batch_size, seq_len, feature_dim)

        Output - sequence of tiled representations
                 shape: (batch_size, seq_len * factor, feature_dim)
        """
        assert (
            len(feature.shape) == 3
        ), "Input argument `feature` has invalid shape: {}".format(feature.shape)
        tiled_feature = feature.repeat(1, 1, self.tile_factor)
        tiled_feature = tiled_feature.reshape(
            feature.size(0), feature.size(1) * self.tile_factor, feature.size(2)
        )
        return tiled_feature

    def output_size(self) -> int:
        return self.featurizer.output_size

    def forward(
        self, input: torch.Tensor, input_lengths: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        feats, feats_lens = self.upstream(input, input_lengths)
        if self.layer != -1:
            layer = self.layer
            feats, feats_lens = feats[layer], feats_lens[layer]
            return feats, feats_lens

        if self.multilayer_feature:
            feats, feats_lens = self.featurizer(feats, feats_lens)
        else:
            feats, feats_lens = self.featurizer(feats[-1:], feats_lens[-1:])

        if self.tile_factor != 1:
            feats = self._tile_representations(feats)

        return feats, feats_lens

    def reload_pretrained_parameters(self):
        self.upstream.load_state_dict(self.pretrained_params)
        logging.info("Pretrained S3PRL frontend model parameters reloaded!")


if __name__ == "__main__":
    source_fairseq_path = "/home/work_nfs15/asr_data/ckpt/origin_chinese_hubert/chinese_hubert_large.pt"
    s3qrl_path = "/home/work_nfs15/asr_data/ckpt/origin_chinese_hubert/chinese_hubert_large_s3qrl.pt"
    # load_and_convert_fairseq_ckpt(source_fairseq_path, s3qrl_path) # 成功转换
    hubert_model = S3prlFrontend(
        frontend_conf={
            "upstream": "hubert_local",
            "upstream_model_config":"",
            "upstream_ckpt": s3qrl_path,
        },
        download_dir="./hub",
    multilayer_feature=True,)
    print(hubert_model)
    input_wav_tensor = torch.randn(5, 16000*10)
    # 随机出5个长度
    input_wav_lens = torch.randint(16000*5, 16000*10, (5,))
    feats, feats_lens = hubert_model(input_wav_tensor, input_wav_lens)
    print(feats.shape, feats_lens.shape)
