# Copyright (C) 2015-2023: 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.pyplot as plt
import numpy as np

from ..Utilities.outputfiles_merge import get_output_data

logger = logging.getLogger(__name__)


def mpl_plot(filename, outputdata, dt, rxnumber, rxcomponent, save=False):
    """Creates a plot of the B-scan.

    Args:
        filename: string of filename (including path) of output file.
        outputdata: array of A-scans, i.e. B-scan data.
        dt: float of temporal resolution of the model.
        rxnumber: int of receiver output number.
        rxcomponent: string of receiver output field/current component.
        save: boolean flag to save plot to file.

    Returns:
        plt: matplotlib plot object.
    """

    file = Path(filename)

    fig = plt.figure(num=file.stem + " - rx" + str(rxnumber), figsize=(20, 10), facecolor="w", edgecolor="w")
    plt.imshow(
        outputdata,
        extent=[0, outputdata.shape[1], outputdata.shape[0] * dt, 0],
        interpolation="nearest",
        aspect="auto",
        cmap="seismic",
        vmin=-np.amax(np.abs(outputdata)),
        vmax=np.amax(np.abs(outputdata)),
    )
    plt.xlabel("Trace number")
    plt.ylabel("Time [s]")

    # Grid properties
    ax = fig.gca()
    ax.grid(which="both", axis="both", linestyle="-.")

    cb = plt.colorbar()
    if "E" in rxcomponent:
        cb.set_label("Field strength [V/m]")
    elif "H" in rxcomponent:
        cb.set_label("Field strength [A/m]")
    elif "I" in rxcomponent:
        cb.set_label("Current [A]")

    if save:
        # Save a PDF of the figure
        fig.savefig(filename[:-3] + ".pdf", dpi=None, format="pdf", bbox_inches="tight", pad_inches=0.1)
        # Save a PNG of the figure
        # fig.savefig(filename[:-3] + '.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 a B-scan image.",
        usage="cd gprMax; python -m toolboxes.Plotting.plot_Bscan outputfile output",
    )
    parser.add_argument("outputfile", help="name of output file including path")
    parser.add_argument(
        "rx_component",
        help="name of output component to be plotted",
        choices=["Ex", "Ey", "Ez", "Hx", "Hy", "Hz", "Ix", "Iy", "Iz"],
    )
    parser.add_argument(
        "-gather", action="store_true", default=False, help="gather together all receiver outputs in file"
    )
    parser.add_argument(
        "-save", action="store_true", default=False, help="save plot directly to file, i.e. do not display"
    )
    args = parser.parse_args()

    # Open output file and read number of outputs (receivers)
    f = h5py.File(args.outputfile, "r")
    nrx = f.attrs["nrx"]
    f.close()

    # Check there are any receivers
    if nrx == 0:
        logger.exception(f"No receivers found in {args.outputfile}")
        raise ValueError

    for rx in range(1, nrx + 1):
        outputdata, dt = get_output_data(args.outputfile, rx, args.rx_component)
        if args.gather:
            if rx == 1:
                rxsgather = outputdata
            rxsgather = np.column_stack((rxsgather, outputdata))
        else:
            plthandle = mpl_plot(args.outputfile, outputdata, dt, rx, args.rx_component, save=args.save)

    # Plot all receivers from single output file together if required
    if args.gather:
        plthandle = mpl_plot(args.outputfile, rxsgather, dt, rx, args.rx_component, save=args.save)

    plthandle.show()