#!/usr/bin/env python3

import numpy as np
import matplotlib.pyplot as plt

from earsim import REvent


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, log_max=False, plot_names=True, plot_max_values=True,**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)
    else:
        max_vals = None

    sc = ax.scatter(x, y, c=max_vals, **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

if __name__ == "__main__":
    fname = "ZH_airshower/mysim.sry"
    i = 0

    ev = REvent(fname)

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

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

    plt.show()