"""
Locations are wrappers around a Numpy N-dimensional
array.
"""

import numpy as np
from functools import partial

try:
    from travelsignal import TravelSignal
except ModuleNotFoundError:
    from .travelsignal import TravelSignal

class Location:
    """
    A location is a point designated by a spatial coordinate x.
    """

    def __init__(self, x):
        self.x = np.asarray(x)

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

    def __getitem__(self, key):
        return self.x[key]

    def __setitem__(self, key, val):
        self.x[key] = val

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

        return self.__class__(self.x + other)

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

        return self.__class__(self.x - other)

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

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

        return np.all(self.x == other)

    # math alias functions
    __radd__ = __add__
    __rsub__ = __sub__
    __rmul__ = __mul__

class Receiver(Location):
    """
    A location able to trace a signal over time.

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

    def recv(self, travel_signal: TravelSignal) -> TravelSignal:
        """
        Return a function that traces the signal as a function of time
        at the receiver's location
        """
        return partial(travel_signal, x_f=self.x)

    receive = recv

class Emitter(Location):
    """
    Emit a signal from position x_0 (and time t_0)
    """
    def emit(self, travel_signal: TravelSignal) -> TravelSignal:
        return partial(travel_signal, x_0=self.x)


if __name__ == "__main__":
    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d import axes3d

    # 2D showcase
    source = Emitter([1,1])
    
    antennae = [
        Receiver([2,3]),
        Receiver([10,10]),
        Receiver([-2,-3]),
        Receiver([-10,0]),
    ]

    fig, ax = plt.subplots()
    
    ax.set_title("Geometry of Emitter(s) and Antennae")
    ax.set_ylabel("y")
    ax.set_xlabel("x")
    ax.plot(*source.x, '*', label="Emitter")
    
    for j, ant in enumerate(antennae):
        ax.plot(*ant.x, '+', label="Antenna {}".format(j))

    ax.legend()
    fig.show()

    # 3D showcase
    source = Emitter([1,1,1])

    antennae = [
        Receiver([2,3,0]),
        Receiver([10,10,-5]),
        Receiver([-2,-3,9]),
        Receiver([-10,0,-5]),
    ]

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    
    ax.set_title("Geometry of Emitter(s) and Antennae")
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    ax.set_zlabel("z")
    ax.plot([source.x[0]], *source.x[1:], '*', label="Emitter")
    
    for j, ant in enumerate(antennae):
        ax.plot([ant.x[0]], *ant.x[1:], '+', label="Antenna {}".format(j))

    ax.legend()
    plt.show()