# AUTOGENERATED! DO NOT EDIT! File to edit: nbdev/05_Masks.ipynb (unless otherwise specified).

__all__ = ['safe_div', 'MaskBase', 'MaskBinary', 'MaskcIRM', 'Maskify']

# Cell
from .imports import *
from .plot import *
from .utils import *
from .base import *
from .core import *

# Cell
def safe_div(x,y):
    return np.divide(x,y, out=np.zeros_like(x.data), where=y!=0)

# Cell
class MaskBase():
    _show_args={}
    def __init__(self, data): store_attr('data', self)
    @property
    def shape(self):
        return self.data.shape
    @classmethod
    def create(cls, specs, mix_spec):
        return [cls.generate(spec, mix_spec) for spec in specs]
    def __mul__(self, spec):
        raise NotImplementedError('This function needs to be implemented before use')
    def __rmul__(self, spec):
        return self*spec
    def generate(self, spec, mix_spec):
        raise NotImplementedError('This function needs to be implemented before use')
    @delegates(setup_graph)
    def show(self, ctx=None, **kwargs): return show_mask(self, ctx=ctx, **merge(self._show_args, kwargs))


# Cell
class MaskBinary(MaskBase):
    def __init__(self, data, threshold=1):
        store_attr('data, threshold', self)
    def __mul__(self, spec):
        if isinstance(spec, torch.Tensor): return spec*(self.data >= self.threshold)
        return SpecBase(spec.data*(self.data >= self.threshold), spec.sr, spec.fn)

    @classmethod
    def generate(cls, spec, mix_spec, threshold=0.9):
        Binary = (abs(spec.data) <= threshold)*1
        return cls(Binary)

    def show_binary(self):
        return show_mask((self.data >= self.threshold)*1)

# Cell
class MaskcIRM(MaskBase):
    def __mul__(self, spec):
        if isinstance(spec, SpecBase):
            return SpecBase(spec.data*self.data, spec.sr)
        elif self.data.requires_grad == True or spec.requires_grad == True:
            return complex_mult(self.data, spec)
        return real2complex(spec)*real2complex(self.data)

    @classmethod
    def generate(cls, spec, mix_spec):
        cIRM = safe_div(spec.data, mix_spec.data)
        return cls(cIRM)

# Cell
class Maskify(Transform):
    as_item_force=True
    def __init__(self, MaskType=MaskcIRM, Aud2Spec=Spectify()):
        store_attr("MaskType, Aud2Spec", self)
    def encodes(self, audioList)->None:
        specList = [self.Aud2Spec(a) for a in audioList]
        joined = AudioBase(join_audios(audioList), audioList[0].sr)
        mix_spec = self.Aud2Spec(joined)
        maskList = self.MaskType.create(specList, mix_spec)
        return mix_spec, maskList
    def decodes(self, spec_mask)->None:
        mix_spec, maskList = spec_mask
        return [self.Aud2Spec.decode(spec*m) for m in maskList]

def mask2tensor(mask:MaskBase):
    data = complex2real(mask.data)
    return TensorMask(data)

MaskBase._tensor_cls = TensorMask

@ToTensor
def encodes(self, o:MaskBase): return o._tensor_cls(mask2tensor(o))