# Copyright (C) 2015-2024: The University of Edinburgh, United Kingdom
#                 Authors: Craig Warren, Antonis Giannopoulos, and John Hartley
#
# This file is part of gprMax.
#
# gprMax is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# gprMax is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with gprMax.  If not, see <http://www.gnu.org/licenses/>.

import argparse
import logging
from pathlib import Path

import h5py
import matplotlib.gridspec as gridspec
import matplotlib.pyplot as plt
import numpy as np

logger = logging.getLogger(__name__)


def calculate_antenna_params(filename, tltxnumber=1, tlrxnumber=None, rxnumber=None, rxcomponent=None):
    """Calculates antenna parameters - incident, reflected and total volatges
        and currents; s11, (s21) and input impedance.

    Args:
        filename: string of filename (including path) of output file.
        tltxnumber: int for transmitter antenna - transmission line number
        tlrxnumber: int for receiver antenna - transmission line number
        rxnumber: int for receiver antenna - output number
        rxcomponent: in for receiver antenna - output electric field component

    Returns:
        antennaparams: dict of antenna parameters.
    """

    # Open output file and read some attributes
    file = Path(filename)
    f = h5py.File(file, "r")
    dxdydz = f.attrs["dx_dy_dz"]
    dt = f.attrs["dt"]
    iterations = f.attrs["Iterations"]

    # Calculate time array and frequency bin spacing
    time = np.linspace(0, (iterations - 1) * dt, num=iterations)
    df = 1 / np.amax(time)

    logger.info(f"Time window: {np.amax(time):g} s ({iterations} iterations)")
    logger.info(f"Time step: {dt:g} s")
    logger.info(f"Frequency bin spacing: {df:g} Hz")

    # Read/calculate voltages and currents from transmitter antenna
    tltxpath = "/tls/tl" + str(tltxnumber) + "/"

    # Incident voltages/currents
    Vinc = f[tltxpath + "Vinc"][:]
    Iinc = f[tltxpath + "Iinc"][:]

    # Total (incident + reflected) voltages/currents
    Vtotal = f[tltxpath + "Vtotal"][:]
    Itotal = f[tltxpath + "Itotal"][:]

    # Reflected voltages/currents
    Vref = Vtotal - Vinc
    Iref = Itotal - Iinc

    # If a receiver antenna is used (with a transmission line or receiver),
    # get received voltage for s21
    if tlrxnumber:
        tlrxpath = "/tls/tl" + str(tlrxnumber) + "/"
        Vrec = f[tlrxpath + "Vtotal"][:]

    elif rxnumber:
        rxpath = "/rxs/rx" + str(rxnumber) + "/"
        availableoutputs = list(f[rxpath].keys())

        if rxcomponent not in availableoutputs:
            logger.exception(
                f"{rxcomponent} output requested, but the available "
                + f"output for receiver {rxnumber} is "
                + f"{', '.join(availableoutputs)}"
            )
            raise ValueError

        rxpath += rxcomponent

        # Received voltage
        if rxcomponent == "Ex":
            Vrec = f[rxpath][:] * -1 * dxdydz[0]
        elif rxcomponent == "Ey":
            Vrec = f[rxpath][:] * -1 * dxdydz[1]
        elif rxcomponent == "Ez":
            Vrec = f[rxpath][:] * -1 * dxdydz[2]
    f.close()

    # Frequency bins
    freqs = np.fft.fftfreq(Vinc.size, d=dt)

    # Delay correction - current lags voltage, so delay voltage to match
    # current timestep
    delaycorrection = np.exp(1j * 2 * np.pi * freqs * (dt / 2))

    # Calculate s11 and (optionally) s21
    s11 = np.abs(np.fft.fft(Vref) / np.fft.fft(Vinc))
    if tlrxnumber or rxnumber:
        s21 = np.abs(np.fft.fft(Vrec) / np.fft.fft(Vinc))

    # Calculate input impedance
    zin = (np.fft.fft(Vtotal) * delaycorrection) / np.fft.fft(Itotal)

    # Calculate input admittance
    yin = np.fft.fft(Itotal) / (np.fft.fft(Vtotal) * delaycorrection)

    # Convert to decibels (ignore warning from taking a log of any zero values)
    with np.errstate(divide="ignore"):
        Vincp = 20 * np.log10(np.abs((np.fft.fft(Vinc) * delaycorrection)))
        Iincp = 20 * np.log10(np.abs(np.fft.fft(Iinc)))
        Vrefp = 20 * np.log10(np.abs((np.fft.fft(Vref) * delaycorrection)))
        Irefp = 20 * np.log10(np.abs(np.fft.fft(Iref)))
        Vtotalp = 20 * np.log10(np.abs((np.fft.fft(Vtotal) * delaycorrection)))
        Itotalp = 20 * np.log10(np.abs(np.fft.fft(Itotal)))
        s11 = 20 * np.log10(s11)

    # Replace any NaNs or Infs from zero division
    Vincp[np.invert(np.isfinite(Vincp))] = 0
    Iincp[np.invert(np.isfinite(Iincp))] = 0
    Vrefp[np.invert(np.isfinite(Vrefp))] = 0
    Irefp[np.invert(np.isfinite(Irefp))] = 0
    Vtotalp[np.invert(np.isfinite(Vtotalp))] = 0
    Itotalp[np.invert(np.isfinite(Itotalp))] = 0
    s11[np.invert(np.isfinite(s11))] = 0

    # Create dictionary of antenna parameters
    antennaparams = {
        "time": time,
        "freqs": freqs,
        "Vinc": Vinc,
        "Vincp": Vincp,
        "Iinc": Iinc,
        "Iincp": Iincp,
        "Vref": Vref,
        "Vrefp": Vrefp,
        "Iref": Iref,
        "Irefp": Irefp,
        "Vtotal": Vtotal,
        "Vtotalp": Vtotalp,
        "Itotal": Itotal,
        "Itotalp": Itotalp,
        "s11": s11,
        "zin": zin,
        "yin": yin,
    }
    if tlrxnumber or rxnumber:
        with np.errstate(divide="ignore"):  # Ignore warning from taking a log of any zero values
            s21 = 20 * np.log10(s21)
        s21[np.invert(np.isfinite(s21))] = 0
        antennaparams["s21"] = s21

    return antennaparams


def mpl_plot(
    filename,
    time,
    freqs,
    Vinc,
    Vincp,
    Iinc,
    Iincp,
    Vref,
    Vrefp,
    Iref,
    Irefp,
    Vtotal,
    Vtotalp,
    Itotal,
    Itotalp,
    s11,
    zin,
    yin,
    s21=None,
    save=False,
):
    """Plots antenna parameters - incident, reflected and total voltages and
        currents; s11, (s21) and input impedance.

    Args:
        filename: string of filename (including path) of output file.
        time: array of simulation time.
        freq: array of frequencies for FFTs.
        Vinc, Vincp, Iinc, Iincp: arrays of time and frequency domain
                                    representations of incident voltage and
                                    current.
        Vref, Vrefp, Iref, Irefp: arrays of time and frequency domain
                                    representations of reflected voltage and
                                    current.
        Vtotal, Vtotalp, Itotal, Itotalp: arrays of time and frequency domain
                                            representations of total voltage and
                                            current.
        s11, s21: array(s) of s11 and, optionally, s21 parameters.
        zin, yin: arrays of input impedance and input admittance parameters.
        save: boolean flag to save plot to file.

    Returns:
        plt: matplotlib plot object.
    """

    # Set plotting range
    pltrangemin = 1
    # To a certain drop from maximum power
    pltrangemax = np.where((np.amax(Vincp[1::]) - Vincp[1::]) > 60)[0][0] + 1
    # To a maximum frequency
    pltrangemax = np.where(freqs > 3e9)[0][0]
    pltrange = np.s_[pltrangemin:pltrangemax]

    # Print some useful values from s11, and input impedance
    s11minfreq = np.where(s11[pltrange] == np.amin(s11[pltrange]))[0][0]
    logger.info(f"s11 minimum: {np.amin(s11[pltrange]):g} dB at " + f"{freqs[s11minfreq + pltrangemin]:g} Hz")
    logger.info(f"At {freqs[s11minfreq + pltrangemin]:g} Hz...")
    logger.info(
        f"Input impedance: {np.abs(zin[s11minfreq + pltrangemin]):.1f}"
        + f"{zin[s11minfreq + pltrangemin].imag:+.1f}j Ohms"
    )
    # logger.info(f'Input admittance (mag): {np.abs(yin[s11minfreq + pltrangemin]):g} S')
    # logger.info(f'Input admittance (phase): {np.angle(yin[s11minfreq + pltrangemin], deg=True):.1f} deg')

    # Figure 1
    # Plot incident voltage
    fig1, ax = plt.subplots(
        num="Transmitter transmission line parameters", figsize=(20, 12), facecolor="w", edgecolor="w"
    )
    gs1 = gridspec.GridSpec(4, 2, hspace=0.7)
    ax = plt.subplot(gs1[0, 0])
    ax.plot(time, Vinc, "r", lw=2, label="Vinc")
    ax.set_title("Incident voltage")
    ax.set_xlabel("Time [s]")
    ax.set_ylabel("Voltage [V]")
    ax.set_xlim([0, np.amax(time)])
    ax.grid(which="both", axis="both", linestyle="-.")

    # Plot frequency spectra of incident voltage
    ax = plt.subplot(gs1[0, 1])
    markerline, stemlines, baseline = ax.stem(freqs[pltrange], Vincp[pltrange], "-.", use_line_collection=True)
    plt.setp(baseline, "linewidth", 0)
    plt.setp(stemlines, "color", "r")
    plt.setp(markerline, "markerfacecolor", "r", "markeredgecolor", "r")
    ax.plot(freqs[pltrange], Vincp[pltrange], "r", lw=2)
    ax.set_title("Incident voltage")
    ax.set_xlabel("Frequency [Hz]")
    ax.set_ylabel("Power [dB]")
    ax.grid(which="both", axis="both", linestyle="-.")

    # Plot incident current
    ax = plt.subplot(gs1[1, 0])
    ax.plot(time, Iinc, "b", lw=2, label="Vinc")
    ax.set_title("Incident current")
    ax.set_xlabel("Time [s]")
    ax.set_ylabel("Current [A]")
    ax.set_xlim([0, np.amax(time)])
    ax.grid(which="both", axis="both", linestyle="-.")

    # Plot frequency spectra of incident current
    ax = plt.subplot(gs1[1, 1])
    markerline, stemlines, baseline = ax.stem(freqs[pltrange], Iincp[pltrange], "-.", use_line_collection=True)
    plt.setp(baseline, "linewidth", 0)
    plt.setp(stemlines, "color", "b")
    plt.setp(markerline, "markerfacecolor", "b", "markeredgecolor", "b")
    ax.plot(freqs[pltrange], Iincp[pltrange], "b", lw=2)
    ax.set_title("Incident current")
    ax.set_xlabel("Frequency [Hz]")
    ax.set_ylabel("Power [dB]")
    ax.grid(which="both", axis="both", linestyle="-.")

    # Plot total voltage
    ax = plt.subplot(gs1[2, 0])
    ax.plot(time, Vtotal, "r", lw=2, label="Vinc")
    ax.set_title("Total (incident + reflected) voltage")
    ax.set_xlabel("Time [s]")
    ax.set_ylabel("Voltage [V]")
    ax.set_xlim([0, np.amax(time)])
    ax.grid(which="both", axis="both", linestyle="-.")

    # Plot frequency spectra of total voltage
    ax = plt.subplot(gs1[2, 1])
    markerline, stemlines, baseline = ax.stem(freqs[pltrange], Vtotalp[pltrange], "-.", use_line_collection=True)
    plt.setp(baseline, "linewidth", 0)
    plt.setp(stemlines, "color", "r")
    plt.setp(markerline, "markerfacecolor", "r", "markeredgecolor", "r")
    ax.plot(freqs[pltrange], Vtotalp[pltrange], "r", lw=2)
    ax.set_title("Total (incident + reflected) voltage")
    ax.set_xlabel("Frequency [Hz]")
    ax.set_ylabel("Power [dB]")
    ax.grid(which="both", axis="both", linestyle="-.")

    # Plot total current
    ax = plt.subplot(gs1[3, 0])
    ax.plot(time, Itotal, "b", lw=2, label="Vinc")
    ax.set_title("Total (incident + reflected) current")
    ax.set_xlabel("Time [s]")
    ax.set_ylabel("Current [A]")
    ax.set_xlim([0, np.amax(time)])
    ax.grid(which="both", axis="both", linestyle="-.")

    # Plot frequency spectra of total current
    ax = plt.subplot(gs1[3, 1])
    markerline, stemlines, baseline = ax.stem(freqs[pltrange], Itotalp[pltrange], "-.", use_line_collection=True)
    plt.setp(baseline, "linewidth", 0)
    plt.setp(stemlines, "color", "b")
    plt.setp(markerline, "markerfacecolor", "b", "markeredgecolor", "b")
    ax.plot(freqs[pltrange], Itotalp[pltrange], "b", lw=2)
    ax.set_title("Total (incident + reflected) current")
    ax.set_xlabel("Frequency [Hz]")
    ax.set_ylabel("Power [dB]")
    ax.grid(which="both", axis="both", linestyle="-.")

    # Plot reflected (reflected) voltage
    # ax = plt.subplot(gs1[4, 0])
    # ax.plot(time, Vref, 'r', lw=2, label='Vref')
    # ax.set_title('Reflected voltage')
    # ax.set_xlabel('Time [s]')
    # ax.set_ylabel('Voltage [V]')
    # ax.set_xlim([0, np.amax(time)])
    # ax.grid(which='both', axis='both', linestyle='-.')

    # Plot frequency spectra of reflected voltage
    # ax = plt.subplot(gs1[4, 1])
    # markerline, stemlines, baseline = ax.stem(freqs[pltrange], Vrefp[pltrange],
    #                                           '-.', use_line_collection=True)
    # plt.setp(baseline, 'linewidth', 0)
    # plt.setp(stemlines, 'color', 'r')
    # plt.setp(markerline, 'markerfacecolor', 'r', 'markeredgecolor', 'r')
    # ax.plot(freqs[pltrange], Vrefp[pltrange], 'r', lw=2)
    # ax.set_title('Reflected voltage')
    # ax.set_xlabel('Frequency [Hz]')
    # ax.set_ylabel('Power [dB]')
    # ax.grid(which='both', axis='both', linestyle='-.')

    # Plot reflected (reflected) current
    # ax = plt.subplot(gs1[5, 0])
    # ax.plot(time, Iref, 'b', lw=2, label='Iref')
    # ax.set_title('Reflected current')
    # ax.set_xlabel('Time [s]')
    # ax.set_ylabel('Current [A]')
    # ax.set_xlim([0, np.amax(time)])
    # ax.grid(which='both', axis='both', linestyle='-.')

    # Plot frequency spectra of reflected current
    # ax = plt.subplot(gs1[5, 1])
    # markerline, stemlines, baseline = ax.stem(freqs[pltrange], Irefp[pltrange],
    #                                           '-.', use_line_collection=True)
    # plt.setp(baseline, 'linewidth', 0)
    # plt.setp(stemlines, 'color', 'b')
    # plt.setp(markerline, 'markerfacecolor', 'b', 'markeredgecolor', 'b')
    # ax.plot(freqs[pltrange], Irefp[pltrange], 'b', lw=2)
    # ax.set_title('Reflected current')
    # ax.set_xlabel('Frequency [Hz]')
    # ax.set_ylabel('Power [dB]')
    # ax.grid(which='both', axis='both', linestyle='-.')

    # Figure 2
    # Plot frequency spectra of s11
    fig2, ax = plt.subplots(num="Antenna parameters", figsize=(20, 12), facecolor="w", edgecolor="w")
    gs2 = gridspec.GridSpec(2, 2, hspace=0.3)
    ax = plt.subplot(gs2[0, 0])
    markerline, stemlines, baseline = ax.stem(freqs[pltrange], s11[pltrange], "-.", use_line_collection=True)
    plt.setp(baseline, "linewidth", 0)
    plt.setp(stemlines, "color", "g")
    plt.setp(markerline, "markerfacecolor", "g", "markeredgecolor", "g")
    ax.plot(freqs[pltrange], s11[pltrange], "g", lw=2)
    ax.set_title("s11")
    ax.set_xlabel("Frequency [Hz]")
    ax.set_ylabel("Power [dB]")
    # ax.set_xlim([0, 5e9])
    # ax.set_ylim([-25, 0])
    ax.grid(which="both", axis="both", linestyle="-.")

    # Plot frequency spectra of s21
    if s21 is not None:
        ax = plt.subplot(gs2[0, 1])
        markerline, stemlines, baseline = ax.stem(freqs[pltrange], s21[pltrange], "-.", use_line_collection=True)
        plt.setp(baseline, "linewidth", 0)
        plt.setp(stemlines, "color", "g")
        plt.setp(markerline, "markerfacecolor", "g", "markeredgecolor", "g")
        ax.plot(freqs[pltrange], s21[pltrange], "g", lw=2)
        ax.set_title("s21")
        ax.set_xlabel("Frequency [Hz]")
        ax.set_ylabel("Power [dB]")
        # ax.set_xlim([0.88e9, 1.02e9])
        # ax.set_ylim([-25, 50])
        ax.grid(which="both", axis="both", linestyle="-.")

    # Plot input resistance (real part of impedance)
    ax = plt.subplot(gs2[1, 0])
    markerline, stemlines, baseline = ax.stem(freqs[pltrange], zin[pltrange].real, "-.", use_line_collection=True)
    plt.setp(baseline, "linewidth", 0)
    plt.setp(stemlines, "color", "g")
    plt.setp(markerline, "markerfacecolor", "g", "markeredgecolor", "g")
    ax.plot(freqs[pltrange], zin[pltrange].real, "g", lw=2)
    ax.set_title("Input impedance (resistive)")
    ax.set_xlabel("Frequency [Hz]")
    ax.set_ylabel("Resistance [Ohms]")
    # ax.set_xlim([0.88e9, 1.02e9])
    ax.set_ylim(bottom=0)
    # ax.set_ylim([0, 300])
    ax.grid(which="both", axis="both", linestyle="-.")

    # Plot input reactance (imaginery part of impedance)
    ax = plt.subplot(gs2[1, 1])
    markerline, stemlines, baseline = ax.stem(freqs[pltrange], zin[pltrange].imag, "-.", use_line_collection=True)
    plt.setp(baseline, "linewidth", 0)
    plt.setp(stemlines, "color", "g")
    plt.setp(markerline, "markerfacecolor", "g", "markeredgecolor", "g")
    ax.plot(freqs[pltrange], zin[pltrange].imag, "g", lw=2)
    ax.set_title("Input impedance (reactive)")
    ax.set_xlabel("Frequency [Hz]")
    ax.set_ylabel("Reactance [Ohms]")
    # ax.set_xlim([0.88e9, 1.02e9])
    # ax.set_ylim([-300, 300])
    ax.grid(which="both", axis="both", linestyle="-.")

    # Plot input admittance (magnitude)
    # ax = plt.subplot(gs2[2, 0])
    # markerline, stemlines, baseline = ax.stem(freqs[pltrange], np.abs(yin[pltrange]),
    #                                           '-.', use_line_collection=True)
    # plt.setp(baseline, 'linewidth', 0)
    # plt.setp(stemlines, 'color', 'g')
    # plt.setp(markerline, 'markerfacecolor', 'g', 'markeredgecolor', 'g')
    # ax.plot(freqs[pltrange], np.abs(yin[pltrange]), 'g', lw=2)
    # ax.set_title('Input admittance (magnitude)')
    # ax.set_xlabel('Frequency [Hz]')
    # ax.set_ylabel('Admittance [Siemens]')
    # ax.set_xlim([0.88e9, 1.02e9])
    # ax.set_ylim([0, 0.035])
    # ax.grid(which='both', axis='both', linestyle='-.')

    # Plot input admittance (phase)
    # ax = plt.subplot(gs2[2, 1])
    # markerline, stemlines, baseline = ax.stem(freqs[pltrange], np.angle(yin[pltrange], deg=True),
    #                                           '-.', use_line_collection=True)
    # plt.setp(baseline, 'linewidth', 0)
    # plt.setp(stemlines, 'color', 'g')
    # plt.setp(markerline, 'markerfacecolor', 'g', 'markeredgecolor', 'g')
    # ax.plot(freqs[pltrange], np.angle(yin[pltrange], deg=True), 'g', lw=2)
    # ax.set_title('Input admittance (phase)')
    # ax.set_xlabel('Frequency [Hz]')
    # ax.set_ylabel('Phase [degrees]')
    # ax.set_xlim([0.88e9, 1.02e9])
    # ax.set_ylim([-40, 100])
    # ax.grid(which='both', axis='both', linestyle='-.')

    if save:
        savename1 = filename.stem + "_tl_params"
        savename1 = filename.parent / savename1
        savename2 = filename.stem + "_ant_params"
        savename2 = filename.parent / savename2
        # Save a PDF of the figure
        fig1.savefig(savename1.with_suffix(".pdf"), dpi=None, format="pdf", bbox_inches="tight", pad_inches=0.1)
        fig2.savefig(savename2.with_suffix(".pdf"), dpi=None, format="pdf", bbox_inches="tight", pad_inches=0.1)
        # Save a PNG of the figure
        # fig1.savefig(savename1.with_suffix('.png'), dpi=150, format='png',
        #             bbox_inches='tight', pad_inches=0.1)
        # fig2.savefig(savename2.with_suffix('.png'), dpi=150, format='png',
        #             bbox_inches='tight', pad_inches=0.1)

    return plt


if __name__ == "__main__":
    # Parse command line arguments
    parser = argparse.ArgumentParser(
        description="Plots antenna parameters - "
        + "incident, reflected and total voltages "
        + "and currents; s11, (s21) and input impedance "
        + "from an output file containing a transmission "
        + "line source.",
        usage="cd gprMax; python -m toolboxes.Plotting.plot_antenna_params outputfile",
    )
    parser.add_argument("outputfile", help="name of output file including path")
    parser.add_argument("--tltx-num", default=1, type=int, help="transmitter antenna - transmission line number")
    parser.add_argument("--tlrx-num", type=int, help="receiver antenna - transmission line number")
    parser.add_argument("--rx-num", type=int, help="receiver antenna - output number")
    parser.add_argument(
        "--rx-component",
        type=str,
        help="receiver antenna - output electric field component",
        choices=["Ex", "Ey", "Ez"],
    )
    parser.add_argument(
        "-save", action="store_true", default=False, help="save plot directly to file, i.e. do not display"
    )
    args = parser.parse_args()

    antennaparams = calculate_antenna_params(
        args.outputfile, args.tltx_num, args.tlrx_num, args.rx_num, args.rx_component
    )
    plthandle = mpl_plot(args.outputfile, **antennaparams, save=args.save)
    plthandle.show()