from functools import partial
import copy
from typing import Union

from .location import Location
from ..signals import Signal

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):
        super().__init__(x)

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

    def emit(self, signal: Union[Signal, callable]) -> Union[Signal, callable]:
        """
        Emit signal from this antenna's location.

        Note that this merely sets a default argument.
        """
        if not isinstance(signal, Signal):
            return partial(signal, x_0=self.x)
        else:
            new_signal = copy.copy(signal)
            new_signal.x_0 = self.x

            return new_signal

    def recv(self, signal: Union[Signal, callable]) -> Union[Signal, callable]:
        """
        Trace signal as a function of time at this antenna's
        location.

        Note that this merely sets a default argument.
        """
        if not isinstance(signal, Signal):
            return partial(signal, x_f=self.x)
        else:
            new_signal = copy.copy(signal)
            new_signal.x_f = self.x

            return new_signal

    receive = recv

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))