#!/usr/bin/env python3

import numpy as np
import matplotlib.pyplot as plt

def plot_antenna_Efields(antenna, ax=None, plot_Ex=True, plot_Ey=True, plot_Ez=True, label_append="",**kwargs):
    """Show waveforms from an antenna"""
    if ax is None:
        ax = plt.gca()

    default_kwargs = dict(alpha=0.6)

    kwargs = {**default_kwargs, **kwargs}

    i = antenna.name

    if plot_Ex:
        ax.plot(antenna.t, antenna.Ex, label=f"E{{x}}{label_append}".format(x='x', i=i), **kwargs)
    if plot_Ey:
        ax.plot(antenna.t, antenna.Ey, label=f"E{{x}}{label_append}".format(x='y', i=i), **kwargs)
    if plot_Ez:
        ax.plot(antenna.t, antenna.Ez, label=f"E{{x}}{label_append}".format(x='z', i=i), **kwargs)

    ax.set_xlabel("[ns]")
    ax.set_ylabel("[$\mu$V/m]")
    ax.set_title("Antenna {i}".format(i=i))
    ax.legend()

    return ax

def plot_antenna_geometry(antennas, ax=None, plot_names=True, plot_max_values=True, log_max=False, colors=None,**kwargs):
    """Show the max values at each antenna."""
    default_kwargs = dict(
            cmap='Spectral_r'
            )

    kwargs = {**default_kwargs, **kwargs}

    x = [ a.x for a in antennas ]
    y = [ a.y for a in antennas ]

    if ax is None:
        ax = plt.gca()

    if plot_max_values:
        max_vals = np.asarray([ np.max([np.max(a.Ex), np.max(a.Ey), np.max(a.Ez)]) for a in antennas ])

        if log_max:
            max_vals = np.log10(max_vals)

        colors = max_vals

    sc = ax.scatter(x, y, c=colors, **kwargs)

    if plot_max_values:
        fig = ax.get_figure()
        fig.colorbar(sc, ax=ax, label="[$\mu$V/m]")

    if plot_names:
        [ ax.annotate(a.name, (a.x, a.y), ha='center',va='center') for a in antennas ]

    ax.set_title("Maximum E field at each antenna")
    ax.set_ylabel("[m]")
    ax.set_xlabel("[m]")


    return ax, sc

if __name__ == "__main__":
    import os.path as path

    from earsim import REvent

    import aa_generate_beacon as beacon

    fname = "ZH_airshower/mysim.sry"
    i = 0

    if True:
        ev = REvent(fname)
        antennas = ev.antennas
    else:

        fname_dir = path.dirname(fname)
        antennas_fname = path.join(fname_dir, beacon.antennas_fname)

        f_beacon, tx, antennas = beacon.read_beacon_hdf5(antennas_fname)

    if True:
       fig, ax1 = plt.subplots()
       plot_antenna_Efields(antennas[i], ax=ax1)

    if True:
       fig2, ax2 = plt.subplots()
       plot_antenna_geometry(antennas, ax=ax2, plot_max_values=True, plot_names=False)
       ax2.set_aspect('equal', 'datalim')

    plt.show()