import math
import random
from typing import Optional, Tuple
from fairseq.checkpoint_utils import load_model_ensemble_and_task
import numpy as np
import torch
import torch.nn.functional as F

# from fairseq.data.data_utils import compute_mask_indices
from fairseq.utils import index_put


# @torch.jit.script
def pad_to_multiple(x, multiple, dim=-1, value=0):
    # Inspired from https://github.com/lucidrains/local-attention/blob/master/local_attention/local_attention.py#L41
    if x is None:
        return None, 0
    tsz = x.size(dim)
    m = tsz / multiple
    remainder = math.ceil(m) * multiple - tsz
    if int(tsz % multiple) == 0:
        return x, 0
    pad_offset = (0,) * (-1 - dim) * 2

    return F.pad(x, (*pad_offset, 0, remainder), value=value), remainder


def extract_features(
    self,
    x,
    padding_mask=None,
    tgt_layer=None,
    min_layer=0,
):
    if padding_mask is not None:
        x = index_put(x, padding_mask, 0)

    x_conv = self.pos_conv(x.transpose(1, 2))
    x_conv = x_conv.transpose(1, 2)
    x = x + x_conv

    if not self.layer_norm_first:
        x = self.layer_norm(x)

    # pad to the sequence length dimension
    x, pad_length = pad_to_multiple(x, self.required_seq_len_multiple, dim=-2, value=0)
    if pad_length > 0 and padding_mask is None:
        padding_mask = x.new_zeros((x.size(0), x.size(1)), dtype=torch.bool)
        padding_mask[:, -pad_length:] = True
    else:
        padding_mask, _ = pad_to_multiple(
            padding_mask, self.required_seq_len_multiple, dim=-1, value=True
        )
    x = F.dropout(x, p=self.dropout, training=self.training)

    # B x T x C -> T x B x C
    x = x.transpose(0, 1)

    layer_results = []
    r = None
    for i, layer in enumerate(self.layers):
        dropout_probability = np.random.random() if self.layerdrop > 0 else 1
        if not self.training or (dropout_probability > self.layerdrop):
            x, (z, lr) = layer(
                x, self_attn_padding_mask=padding_mask, need_weights=False
            )
            if i >= min_layer:
                layer_results.append((x, z, lr))
        if i == tgt_layer:
            r = x
            break

    if r is not None:
        x = r

    # T x B x C -> B x T x C
    x = x.transpose(0, 1)

    # undo paddding
    if pad_length > 0:
        x = x[:, :-pad_length]

        def undo_pad(a, b, c):
            return (
                a[:-pad_length],
                b[:-pad_length] if b is not None else b,
                c[:-pad_length],
            )

        layer_results = [undo_pad(*u) for u in layer_results]

    return x, layer_results


def compute_mask_indices(
    shape: Tuple[int, int],
    padding_mask: Optional[torch.Tensor],
    mask_prob: float,
    mask_length: int,
    mask_type: str = "static",
    mask_other: float = 0.0,
    min_masks: int = 0,
    no_overlap: bool = False,
    min_space: int = 0,
    require_same_masks: bool = True,
    mask_dropout: float = 0.0,
) -> torch.Tensor:
    """
    Computes random mask spans for a given shape

    Args:
        shape: the the shape for which to compute masks.
            should be of size 2 where first element is batch size and 2nd is timesteps
        padding_mask: optional padding mask of the same size as shape, which will prevent masking padded elements
        mask_prob: probability for each token to be chosen as start of the span to be masked. this will be multiplied by
            number of timesteps divided by length of mask span to mask approximately this percentage of all elements.
            however due to overlaps, the actual number will be smaller (unless no_overlap is True)
        mask_type: how to compute mask lengths
            static = fixed size
            uniform = sample from uniform distribution [mask_other, mask_length*2]
            normal = sample from normal distribution with mean mask_length and stdev mask_other. mask is min 1 element
            poisson = sample from possion distribution with lambda = mask length
        min_masks: minimum number of masked spans
        no_overlap: if false, will switch to an alternative recursive algorithm that prevents spans from overlapping
        min_space: only used if no_overlap is True, this is how many elements to keep unmasked between spans
        require_same_masks: if true, will randomly drop out masks until same amount of masks remains in each sample
        mask_dropout: randomly dropout this percentage of masks in each example
    """

    bsz, all_sz = shape
    mask = torch.full((bsz, all_sz), False)

    all_num_mask = int(
        # add a random number for probabilistic rounding
        mask_prob * all_sz / float(mask_length)
        + torch.rand([1]).item()
    )

    all_num_mask = max(min_masks, all_num_mask)

    mask_idcs = []
    for i in range(bsz):
        if padding_mask is not None:
            sz = all_sz - padding_mask[i].long().sum().item()
            num_mask = int(mask_prob * sz / float(mask_length) + np.random.rand())
            num_mask = max(min_masks, num_mask)
        else:
            sz = all_sz
            num_mask = all_num_mask

        if mask_type == "static":
            lengths = torch.full([num_mask], mask_length)
        elif mask_type == "uniform":
            lengths = torch.randint(mask_other, mask_length * 2 + 1, size=[num_mask])
        elif mask_type == "normal":
            lengths = torch.normal(mask_length, mask_other, size=[num_mask])
            lengths = [max(1, int(round(x))) for x in lengths]
        else:
            raise Exception("unknown mask selection " + mask_type)

        if sum(lengths) == 0:
            lengths[0] = min(mask_length, sz - 1)

        if no_overlap:
            mask_idc = []

            def arrange(s, e, length, keep_length):
                span_start = torch.randint(low=s, high=e - length, size=[1]).item()
                mask_idc.extend(span_start + i for i in range(length))

                new_parts = []
                if span_start - s - min_space >= keep_length:
                    new_parts.append((s, span_start - min_space + 1))
                if e - span_start - length - min_space > keep_length:
                    new_parts.append((span_start + length + min_space, e))
                return new_parts

            parts = [(0, sz)]
            min_length = min(lengths)
            for length in sorted(lengths, reverse=True):
                t = [e - s if e - s >= length + min_space else 0 for s, e in parts]
                lens = torch.asarray(t, dtype=torch.int)
                l_sum = torch.sum(lens)
                if l_sum == 0:
                    break
                probs = lens / torch.sum(lens)
                c = torch.multinomial(probs.float(), len(parts)).item()
                s, e = parts.pop(c)
                parts.extend(arrange(s, e, length, min_length))
            mask_idc = torch.asarray(mask_idc)
        else:
            min_len = min(lengths)
            if sz - min_len <= num_mask:
                min_len = sz - num_mask - 1
            mask_idc = torch.asarray(
                random.sample([i for i in range(sz - min_len)], num_mask)
            )
            mask_idc = torch.asarray(
                [
                    mask_idc[j] + offset
                    for j in range(len(mask_idc))
                    for offset in range(lengths[j])
                ]
            )

        mask_idcs.append(torch.unique(mask_idc[mask_idc < sz]))

    min_len = min([len(m) for m in mask_idcs])
    for i, mask_idc in enumerate(mask_idcs):
        if isinstance(mask_idc, torch.Tensor):
            mask_idc = torch.asarray(mask_idc, dtype=torch.float)
        if len(mask_idc) > min_len and require_same_masks:
            mask_idc = torch.asarray(
                random.sample([i for i in range(mask_idc)], min_len)
            )
        if mask_dropout > 0:
            num_holes = int(round(len(mask_idc) * mask_dropout))
            mask_idc = torch.asarray(
                random.sample([i for i in range(mask_idc)], len(mask_idc) - num_holes)
            )

        mask[i, mask_idc.int()] = True

    return mask


def apply_mask(self, x, padding_mask, target_list):
    B, T, C = x.shape
    torch.zeros_like(x)
    if self.mask_prob > 0:
        mask_indices = compute_mask_indices(
            (B, T),
            padding_mask,
            self.mask_prob,
            self.mask_length,
            self.mask_selection,
            self.mask_other,
            min_masks=2,
            no_overlap=self.no_mask_overlap,
            min_space=self.mask_min_space,
        )
        mask_indices = mask_indices.to(x.device)
        x[mask_indices] = self.mask_emb
    else:
        mask_indices = None

    if self.mask_channel_prob > 0:
        mask_channel_indices = compute_mask_indices(
            (B, C),
            None,
            self.mask_channel_prob,
            self.mask_channel_length,
            self.mask_channel_selection,
            self.mask_channel_other,
            no_overlap=self.no_mask_channel_overlap,
            min_space=self.mask_channel_min_space,
        )
        mask_channel_indices = (
            mask_channel_indices.to(x.device).unsqueeze(1).expand(-1, T, -1)
        )
        x[mask_channel_indices] = 0

    return x, mask_indices


def get_hubert_model(
    model_path="assets/hubert/hubert_base.pt", device=torch.device("cpu")
):
    models, _, _ = load_model_ensemble_and_task(
        [model_path],
        suffix="",
    )
    hubert_model = models[0]
    hubert_model = hubert_model.to(device)

    def _apply_mask(x, padding_mask, target_list):
        return apply_mask(hubert_model, x, padding_mask, target_list)

    hubert_model.apply_mask = _apply_mask

    def _extract_features(
        x,
        padding_mask=None,
        tgt_layer=None,
        min_layer=0,
    ):
        return extract_features(
            hubert_model.encoder,
            x,
            padding_mask=padding_mask,
            tgt_layer=tgt_layer,
            min_layer=min_layer,
        )

    hubert_model.encoder.extract_features = _extract_features

    hubert_model._forward = hubert_model.forward

    def hubert_extract_features(
        self,
        source: torch.Tensor,
        padding_mask: Optional[torch.Tensor] = None,
        mask: bool = False,
        ret_conv: bool = False,
        output_layer: Optional[int] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        res = self._forward(
            source,
            padding_mask=padding_mask,
            mask=mask,
            features_only=True,
            output_layer=output_layer,
        )
        feature = res["features"] if ret_conv else res["x"]
        return feature, res["padding_mask"]

    def _hubert_extract_features(
        source: torch.Tensor,
        padding_mask: Optional[torch.Tensor] = None,
        mask: bool = False,
        ret_conv: bool = False,
        output_layer: Optional[int] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        return hubert_extract_features(
            hubert_model, source, padding_mask, mask, ret_conv, output_layer
        )

    hubert_model.extract_features = _hubert_extract_features

    def infer(source, padding_mask, output_layer: torch.Tensor):
        output_layer = output_layer.item()
        logits = hubert_model.extract_features(
            source=source, padding_mask=padding_mask, output_layer=output_layer
        )
        feats = hubert_model.final_proj(logits[0]) if output_layer == 9 else logits[0]
        return feats

    hubert_model.infer = infer
    # hubert_model.forward=infer
    # hubert_model.forward

    return hubert_model