#!/usr/bin/env python3

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d

# fix package-internal importing
if __name__ == "__main__" and __package__ is None:
    import sys
    sys.path.append("../../")
    __package__ = "lib.location"

from . import location as loc
from ..location.antenna import Receiver, Emitter

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

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

fig, ax = plt.subplots()
loc.plot_geometry(ax, [source], antennae)
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()