#!/usr/bin/env python3

__doc__ = \
"""
Show geometry and time delay between radio antennas 
and a source (true, and expected).
"""

import matplotlib.pyplot as plt
from matplotlib.path import Path
import matplotlib.patches as patches

import numpy as np

def antenna_path(loc, size=1, height=(.5)**(.5), width=1, stem_height=None):
    stem_height = stem_height if stem_height is not None else height

    vertices = [
        (0, 0), # center, middle
        ( width/2*size,  height/2*size), #right, top
        ( width/2*size, -height/2*size), #right, bottom
        (-width/2*size,  height/2*size), #left, top
        (-width/2*size, -height/2*size), #left, bottom
        (0, 0), # back to center
        (0, -stem_height), # stem
        (0, 0), # back to center
    ]

    codes = [
            Path.MOVETO,
            Path.LINETO,
            Path.LINETO,
            Path.LINETO,
            Path.LINETO,
            Path.LINETO,
            Path.LINETO,
            Path.CLOSEPOLY,
    ]

    # modify vertices
    for i, v in enumerate(vertices):
        vertices[i] = ( loc[0] + v[0], loc[1] + v[1] )

    return Path(vertices, codes)


def radio_interferometry_figure(emit_loc=(3,8), resolve_loc=None, N_antenna=4, ant_size=0.5, **fig_kwargs):
    fig, ax = plt.subplots(**fig_kwargs)

    annot_kwargs = dict(
            color='red',
            fontsize=15,
    )

    ray_kwargs = dict(
            marker=None,
            ls='solid',
            color='red',
            alpha=0.8
    )
    antenna_patch_kwargs = dict(
            edgecolor='k',
            facecolor='none',
            lw=2
    )

    antenna_patches = []
    dx_antenna = 1.5
    stem_height = ant_size
    for i in range(N_antenna):
        path = antenna_path( (4+dx_antenna*i, stem_height), size=ant_size, stem_height=stem_height)
        patch = patches.PathPatch(path, **antenna_patch_kwargs)
        ax.add_patch(patch)
        antenna_patches.append(patch)

        if i == N_antenna - 1:
            ant_loc = path.vertices[0]
            ax.annotate("$\\vec{a_i}$", (ant_loc[0]+0.4, 0.8), va='top', ha='left', **{**annot_kwargs, **dict(color='k')})

    # ground level
    ax.axhline(0, color='k')

    # indicate antenna signal
    ant_loc = antenna_patches[int(2/4*N_antenna)].get_path().vertices[0]
    ax.annotate("$S_i(t)$", (ant_loc[0], ant_loc[1]-stem_height-0.2), va='top', ha='center', **annot_kwargs)


    if emit_loc or resolve_loc:
        # resolve_loc
        if resolve_loc is None:
            resolve_loc = emit_loc

        if resolve_loc:
            for i, antenna in enumerate(antenna_patches):
                    ant_loc = antenna.get_path().vertices[0]

                    # rays from the antenna to resolve_loc
                    ax.plot( (ant_loc[0], resolve_loc[0]), (ant_loc[1], resolve_loc[1]), **ray_kwargs)
                    if i == N_antenna - 1:
                        ax.annotate("$\Delta_i$", ( (ant_loc[0]+resolve_loc[0])/2, (ant_loc[1]+resolve_loc[1])/2 ), va='bottom', ha='left', **annot_kwargs)

            ax.plot(*resolve_loc, 'ro')
            ax.annotate("$S(\\vec{x}, t)$", resolve_loc, ha='left', va='bottom', **annot_kwargs)

        # emit loc
        if emit_loc:
            ax.plot(*emit_loc, 'ko')

            ax.annotate('$S_0$', emit_loc, ha='right', va='top', **{**annot_kwargs, **dict(color='k')})

    ax.set_xlim(-1, 10)
    ax.set_ylim(-1, 10)

    ax.axis('off')

    fig.tight_layout()

    return fig


if __name__ == "__main__":
    emit_loc = (2.5, 8)
    figsize = (6,6)

    from argparse import ArgumentParser
    import os.path as path

    parser = ArgumentParser(description=__doc__)
    parser.add_argument('scenario', choices=['base', 'true', 'closeby', 'far-away'], default='true')
    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.")

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

    ###
    emit_loc = None
    resolve_loc = None
    if args.scenario != 'base':
        emit_loc = (2.5, 8)

        if args.scenario == 'closeby':
            resolve_loc = (3, 8)
        elif args.scenario == 'far-away':
            resolve_loc = (8, 8)
    
    fig = radio_interferometry_figure(emit_loc=emit_loc, resolve_loc=resolve_loc, figsize=figsize)

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