import os
import sys
import math
import torch
import librosa
import torchaudio

import numpy as np

sys.path.append(os.getcwd())

from main.library.predictors.FCN.model import MODEL
from main.library.predictors.FCN.convert import frequency_to_bins, seconds_to_samples, bins_to_frequency

CENTS_PER_BIN, PITCH_BINS, SAMPLE_RATE, WINDOW_SIZE = 5, 1440, 16000, 1024

class FCN:
    def __init__(self, model_path, hop_length=160, batch_size=None, f0_min=50, f0_max=1100, device=None, sample_rate=16000, providers=None, onnx=False):
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.hopsize = hop_length / SAMPLE_RATE
        self.batch_size = batch_size
        self.sample_rate = sample_rate
        self.onnx = onnx
        self.f0_min = f0_min
        self.f0_max = f0_max

        if self.onnx:
            import onnxruntime as ort

            sess_options = ort.SessionOptions()
            sess_options.log_severity_level = 3
            self.model = ort.InferenceSession(model_path, sess_options=sess_options, providers=providers)
        else:
            model = MODEL()
            ckpt = torch.load(model_path, map_location="cpu", weights_only=True)
            model.load_state_dict(ckpt['model'])
            model.eval()
            self.model = model.to(device)
    
    def entropy(self, logits):
        distribution = torch.nn.functional.softmax(logits, dim=1)
        return (1 + 1 / math.log(PITCH_BINS) * (distribution * torch.log(distribution + 1e-7)).sum(dim=1))
    
    def expected_frames(self, samples, center):
        hopsize_resampled = seconds_to_samples(self.hopsize, self.sample_rate)

        if center == 'half-window':
            window_size_resampled = WINDOW_SIZE / SAMPLE_RATE * self.sample_rate
            samples = samples - (window_size_resampled - hopsize_resampled)
        elif center == 'half-hop':
            samples = samples
        elif center == 'zero':
            samples = samples + hopsize_resampled
        else: raise ValueError

        return max(1, int(samples / hopsize_resampled))

    def resample(self, audio, target_rate=SAMPLE_RATE):
        if self.sample_rate == target_rate: return audio

        resampler = torchaudio.transforms.Resample(self.sample_rate, target_rate)
        resampler = resampler.to(audio.device)

        return resampler(audio)
    
    def preprocess(self, audio, center='half-window'):
        total_frames = self.expected_frames(audio.shape[-1], center)
        if self.sample_rate != SAMPLE_RATE: audio = self.resample(audio)

        hopsize = seconds_to_samples(self.hopsize, SAMPLE_RATE)
        if center in ['half-hop', 'zero']:
            if center == 'half-hop': padding = int((WINDOW_SIZE - hopsize) / 2)
            else: padding = int(WINDOW_SIZE / 2)

            audio = torch.nn.functional.pad(audio, (padding, padding), mode='reflect')

        if isinstance(hopsize, int) or hopsize.is_integer():
            hopsize = int(round(hopsize))
            start_idxs = None
        else:
            start_idxs = torch.round(torch.tensor([hopsize * i for i in range(total_frames + 1)])).int()

        batch_size = total_frames if self.batch_size is None else self.batch_size

        for i in range(0, total_frames, batch_size):
            batch = min(total_frames - i, batch_size)

            if start_idxs is None:
                start = i * hopsize
                end = start + int((batch - 1) * hopsize) + WINDOW_SIZE
                end = min(end, audio.shape[-1])
                batch_audio = audio[:, start:end]

                if end - start < WINDOW_SIZE:
                    padding = WINDOW_SIZE - (end - start)
                    remainder = (end - start) % hopsize

                    if remainder: padding += end - start - hopsize
                    batch_audio = torch.nn.functional.pad(batch_audio, (0, padding))

                frames = torch.nn.functional.unfold(batch_audio[:, None, None], kernel_size=(1, WINDOW_SIZE), stride=(1, hopsize)).permute(2, 0, 1)
            else:
                frames = torch.zeros(batch, 1, WINDOW_SIZE)

                for j in range(batch):
                    start = start_idxs[i + j]
                    end = min(start + WINDOW_SIZE, audio.shape[-1])
                    frames[j, :, : end - start] = audio[:, start:end]

            yield frames

    def viterbi(self, logits):
        if not hasattr(self, 'transition'):
            xx, yy = np.meshgrid(range(PITCH_BINS), range(PITCH_BINS))
            transition = np.maximum(12 - abs(xx - yy), 0)
            self.transition = transition / transition.sum(axis=1, keepdims=True)

        with torch.no_grad():
            probs = torch.nn.functional.softmax(logits, dim=1)

        bins = torch.tensor(np.array([librosa.sequence.viterbi(sequence, self.transition).astype(np.int64) for sequence in probs.cpu().numpy()]), device=probs.device)
        return bins_to_frequency(bins)

    def postprocess(self, logits):
        with torch.inference_mode():
            minidx = frequency_to_bins(torch.tensor(self.f0_min))
            maxidx = frequency_to_bins(torch.tensor(self.f0_max), torch.ceil)

            logits[:, :minidx] = -float('inf')
            logits[:, maxidx:] = -float('inf')

            pitch = self.viterbi(logits)
            periodicity = self.entropy(logits)

            return pitch.T, periodicity.T

    def compute_f0(self, audio, center = 'half-window'):
        if self.batch_size is not None: logits = []

        for frames in self.preprocess(audio, center):
            if self.onnx:
                inferred = torch.tensor(
                    self.model.run(
                        [self.model.get_outputs()[0].name], 
                        {
                            self.model.get_inputs()[0].name: frames.cpu().numpy()
                        }
                    )[0]
                ).detach()
            else:
                with torch.no_grad():
                    inferred = self.model(frames.to(self.device)).detach()

            logits.append(inferred)

        pitch, periodicity = self.postprocess(torch.cat(logits, 0).to(self.device))
        return pitch, periodicity