#!/usr/bin/env python3

__doc__ = \
"""
Create a figure showing the timing and geometry of 3+1 antennas,
and the triangles between them.
The additional antenna is to show that baselines are shared between
triangles.

In code (and on stdout), the antennas have time offsets which can be determined iteratively up to an overall offset.
"""

import matplotlib.pyplot as plt
import numpy as np
from itertools import combinations

c_light = 3e8 # m/s

class Antenna:
    """
    Simple Antenna class
    """
    def __init__(self,x=0,y=0,z=0,t0=0,name=""):
        self.x = x
        self.y = y
        self.z = z
        self.t = t0
        self.name = name

    def __repr__(self):
        cls = self.__class__.__name__

        return f'{cls}(x={self.x!r},y={self.y!r},z={self.z!r},t0={self.t!r},name={self.name!r})'


def distance(x1, x2):
    """
    Calculate the Euclidean distance between two locations x1 and x2
    """

    assert type(x1) in [Antenna]
    x1 = np.array([x1.x, x1.y, x1.z])

    assert type(x2) in [Antenna]
    x2 = np.array([x2.x, x2.y, x2.z])

    return np.sqrt( np.sum( (x1-x2)**2 ) )


def phase_mod(phase, low=np.pi):
    """
    Modulo phase such that it falls within the
    interval $[-low, 2\pi - low)$.
    """
    return (phase + low) % (2*np.pi) - low


def antenna_triangles(antennas):
    return combinations(antennas, 3)


def antenna_baselines(antennas):
    return combinations(antennas, 2)


def plot_four_antenna_geometry(tx, ants, extra_ant=None, ax=None, line_kwargs={}, scatter_kwargs={}, scatter_zorder=5, annotate_td=True):

    default_line_kwargs = dict( color='grey', lw=3, alpha=0.7)
    default_scatter_kwargs = dict( color='grey', s=200)

    if ax is None:
        ax = plt.gca()
        ax.set_aspect('equal')

    line_kwargs = {**default_line_kwargs, **line_kwargs}
    scatter_kwargs = {**default_scatter_kwargs, **scatter_kwargs}

    if extra_ant is not None:
        all_ants = ants + [extra_ant]
    else:
        all_ants = ants


    # Plot Antennas + Tx
    for i, ant in enumerate([tx] + all_ants):
        ax.scatter(ant.x, ant.y, zorder=scatter_zorder, **scatter_kwargs)
        ax.annotate(ant.name, (ant.x, ant.y), ha='center', va='center',zorder=scatter_zorder)

    # Lines connecting Tx and ants
    tmp_line_kwargs = line_kwargs
    tmp_line_kwargs['linestyle'] = '--'
    for ant in ants:
        ax.plot([tx.x, ant.x], [tx.y, ant.y], **tmp_line_kwargs)

        if annotate_td:
            annot_kwargs = dict(va='bottom', ha='center')
            dx, dy = 0, 0.1
            x = (tx.x + ant.x)/2 + dx
            y = (tx.y + ant.y)/2 + dy


            ax.annotate(f'$t_{{d{ant.name}}}$', (x,y), **annot_kwargs)

    # Lines due to all Antennas (including extra_ant)
    if extra_ant is not None:
        line_offset = 0.08*np.array([1,1])
        for i, ant_triangle in enumerate(antenna_triangles(all_ants)):
            tmp_line_kwargs['color'] = None
            tmp_line_kwargs['linestyle'] = '--'
            tmp_line_kwargs['alpha'] = 0.4
            for j, ant in enumerate(antenna_baselines(ant_triangle)):
                a, b = ant[0], ant[1]
                if j == 1: # fix ordering
                    a, b == b, a

                dx, dy = (i-1)*line_offset

                l = ax.plot([ a.x+dx, b.x+dx], [a.y+dy, b.y+dy], **tmp_line_kwargs)
                line_kwargs['color'] = l[0].get_color()

    # Lines internal to ants triangle
    tmp_line_kwargs = line_kwargs
    tmp_line_kwargs['color'] = 'green'
    tmp_line_kwargs['linestyle'] = '-'
    tmp_line_kwargs['alpha'] = 0.7
    for j, ant_pair in enumerate(combinations(ants,2)):
        a, b = ant_pair[0], ant_pair[1]
        if j == 1: # fix ordering
            a, b = b, a

        ax.plot([ a.x, b.x], [a.y, b.y], **tmp_line_kwargs)

    return ax


if __name__ == "__main__":
    from argparse import ArgumentParser
    import os.path as path
    import sys

    parser = ArgumentParser(description=__doc__)
    parser.add_argument("fname", metavar="path/to/figure[/]", nargs="?", help="Location for generated figure, will append __file__ if a directory. If not supplied, figure is shown.")
    parser.add_argument('num_ant', help='Number of antennas to show 2, 3', default=3, type=int)
    parser.add_argument("--no-extra", dest='extra', action='store_false', help='Disable the extra (fourth) antenna')

    args = parser.parse_args()

    if args.fname is not None and path.isdir(args.fname):
        args.fname = path.join(args.fname, path.splitext(path.basename(__file__))[0] + ".pdf")

    tx = Antenna(name="T", x=-8, y=2, t0=0)

    ants = [
        Antenna(name='1', x=0, y= 0, t0=1 ),
        Antenna(name='2', x=2, y=-3, t0=4 ),
        Antenna(name='3', x=1, y= 3, t0=10 ),
        ]

    if args.num_ant < 3:
        ants = [ants[0], ants[1]]

    if args.extra:
        extra_ant = Antenna(name='4', x=4, y=-1, t0=-6)
        all_ants = ants + [extra_ant]
    else:
        extra_ant = None
        all_ants = ants

    ax = plot_four_antenna_geometry(tx, ants, extra_ant=extra_ant)

    if True:
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        ax.spines['bottom'].set_visible(False)
        ax.spines['left'].set_visible(False)

        ax.tick_params(left=False, bottom=False, labelleft=False, labelbottom=False)


    fig = ax.get_figure()

    if args.fname is not None:
        plt.savefig(args.fname)
    else:
        plt.show()

    if True:
        sys.exit(0)
    ########################
    ### Lol, show my calculations are right

    use_phase = False
    correct_time = True

    if args.extra:
        all_ants = ants + [extra_ant]
    else:
        all_ants = ants

    for ant in enumerate(all_ants):
        ant.offsets = []
        ant.backup_t = []

    for i, triangle in enumerate(antenna_triangles(all_ants)):
        print("Antenna Triangle({},{},{})".format(*[ant.name for ant in triangle]))

        sigma = np.zeros((3))
        for j, ant_pair in enumerate(antenna_baselines(triangle)):
            a, b = ant_pair[0], ant_pair[1]
            if j == 1: # fix ordering
                a, b = b, a

            if i == 0: # print sigma pairing for first triangle
                print('j={}: {},{}'.format(j, a.name, b.name))

            phys_Dt = (distance(tx, a) - distance(tx, b))/c_light
            meas_Dt = a.t - b.t

            if use_phase:
                f_beacon = 50e6 # Hz
                to_phase = lambda t: phase_mod(2*np.pi*f_beacon*t)

                phys_Dt = to_phase(phys_Dt)
                meas_Dt = to_phase(meas_Dt)

            sigma[j] = meas_Dt - phys_Dt

            if False:
                print(
                    "Dt'_{},{} = ".format(a.name, b.name)
                    + "{}".format(meas_Dt)
                    + " = {} - {}".format(a.t, b.t)
                    )
                print(
                    "Dt_{},{} = ".format(a.name, b.name)
                    + "{}".format(phys_Dt)
                    + " = {} - {}".format(distance(tx, a), distance(tx, b))
                    )
                print(
                    "sigma_{},{} = ".format(a.name, b.name)
                    + "{}".format(sigma[j])
                    )

        print("sigmas:", sigma)
        if use_phase:
            print("sigmas sum:", phase_mod(np.sum(sigma)))
        else:
            print("sigmas sum:", np.sum(sigma))

        # Take the first antenna as reference
        ref_idx = 0
        ref_sigma = sigma[ref_idx]
        sigma = sigma - ref_sigma

        ant_sigma = 1/3*np.array([0, sigma[1] + sigma[2], 2*sigma[1] - sigma[2]])

        for i in [1,2]:
            triangle[i].offsets.append(-ant_sigma[i])
            if correct_time:
                triangle[i].backup_t.append(triangle[i].t)
                triangle[i].t += -ant_sigma[i]

    if correct_time:
        for i, ant in enumerate(all_ants):
            print(i, ant.name, ant.offsets)
