import os
import sys
import torch

import numpy as np
import torch.nn as nn
import onnxruntime as ort
import torch.nn.functional as F

from einops import rearrange
from torch.nn.utils.parametrizations import weight_norm

sys.path.append(os.getcwd())
os.environ["LRU_CACHE_CAPACITY"] = "3"

from infer.lib.predictors.FCPE.wav2mel import Wav2Mel
from infer.lib.predictors.FCPE.encoder import EncoderLayer, ConformerNaiveEncoder
from infer.lib.predictors.FCPE.utils import batch_interp_with_replacement_detach, decrypt_model, DotDict

@torch.no_grad()
def cent_to_f0(cent):
    return 10 * 2 ** (cent / 1200)

@torch.no_grad()
def f0_to_cent(f0):
    return 1200 * (f0 / 10).log2()

@torch.no_grad()
def latent2cents_decoder(cent_table, y, threshold = 0.05, mask = True):
    if str(y.device).startswith("privateuseone"): 
        cent_table = cent_table.cpu()
        y = y.cpu()

    B, N, _ = y.size()
    ci = cent_table[None, None, :].expand(B, N, -1)
    rtn = (ci * y).sum(dim=-1, keepdim=True) / y.sum(dim=-1, keepdim=True)  

    if mask:
        confident = y.max(dim=-1, keepdim=True)[0]
        confident_mask = torch.ones_like(confident)
        confident_mask[confident <= threshold] = float("-INF")
        rtn = rtn * confident_mask

    return rtn

@torch.no_grad()
def latent2cents_local_decoder(cent_table, out_dims, y, threshold = 0.05, mask = True):
    if str(y.device).startswith("privateuseone"): 
        cent_table = cent_table.cpu()
        y = y.cpu()

    B, N, _ = y.size()
    ci = cent_table[None, None, :].expand(B, N, -1)
    confident, max_index = y.max(dim=-1, keepdim=True)

    local_argmax_index = torch.arange(0, 9).to(max_index.device) + (max_index - 4)
    local_argmax_index[local_argmax_index < 0] = 0
    local_argmax_index[local_argmax_index >= out_dims] = out_dims - 1

    y_l = y.gather(-1, local_argmax_index)
    rtn = (ci.gather(-1, local_argmax_index) * y_l).sum(dim=-1, keepdim=True) / y_l.sum(dim=-1, keepdim=True) 

    if mask:
        confident_mask = torch.ones_like(confident)
        confident_mask[confident <= threshold] = float("-INF")
        rtn = rtn * confident_mask

    return rtn

def cents_decoder(cent_table, y, confidence, threshold = 0.05, mask=True):
    if str(y.device).startswith("privateuseone"): 
        cent_table = cent_table.cpu()
        y = y.cpu()

    B, N, _ = y.size()
    rtn = (cent_table[None, None, :].expand(B, N, -1) * y).sum(dim=-1, keepdim=True) / y.sum(dim=-1, keepdim=True)

    if mask:
        confident = y.max(dim=-1, keepdim=True)[0]
        confident_mask = torch.ones_like(confident)
        confident_mask[confident <= threshold] = float("-INF")
        rtn = rtn * confident_mask

    return (rtn, confident) if confidence else rtn

def cents_local_decoder(cent_table, y, n_out, confidence, threshold = 0.05, mask=True):
    if str(y.device).startswith("privateuseone"): 
        cent_table = cent_table.cpu()
        y = y.cpu()

    B, N, _ = y.size()
    confident, max_index = y.max(dim=-1, keepdim=True)
    local_argmax_index = (torch.arange(0, 9).to(max_index.device) + (max_index - 4)).clamp(0, n_out - 1)
    y_l = y.gather(-1, local_argmax_index)
    rtn = (cent_table[None, None, :].expand(B, N, -1).gather(-1, local_argmax_index) * y_l).sum(dim=-1, keepdim=True) / y_l.sum(dim=-1, keepdim=True)

    if mask:
        confident_mask = torch.ones_like(confident)
        confident_mask[confident <= threshold] = float("-INF")
        rtn = rtn * confident_mask

    return (rtn, confident) if confidence else rtn

class PCmer(nn.Module):
    def __init__(
        self, 
        num_layers, 
        num_heads, 
        dim_model, 
        dim_keys, 
        dim_values, 
        residual_dropout, 
        attention_dropout
    ):
        super().__init__()
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.dim_model = dim_model
        self.dim_values = dim_values
        self.dim_keys = dim_keys
        self.residual_dropout = residual_dropout
        self.attention_dropout = attention_dropout
        self._layers = nn.ModuleList([EncoderLayer(self) for _ in range(num_layers)])

    def forward(self, phone, mask=None):
        for layer in self._layers:
            phone = layer(phone, mask)

        return phone

class CFNaiveMelPE(nn.Module):
    def __init__(
        self, 
        input_channels, 
        out_dims, 
        hidden_dims = 512, 
        n_layers = 6, 
        n_heads = 8, 
        f0_max = 1975.5, 
        f0_min = 32.70, 
        use_fa_norm = False, 
        conv_only = False, 
        conv_dropout = 0, 
        atten_dropout = 0, 
        use_harmonic_emb = False
    ):
        super().__init__()
        self.input_channels = input_channels
        self.out_dims = out_dims
        self.hidden_dims = hidden_dims
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.f0_max = f0_max
        self.f0_min = f0_min
        self.use_fa_norm = use_fa_norm

        self.harmonic_emb = nn.Embedding(9, hidden_dims) if use_harmonic_emb else None
        self.input_stack = nn.Sequential(
            nn.Conv1d(
                input_channels, 
                hidden_dims, 
                3, 
                1, 
                1
            ), 
            nn.GroupNorm(
                4, 
                hidden_dims
            ), 
            nn.LeakyReLU(), 
            nn.Conv1d(
                hidden_dims, 
                hidden_dims, 
                3, 
                1, 
                1
            )
        )
        self.net = ConformerNaiveEncoder(
            num_layers=n_layers, 
            num_heads=n_heads, 
            dim_model=hidden_dims, 
            use_norm=use_fa_norm, 
            conv_only=conv_only, 
            conv_dropout=conv_dropout, 
            atten_dropout=atten_dropout
        )
        self.norm = nn.LayerNorm(hidden_dims)
        self.output_proj = weight_norm(
            nn.Linear(
                hidden_dims, 
                out_dims
            )
        )

        self.cent_table_b = torch.linspace(
            f0_to_cent(torch.Tensor([f0_min]))[0], 
            f0_to_cent(torch.Tensor([f0_max]))[0], 
            out_dims
        ).detach()
        self.gaussian_blurred_cent_mask_b = (
            1200 * torch.Tensor([self.f0_max / 10.]).log2()
        )[0].detach()

        self.register_buffer("cent_table", self.cent_table_b)
        self.register_buffer("gaussian_blurred_cent_mask", self.gaussian_blurred_cent_mask_b)

    def forward(self, x, _h_emb=None):
        x = self.input_stack(x.transpose(-1, -2)).transpose(-1, -2)

        if self.harmonic_emb is not None:
            if _h_emb is None:
                x += self.harmonic_emb(torch.LongTensor([0]).to(x.device))
            else:
                x += self.harmonic_emb(torch.LongTensor([int(_h_emb)]).to(x.device))

        return self.output_proj(self.norm(self.net(x))).sigmoid()

    @torch.no_grad()
    def infer(self, mel, decoder = "local_argmax", threshold = 0.05):
        latent = self.forward(mel)

        return cent_to_f0(
            (
                latent2cents_decoder(
                    self.cent_table, 
                    latent, 
                    threshold=threshold
                )
            ) if decoder == "argmax" else (
                latent2cents_local_decoder(
                    self.cent_table, 
                    self.out_dims, 
                    latent, 
                    threshold=threshold
                )
            )
        )

class FCPE_LEGACY(nn.Module):
    def __init__(
        self, 
        input_channel=128, 
        out_dims=360, 
        n_layers=12, 
        n_chans=512, 
        f0_max=1975.5, 
        f0_min=32.70, 
        confidence=False, 
        threshold=0.05, 
        use_input_conv=True
    ):
        super().__init__()
        self.n_out = out_dims
        self.f0_max = f0_max
        self.f0_min = f0_min
        self.confidence = confidence
        self.threshold = threshold
        self.use_input_conv = use_input_conv

        self.cent_table_b = torch.Tensor(
            np.linspace(
                f0_to_cent(torch.Tensor([f0_min]))[0], 
                f0_to_cent(torch.Tensor([f0_max]))[0], 
                out_dims
            )
        )
        self.register_buffer("cent_table", self.cent_table_b)

        self.stack = nn.Sequential(
            nn.Conv1d(
                input_channel, 
                n_chans, 
                3, 
                1, 
                1
            ), 
            nn.GroupNorm(
                4, 
                n_chans
            ), 
            nn.LeakyReLU(), 
            nn.Conv1d(
                n_chans, 
                n_chans, 
                3, 
                1, 
                1
            )
        )
        self.decoder = PCmer(
            num_layers=n_layers, 
            num_heads=8, 
            dim_model=n_chans, 
            dim_keys=n_chans, 
            dim_values=n_chans, 
            residual_dropout=0.1, 
            attention_dropout=0.1
        )
        self.norm = nn.LayerNorm(n_chans)
        self.dense_out = weight_norm(
            nn.Linear(
                n_chans, 
                self.n_out
            )
        )

    def forward(self, mel, return_hz_f0=False, cdecoder="local_argmax", output_interp_target_length=None):
        x = self.decoder(self.stack(mel.transpose(1, 2)).transpose(1, 2) if self.use_input_conv else mel)
        x = self.dense_out(self.norm(x)).sigmoid()

        x = cent_to_f0(
            (
                cents_decoder(
                    self.cent_table, 
                    x, 
                    self.confidence, 
                    threshold=self.threshold, 
                    mask=True
                )
            ) if cdecoder == "argmax" else (
                cents_local_decoder(
                    self.cent_table, 
                    x, 
                    self.n_out, 
                    self.confidence, 
                    threshold=self.threshold, 
                    mask=True
                )
            )
        )

        x = (1 + x / 700).log() if not return_hz_f0 else x

        if output_interp_target_length is not None: 
            x = F.interpolate(
                torch.where(x == 0, float("nan"), x).transpose(1, 2), 
                size=int(output_interp_target_length), 
                mode="linear"
            ).transpose(1, 2)

            x = torch.where(x.isnan(), float(0.0), x)

        return x

    def gaussian_blurred_cent(self, cents):
        B, N, _ = cents.size()

        return (
            -(self.cent_table[None, None, :].expand(B, N, -1) - cents).square() / 1250
        ).exp() * (cents > 0.1) & (
            cents < (1200.0 * np.log2(self.f0_max / 10.0))
        ).float()

class InferCFNaiveMelPE(torch.nn.Module):
    def __init__(
        self, 
        args, 
        state_dict
    ):
        super().__init__()
        self.model = CFNaiveMelPE(
            input_channels=args.mel.num_mels, 
            out_dims=args.model.out_dims, 
            hidden_dims=args.model.hidden_dims, 
            n_layers=args.model.n_layers, 
            n_heads=args.model.n_heads, 
            f0_max=args.model.f0_max, 
            f0_min=args.model.f0_min, 
            use_fa_norm=args.model.use_fa_norm, 
            conv_only=args.model.conv_only, 
            conv_dropout=args.model.conv_dropout, 
            atten_dropout=args.model.atten_dropout, 
            use_harmonic_emb=False
        )
        self.model.load_state_dict(state_dict)
        self.model.eval()
        self.register_buffer("tensor_device_marker", torch.tensor(1.0).float(), persistent=False)

    def forward(self, mel, decoder_mode = "local_argmax", threshold = 0.006):
        with torch.no_grad():
            mels = rearrange(torch.stack([mel], -1), "B T C K -> (B K) T C")
            f0s = rearrange(self.model.infer(mels, decoder=decoder_mode, threshold=threshold), "(B K) T 1 -> B T (K 1)", K=1)

        return f0s 

    def infer(
        self, 
        mel, 
        decoder_mode = "local_argmax", 
        threshold = 0.006, 
        f0_min = None, 
        f0_max = None, 
        interp_uv = False, 
        output_interp_target_length = None, 
        return_uv = False
    ):
        f0 = self.__call__(mel, decoder_mode, threshold)
        f0_for_uv = f0

        uv = (f0_for_uv < f0_min).type(f0_for_uv.dtype)
        f0 = f0 * (1 - uv)

        if interp_uv: 
            f0 = batch_interp_with_replacement_detach(
                uv.squeeze(-1).bool(), 
                f0.squeeze(-1)
            ).unsqueeze(-1)

        if f0_max is not None: f0[f0 > f0_max] = f0_max

        if output_interp_target_length is not None: 
            f0 = F.interpolate(
                torch.where(f0 == 0, float("nan"), f0).transpose(1, 2), 
                size=int(output_interp_target_length), 
                mode="linear"
            ).transpose(1, 2)

            f0 = torch.where(f0.isnan(), float(0.0), f0)

        if return_uv: return f0, F.interpolate(uv.transpose(1, 2), size=int(output_interp_target_length), mode="nearest").transpose(1, 2)
        else: return f0

class FCPEInfer_LEGACY:
    def __init__(
        self, 
        configs, 
        model_path, 
        device=None, 
        dtype=torch.float32, 
        providers=None, 
        onnx=False, 
        f0_min=50, 
        f0_max=1100
    ):
        if device is None: device = "cuda" if torch.cuda.is_available() else "cpu"
        self.device = device
        self.dtype = dtype
        self.onnx = onnx
        self.f0_min = f0_min
        self.f0_max = f0_max
        self.wav2mel = Wav2Mel(device=self.device, dtype=self.dtype)

        if self.onnx:
            sess_options = ort.SessionOptions()
            sess_options.log_severity_level = 3
            self.model = ort.InferenceSession(decrypt_model(configs, model_path), sess_options=sess_options, providers=providers)
        else:
            ckpt = torch.load(model_path, map_location="cpu", weights_only=True)
            self.args = DotDict(ckpt["config"])
            model = FCPE_LEGACY(
                input_channel=self.args.model.input_channel, 
                out_dims=self.args.model.out_dims, 
                n_layers=self.args.model.n_layers, 
                n_chans=self.args.model.n_chans,  
                f0_max=self.f0_max, 
                f0_min=self.f0_min, 
                confidence=self.args.model.confidence
            )
            model.to(self.device).to(self.dtype)
            model.load_state_dict(ckpt["model"])
            model.eval()
            self.model = model

    @torch.no_grad()
    def __call__(self, audio, sr, threshold=0.05, p_len=None):
        if not self.onnx: self.model.threshold = threshold
        if not hasattr(self, "numpy_threshold") and self.onnx: self.numpy_threshold = np.array(threshold, dtype=np.float32)

        mel = self.wav2mel(audio=audio[None, :], sample_rate=sr).to(self.dtype)

        if self.onnx:
            return torch.as_tensor(
                self.model.run(
                    [self.model.get_outputs()[0].name], 
                    {
                        self.model.get_inputs()[0].name: mel.detach().cpu().numpy(), 
                        self.model.get_inputs()[1].name: self.numpy_threshold
                    }
                )[0], 
                dtype=self.dtype, 
                device=self.device
            )
        else: 
            return self.model(
                mel=mel, 
                return_hz_f0=True, 
                output_interp_target_length=p_len
            )

class FCPEInfer:
    def __init__(
        self, 
        configs, 
        model_path, 
        device=None, 
        dtype=torch.float32, 
        providers=None, 
        onnx=False, 
        f0_min=50, 
        f0_max=1100
    ):
        if device is None: device = "cuda" if torch.cuda.is_available() else "cpu"
        self.device = device
        self.dtype = dtype
        self.onnx = onnx
        self.f0_min = f0_min
        self.f0_max = f0_max
        self.wav2mel = Wav2Mel(device=self.device, dtype=self.dtype)

        if self.onnx:
            sess_options = ort.SessionOptions()
            sess_options.log_severity_level = 3
            self.model = ort.InferenceSession(decrypt_model(configs, model_path), sess_options=sess_options, providers=providers)
        else:
            ckpt = torch.load(model_path, map_location="cpu", weights_only=True)
            ckpt["config_dict"]["model"]["conv_dropout"] = ckpt["config_dict"]["model"]["atten_dropout"] = 0.0
            self.args = DotDict(ckpt["config_dict"])
            model = InferCFNaiveMelPE(self.args, ckpt["model"])
            self.model = model.to(device).to(self.dtype).eval()

    @torch.no_grad()
    def __call__(self, audio, sr, threshold=0.05, p_len=None):
        if not hasattr(self, "numpy_threshold") and self.onnx: self.numpy_threshold = np.array(threshold, dtype=np.float32)
        mel = self.wav2mel(audio=audio[None, :], sample_rate=sr).to(self.dtype)

        if self.onnx:
            return torch.as_tensor(
                self.model.run(
                    [self.model.get_outputs()[0].name], 
                    {
                        self.model.get_inputs()[0].name: mel.detach().cpu().numpy(), 
                        self.model.get_inputs()[1].name: self.numpy_threshold
                    }
                )[0], 
                dtype=self.dtype, 
                device=self.device
            ) 
        else: 
            return self.model.infer(
                mel, 
                threshold=threshold, 
                f0_min=self.f0_min, 
                f0_max=self.f0_max, 
                output_interp_target_length=p_len
            )

class FCPE:
    def __init__(
        self, 
        configs, 
        model_path, 
        hop_length=512, 
        f0_min=50, 
        f0_max=1100, 
        dtype=torch.float32, 
        device=None, 
        sample_rate=16000, 
        threshold=0.05, 
        providers=None, 
        onnx=False, 
        legacy=False
    ):
        self.model = FCPEInfer_LEGACY if legacy else FCPEInfer
        self.fcpe = self.model(configs, model_path, device=device, dtype=dtype, providers=providers, onnx=onnx, f0_min=f0_min, f0_max=f0_max)
        self.hop_length = hop_length
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.threshold = threshold
        self.sample_rate = sample_rate
        self.dtype = dtype
        self.legacy = legacy

    def compute_f0(self, wav, p_len=None):
        x = torch.FloatTensor(wav).to(self.dtype).to(self.device)
        p_len = (x.shape[0] // self.hop_length) if p_len is None else p_len

        f0 = self.fcpe(x, sr=self.sample_rate, threshold=self.threshold, p_len=p_len)
        f0 = f0[:] if f0.dim() == 1 else f0[0, :, 0]

        if torch.all(f0 == 0): return f0.cpu().numpy() if p_len is None else np.zeros(p_len)
        return f0.cpu().numpy()