import numpy as np
from functools import wraps, partial

from . import sampling as smp
from .sampler import Sampler

class Digitizer(Sampler):
    """
    Digitizer that takes in a signal and resamples and quantises the signal.
    """

    def __init__(self, resolution=0.1, bias=0, sampling_frequency=None):
        """

        Parameters
        ##########
        resolution - float
            Resolution of the digitizer
        sampling_frequency - float
            Frequency this digitizer will sample a signal
        """
        super().__init__(sampling_frequency)

        self.resolution = resolution
        self.bias = bias

    def digitise(self, signal, signal_sample_frequency=None):
        """
        Digitize signal according to the specs of this digitizer.

        Effectively resamples signal
        """

        if callable(signal):
            # if signal is already a partial,
            # try to rebuild it after setting the wrapper
            if isinstance(signal, partial):
                rebuild_partial = True
                p_args = signal.args
                p_kwargs = signal.keywords
                signal = signal.func
            else:
                rebuild_partial = False

            @wraps(signal)
            def wrapper(*args, **kwargs):
                return smp.quantise(
                        self.sample(signal(*args, **kwargs), signal_sample_frequency),
                        self.resolution,
                        self.bias
                        )

            # rebuild the partial if applicable
            if rebuild_partial:
                wrapper = partial(wrapper, *p_args, **p_kwargs)

            return wrapper

        else:
            signal = np.asarray(signal)

            return smp.quantise(
                self.sample(signal, signal_sample_frequency),
                self.resolution,
                self.bias
            )
