import importlib
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio
import sentencepiece as spm


def _sib(mod, *names):
    try:
        m = importlib.import_module("." + mod, __package__ or None)
    except (ImportError, TypeError, ValueError):
        m = importlib.import_module(mod)
    return [getattr(m, n) for n in names]


CPUGPTBlock, GLAMixer = _sib("model_cpu_gpt2", "CPUGPTBlock", "GLAMixer")

BLANK = 0


class BPE:
    def __init__(self, model_file="/workspace/asr_assets/bpe256.model"):
        self.sp = spm.SentencePieceProcessor(model_file=model_file)
        self.n = self.sp.get_piece_size()
        self.vocab = self.n + 1

    def encode(self, text):
        return [i + 1 for i in self.sp.encode(text.lower().strip())]

    def decode_ids(self, ids):
        pieces = [i - 1 for i in ids if i != BLANK]
        return self.sp.decode(pieces)


def greedy_decode_bpe(logp, bpe):
    ids = logp.argmax(-1).tolist()
    out, prev = [], None

    for i in ids:
        if i != prev and i != BLANK:
            out.append(i)
        prev = i
    return bpe.decode_ids(out)


class SpecAugment(nn.Module):
    def __init__(self, freq_mask=27, time_mask=70, n_freq=2, n_time=2, p=1.0):
        super().__init__()
        self.fm, self.tm, self.nf, self.nt, self.p = (
            freq_mask,
            time_mask,
            n_freq,
            n_time,
            p,
        )

    def forward(self, x):
        if not self.training:
            return x
        B, T, Fdim = x.shape
        for b in range(B):
            for _ in range(self.nf):
                f = torch.randint(0, self.fm + 1, (1,)).item()
                if f and Fdim - f > 0:
                    f0 = torch.randint(0, Fdim - f, (1,)).item()
                    x[b, :, f0 : f0 + f] = 0

            for _ in range(self.nt):
                tmax = min(self.tm, T // 2)
                if tmax <= 0:
                    continue
                t = torch.randint(0, tmax + 1, (1,)).item()

                if t and T - t > 0:
                    t0 = torch.randint(0, T - t, (1,)).item()
                    x[b, t0 : t0 + t, :] = 0
        return x


class ConvSubsample(nn.Module):
    def __init__(self, n_mels, d_model):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1, d_model, 3, stride=2, padding=1),
            nn.GELU(),
            nn.Conv2d(d_model, d_model, 3, stride=2, padding=1),
            nn.GELU(),
        )
        f = n_mels
        for _ in range(2):
            f = (f + 2 * 1 - 3) // 2 + 1
        self.lin = nn.Linear(d_model * f, d_model)

    def forward(self, x):
        x = x.unsqueeze(1)
        x = self.conv(x)
        b, c, t, fr = x.shape
        return self.lin(x.permute(0, 2, 1, 3).reshape(b, t, c * fr))


class FELACTC2(nn.Module):
    def __init__(self, cfg, vocab, n_mels=80, sr=16000):
        super().__init__()
        self.cfg = cfg
        self.vocab = vocab
        self.n_mels = n_mels
        self.melspec = torchaudio.transforms.MelSpectrogram(
            sample_rate=sr,
            n_fft=400,
            win_length=400,
            hop_length=160,
            n_mels=n_mels,
            power=2.0,
        )
        self.specaug = SpecAugment()
        self.subsample = ConvSubsample(n_mels, cfg.n_embd)
        self.blocks = nn.ModuleList([CPUGPTBlock(cfg, i) for i in range(cfg.n_layer)])
        self.ln_out = nn.RMSNorm(cfg.n_embd)
        self.head = nn.Linear(cfg.n_embd, vocab)
        self._init()

    def _init(self):
        for block in self.blocks:
            if isinstance(block.mixer, GLAMixer) and hasattr(block.mixer, "q_proj"):
                for p in [
                    block.mixer.q_proj,
                    block.mixer.k_proj,
                    block.mixer.v_proj,
                    block.mixer.out_proj,
                ]:
                    nn.init.normal_(p.weight, std=0.02)
            nn.init.normal_(block.ffn.gate.weight, std=0.02)
            nn.init.normal_(block.ffn.up.weight, std=0.02)
            nn.init.zeros_(block.ffn.down.weight)

    def features(self, wav, augment=False):
        m = self.melspec(wav)
        m = torch.log(m + 1e-6)
        m = (m - m.mean(dim=(1, 2), keepdim=True)) / (
            m.std(dim=(1, 2), keepdim=True) + 1e-5
        )
        m = m.transpose(1, 2)
        if augment:
            m = self.specaug(m)
        return m

    def forward(self, wav, augment=False):
        x = self.features(wav, augment=augment)
        x = self.subsample(x)
        for block in self.blocks:
            x = block(x)
        return F.log_softmax(self.head(self.ln_out(x)), dim=-1)

    def param_count(self):
        return sum(p.numel() for p in self.parameters())

    @torch.no_grad()
    def stream_logits(self, wav, frame_chunk=100, reset_frames=0):
        x = self.features(wav)
        x = self.subsample(x)
        B, Tf, d = x.shape
        states = [dict() for _ in self.blocks]
        since_reset = 0

        for c0 in range(0, Tf, frame_chunk):
            if reset_frames and since_reset >= reset_frames:
                states = [dict() for _ in self.blocks]
                since_reset = 0
            xc = x[:, c0 : c0 + frame_chunk]

            for l, b in enumerate(self.blocks):
                xc, states[l] = b.forward_chunk(xc, states[l])
            since_reset += xc.shape[1]

            yield F.log_softmax(self.head(self.ln_out(xc)), dim=-1)


def wer(ref, hyp):
    r, h = ref.split(), hyp.split()
    d = [[0] * (len(h) + 1) for _ in range(len(r) + 1)]
    for i in range(len(r) + 1):
        d[i][0] = i
    for j in range(len(h) + 1):
        d[0][j] = j
    for i in range(1, len(r) + 1):
        for j in range(1, len(h) + 1):
            c = 0 if r[i - 1] == h[j - 1] else 1
            d[i][j] = min(d[i - 1][j] + 1, d[i][j - 1] + 1, d[i - 1][j - 1] + c)

    return d[len(r)][len(h)], len(r)