"""Demucsの実際の区間完了通知を、単調な作業進捗に変換する。音声演算は変更しない。"""
import math
import time
from contextlib import contextmanager


class InferenceProgress:
    def __init__(self, callback, models, shifts, length, samplerate, overlap):
        self.callback = callback
        self.shifts = max(1, shifts)
        self.models = models
        self.length = length + int(samplerate * 0.5)
        self.samplerate = samplerate
        self.overlap = overlap
        self.last = -1.0
        self.last_at = 0.0
        self.completed = [0] * len(models)

    def advance(self, index):
        if not self.callback:
            return
        if not 0 <= index < len(self.models):
            return
        stride = max(1, int((1 - self.overlap) * self.samplerate * self.models[index].segment))
        segments = max(1, math.ceil(self.length / stride))
        self.completed[index] += 1
        value = (index + min(1.0, self.completed[index] / (segments * self.shifts))) / len(self.models)
        value = min(0.999, max(self.last, value))
        now = time.monotonic()
        if value - self.last >= 0.005 or now - self.last_at >= 1:
            self.last, self.last_at = value, now
            self.callback({'progress': value})

    @contextmanager
    def track(self):
        # Demucs 4.0.1にはcallback引数がないため、公開PyTorch hookで区間完了だけを観測。
        handles = []
        try:
            if self.callback:
                for index, model in enumerate(self.models):
                    def observed(_module, _inputs, _output, index=index):
                        self.advance(index)
                    handles.append(model.register_forward_hook(observed))
            yield
        finally:
            for handle in handles:
                handle.remove()

    def finish(self):
        if self.callback:
            self.callback({'progress': 1.0})