#!/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 time
import torch
import logging
from typing import Dict, Tuple
from contextlib import contextmanager
from distutils.version import LooseVersion

from funasr.register import tables
from funasr.models.ctc.ctc import CTC
from funasr.utils import postprocess_utils
from funasr.metrics.compute_acc import th_accuracy
from funasr.utils.datadir_writer import DatadirWriter
from funasr.models.sanm_kws.model import SanmKWS
from funasr.models.paraformer.search import Hypothesis
from funasr.models.paraformer.cif_predictor import mae_loss
from funasr.train_utils.device_funcs import force_gatherable
from funasr.losses.label_smoothing_loss import LabelSmoothingLoss
from funasr.models.transformer.utils.add_sos_eos import add_sos_eos
from funasr.models.transformer.utils.nets_utils import make_pad_mask, pad_list
from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank


if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):
    from torch.cuda.amp import autocast
else:
    # Nothing to do if torch<1.6.0
    @contextmanager
    def autocast(enabled=True):
        yield


@tables.register("model_classes", "SanmKWSStreaming")
class SanmKWSStreaming(SanmKWS):
    """
    Author: Speech Lab of DAMO Academy, Alibaba Group
    Paraformer: Fast and Accurate Parallel Transformer for Non-autoregressive End-to-End Speech Recognition
    https://arxiv.org/abs/2206.08317
    """

    def __init__(
        self,
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)

    def forward(
        self,
        speech: torch.Tensor,
        speech_lengths: torch.Tensor,
        text: torch.Tensor,
        text_lengths: torch.Tensor,
        **kwargs,
    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]:
        """Encoder + Decoder + Calc loss
        Args:
                speech: (Batch, Length, ...)
                speech_lengths: (Batch, )
                text: (Batch, Length)
                text_lengths: (Batch,)
        """
        decoding_ind = kwargs.get("decoding_ind")
        if len(text_lengths.size()) > 1:
            text_lengths = text_lengths[:, 0]
        if len(speech_lengths.size()) > 1:
            speech_lengths = speech_lengths[:, 0]

        batch_size = speech.shape[0]

        # Encoder
        if hasattr(self.encoder, "overlap_chunk_cls"):
            ind = self.encoder.overlap_chunk_cls.random_choice(self.training, decoding_ind)
            encoder_out, encoder_out_lens = self.encode(speech, speech_lengths, ind=ind)
        else:
            encoder_out, encoder_out_lens = self.encode(speech, speech_lengths)

        # decoder: CTC branch
        if hasattr(self.encoder, "overlap_chunk_cls"):
            encoder_out_ctc, encoder_out_lens_ctc = self.encoder.overlap_chunk_cls.remove_chunk(
                encoder_out, encoder_out_lens, chunk_outs=None
            )
        else:
            encoder_out_ctc, encoder_out_lens_ctc = encoder_out, encoder_out_lens

        loss_ctc, cer_ctc = self._calc_ctc_loss(
            encoder_out_ctc, encoder_out_lens_ctc, text, text_lengths
        )

        # Collect CTC branch stats
        stats = dict()
        stats["loss_ctc"] = loss_ctc.detach() if loss_ctc is not None else None
        stats["cer_ctc"] = cer_ctc

        loss = loss_ctc

        stats["cer"] = cer_ctc
        stats["loss"] = torch.clone(loss.detach())

        # force_gatherable: to-device and to-tensor if scalar for DataParallel
        loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device)
        return loss, stats, weight

    def encode_chunk(
        self,
        speech: torch.Tensor,
        speech_lengths: torch.Tensor,
        cache: dict = None,
        **kwargs,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Frontend + Encoder. Note that this method is used by asr_inference.py
        Args:
                speech: (Batch, Length, ...)
                speech_lengths: (Batch, )
                ind: int
        """
        with autocast(False):
            # Data augmentation
            if self.specaug is not None and self.training:
                speech, speech_lengths = self.specaug(speech, speech_lengths)

            # Normalization for feature: e.g. Global-CMVN, Utterance-CMVN
            if self.normalize is not None:
                speech, speech_lengths = self.normalize(speech, speech_lengths)

        # Forward encoder
        encoder_out, encoder_out_lens, _ = self.encoder.forward_chunk(
            speech, speech_lengths, cache=cache["encoder"]
        )

        if isinstance(encoder_out, tuple):
            encoder_out = encoder_out[0]

        return encoder_out, torch.tensor([encoder_out.size(1)])

    def init_cache(self, cache: dict = {}, **kwargs):
        chunk_size = kwargs.get("chunk_size", [0, 10, 5])
        encoder_chunk_look_back = kwargs.get("encoder_chunk_look_back", 0)
        decoder_chunk_look_back = kwargs.get("decoder_chunk_look_back", 0)
        batch_size = 1

        enc_output_size = kwargs["encoder_conf"]["output_size"]
        feats_dims = kwargs["frontend_conf"]["n_mels"] * kwargs["frontend_conf"]["lfr_m"]
        cache_encoder = {
            "start_idx": 0,
            "cif_hidden": torch.zeros((batch_size, 1, enc_output_size)),
            "cif_alphas": torch.zeros((batch_size, 1)),
            "encoder_out": None,
            "encoder_out_lens": None,
            "chunk_size": chunk_size,
            "encoder_chunk_look_back": encoder_chunk_look_back,
            "last_chunk": False,
            "opt": None,
            "feats": torch.zeros((batch_size, chunk_size[0] + chunk_size[2], feats_dims)),
            "tail_chunk": False,
        }
        cache["encoder"] = cache_encoder

        cache_decoder = {
            "decode_fsmn": None,
            "decoder_chunk_look_back": decoder_chunk_look_back,
            "opt": None,
            "chunk_size": chunk_size,
        }
        cache["decoder"] = cache_decoder
        cache["frontend"] = {}
        cache["prev_samples"] = torch.empty(0)

        return cache

    def generate_chunk(
        self,
        speech,
        speech_lengths=None,
        key: list = None,
        tokenizer=None,
        frontend=None,
        **kwargs,
    ):
        cache = kwargs.get("cache", {})
        speech = speech.to(device=kwargs["device"])
        speech_lengths = speech_lengths.to(device=kwargs["device"])

        # Encoder
        is_final = kwargs.get("is_final", False)
        encoder_out, encoder_out_lens = self.encode_chunk(
            speech, speech_lengths, cache=cache, is_final=is_final
        )
        if isinstance(encoder_out, tuple):
            encoder_out = encoder_out[0]

        chunk_size = cache["encoder"]["chunk_size"]
        real_start_pos = chunk_size[0]

        if encoder_out_lens[0] > chunk_size[0] + chunk_size[1] + chunk_size[2]:
            assert False, print("impossible case 1 !")
        if encoder_out_lens[0] == chunk_size[0] + chunk_size[1] + chunk_size[2]:
            real_end_pos = chunk_size[0] + chunk_size[1]
        elif encoder_out_lens[0] > chunk_size[0] + chunk_size[1]:
            real_end_pos = chunk_size[0] + chunk_size[1]
        elif encoder_out_lens[0] > chunk_size[0]:
            real_end_pos = encoder_out_lens[0]
        else:
            assert False, print("impossible case 2 !")

        encoder_out_accum = cache["encoder"]["encoder_out"]
        if encoder_out_accum is not None:
            encoder_out_accum = torch.cat((encoder_out_accum, encoder_out[:, real_start_pos:real_end_pos, :]), dim=1)
        else:
            encoder_out_accum = encoder_out[:, real_start_pos:real_end_pos, :]
        cache["encoder"]["encoder_out"] = encoder_out_accum

        if cache["encoder"]["encoder_out_lens"] is not None:
            cache["encoder"]["encoder_out_lens"][0] += real_end_pos - real_start_pos
        else:
            cache["encoder"]["encoder_out_lens"] = encoder_out_lens
            cache["encoder"]["encoder_out_lens"][0] = real_end_pos - real_start_pos

        if is_final:
            if kwargs.get("output_dir") is not None:
                if not hasattr(self, "writer"):
                    self.writer = DatadirWriter(kwargs.get("output_dir"))

            results = []
            for i in range(encoder_out_accum.size(0)):
                x = encoder_out_accum[i, : cache["encoder"]["encoder_out_lens"][i], :]
                detect_result = self.kws_decoder.decode(x)
                is_deted, det_keyword, det_score = detect_result[0], detect_result[1], detect_result[2]

                if is_deted:
                    self.writer["detect"][key[i]] = "detected " + det_keyword + " " + str(det_score)
                    det_info = "detected " + det_keyword + " " + str(det_score)
                else:
                    self.writer["detect"][key[i]] = "rejected"
                    det_info = "rejected"

                result_i = {"key": key[i], "text": det_info}
                results.append(result_i)

            return results
        else:
            return None

    def inference(
        self,
        data_in,
        data_lengths=None,
        key: list = None,
        tokenizer=None,
        frontend=None,
        cache: dict = {},
        **kwargs,
    ):
        keywords = kwargs.get("keywords")
        from funasr.utils.kws_utils import KwsCtcPrefixDecoder
        self.kws_decoder = KwsCtcPrefixDecoder(
            ctc=self.ctc,
            keywords=keywords,
            token_list=tokenizer.token_list,
            seg_dict=tokenizer.seg_dict,
        )

        meta_data = {}
        chunk_size = kwargs["chunk_size"]
        chunk_stride_samples = int(chunk_size[1] * 960)  # 600ms
        first_chunk_padding_samples = int(chunk_size[2] * 960)  # 600ms

        if len(cache) == 0:
            self.init_cache(cache, **kwargs)

        time1 = time.perf_counter()
        cfg = {"is_final": kwargs.get("is_final", False)}
        audio_sample_list = load_audio_text_image_video(
            data_in,
            fs=frontend.fs,
            audio_fs=kwargs.get("fs", 16000),
            data_type=kwargs.get("data_type", "sound"),
            tokenizer=tokenizer,
            cache=cfg,
        )
        _is_final = cfg["is_final"]  # if data_in is a file or url, set is_final=True

        time2 = time.perf_counter()
        meta_data["load_data"] = f"{time2 - time1:0.3f}"
        assert len(audio_sample_list) == 1, "batch_size must be set 1"

        audio_sample = torch.cat((cache["prev_samples"], audio_sample_list[0]))

        if len(audio_sample) < first_chunk_padding_samples:
            print("key: {}, audio is too short for inference {}".format(key, len(audio_sample)))

        audio_sample_pre = audio_sample[0 : first_chunk_padding_samples]
        feat_pre, feat_pre_lengths = extract_fbank(
            [audio_sample_pre],
            data_type=kwargs.get("data_type", "sound"),
            frontend=frontend,
            cache=cache["frontend"],
            is_final=False,
        )

        audio_sample = audio_sample[first_chunk_padding_samples:]
        audio_chunks = int(len(audio_sample) // chunk_stride_samples)

        for i in range(audio_chunks):
            if i == 0:
                cache["encoder"]["feats"][:, chunk_size[2] :, :] = feat_pre

            kwargs["is_final"] = False
            audio_sample_i = audio_sample[i * chunk_stride_samples : (i + 1) * chunk_stride_samples]

            if kwargs["is_final"] and len(audio_sample_i) < 960:
                cache["encoder"]["tail_chunk"] = True
                speech = cache["encoder"]["feats"]
                speech_lengths = torch.tensor([speech.shape[1]], dtype=torch.int64).to(
                    speech.device
                )
            else:
                # extract fbank feats
                speech, speech_lengths = extract_fbank(
                    [audio_sample_i],
                    data_type=kwargs.get("data_type", "sound"),
                    frontend=frontend,
                    cache=cache["frontend"],
                    is_final=kwargs["is_final"],
                )
            time3 = time.perf_counter()
            meta_data["extract_feat"] = f"{time3 - time2:0.3f}"
            meta_data["batch_data_time"] = (
                speech_lengths.sum().item() * frontend.frame_shift * frontend.lfr_n / 1000
            )

            results_chunk_i = self.generate_chunk(
                speech,
                speech_lengths,
                key=key,
                tokenizer=tokenizer,
                cache=cache,
                frontend=frontend,
                **kwargs,
            )

            # results_chunk_i must be None when is_final=False
            assert results_chunk_i is None

        # process tail samples
        tail_audio_sample = audio_sample[ audio_chunks * chunk_stride_samples: ]
        if len(tail_audio_sample) < 960:
            kwargs["is_final"] = _is_final
            cache["encoder"]["tail_chunk"] = True
            speech = cache["encoder"]["feats"]
            speech_lengths = torch.tensor([speech.shape[1]], dtype=torch.int64).to(
                speech.device
            )
            results_chunk_tail = self.generate_chunk(
                speech,
                speech_lengths,
                key=key,
                tokenizer=tokenizer,
                cache=cache,
                frontend=frontend,
                **kwargs,
            )
        elif len(tail_audio_sample) <= first_chunk_padding_samples:
            kwargs["is_final"] = _is_final
            # extract fbank feats
            # cache["encoder"]["tail_chunk"] = True  # cannot be true
            speech, speech_lengths = extract_fbank(
                [ tail_audio_sample ],
                data_type=kwargs.get("data_type", "sound"),
                frontend=frontend,
                cache=cache["frontend"],
                is_final=kwargs["is_final"],
            )
            results_chunk_tail = self.generate_chunk(
                speech,
                speech_lengths,
                key=key,
                tokenizer=tokenizer,
                cache=cache,
                frontend=frontend,
                **kwargs,
            )
        elif len(tail_audio_sample) > first_chunk_padding_samples and \
             len(tail_audio_sample) < chunk_stride_samples:
            kwargs["is_final"] = False
            # extract fbank feats
            speech, speech_lengths = extract_fbank(
                [ tail_audio_sample ],
                data_type=kwargs.get("data_type", "sound"),
                frontend=frontend,
                cache=cache["frontend"],
                is_final=kwargs["is_final"],
            )
            results_chunk = self.generate_chunk(
                speech,
                speech_lengths,
                key=key,
                tokenizer=tokenizer,
                cache=cache,
                frontend=frontend,
                **kwargs,
            )
            # results_chunk must be None when is_final=False
            assert results_chunk is None

            # push tail chunk
            kwargs["is_final"] = _is_final
            cache["encoder"]["tail_chunk"] = True
            speech = cache["encoder"]["feats"]
            speech_lengths = torch.tensor([speech.shape[1]], dtype=torch.int64).to(
                speech.device
            )
            results_chunk_tail = self.generate_chunk(
                speech,
                speech_lengths,
                key=key,
                tokenizer=tokenizer,
                cache=cache,
                frontend=frontend,
                **kwargs,
            )

        result = results_chunk_tail

        if _is_final:
            self.init_cache(cache, **kwargs)

        if kwargs.get("output_dir"):
            if not hasattr(self, "writer"):
                self.writer = DatadirWriter(kwargs.get("output_dir"))

        return result, meta_data

    def export(self, **kwargs):
        from .export_meta import export_rebuild_model

        models = export_rebuild_model(model=self, **kwargs)
        return models