from functools import partial

from .location import Location
from ..sampling import Digitizer

class Antenna(Location):
    """
    A location able to interact with a signal.

    Either emitting or receiving.

    Optionally uses digitizer to transform the signal
    when receiving.
    """
    def __init__(self, x, digitizer=None):
        super().__init__(x)

        self.digitizer = digitizer

    def __repr__(self):
        return "Antenna({}, {})".format(repr(self.x), repr(self.x))

    def _digitise_partial(self, signal: callable, *args, digitise=True, **kwargs) -> callable:
        """
        A wrapper around functools.partial to support optionally
        digitising the returned signal.
        """
        if self.digitizer and digitise:
            signal = self.digitizer.digitise(signal)

        return partial(signal, *args, **kwargs)


    def emit(self, signal: callable, digitise=False) -> callable:
        """
        Return a function that emits a signal from the antenna's location
        """
        return self._digitise_partial(signal, x_0=self.x, digitise=digitise)

    def recv(self, signal: callable, digitise=True) -> callable:
        """
        Return a function that traces the signal as a function of time
        at the antenna's location
        """
        return self._digitise_partial(signal, x_f=self.x, digitise=digitise)

    receive = recv

    # math
    def __add__(self, other):
        if isinstance(other, Location):
            other = other.x

        return self.__class__(self.x + other, self.digitizer)

    def __sub__(self, other):
        if isinstance(other, Location):
            other = other.x

        return self.__class__(self.x - other, self.digitizer)

    def __mul__(self, other):
        return self.__class__(self.x * other, self.digitizer)

class Receiver(Antenna):
    """
    An antenna which main purpose is to trace a signal over time.

    Optionally applies a transformation to the traced signal.
    """
    def __repr__(self):
        return "Receiver({})".format(repr(self.x))

class Emitter(Antenna):
    """
    An antenna which main purpose is to emit a signal.
    """
    def __repr__(self):
        return "Emitter({})".format(repr(self.x))