#!/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, **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="E{x}".format(x='x', i=i), **kwargs)
    if plot_Ey:
        ax.plot(antenna.t, antenna.Ey, label="E{x}".format(x='y', i=i), **kwargs)
    if plot_Ez:
        ax.plot(antenna.t, antenna.Ez, label="E{x}".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_max_values(antennas, ax=None, log_max=False, plot_names=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 ]
    
    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)


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

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

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

    fig2, ax2 = plt.subplots()
    plot_antenna_max_values(ev.antennas, ax=ax2)

    plt.show()