"""
Created on Dec 20, 2019

@author: @sahin-abdullah

This submodule produces two diffent figures: (1) emulator vs simulator,
(2) regression lines of first order component functions
"""
import matplotlib.pyplot as plt
import numpy as np


def plot(Si):
    # Close all figures
    plt.close("all")

    # Extract necessary data from Si
    try:
        problem = Si.problem
    except AttributeError:
        problem = Si

    Em = Si["Em"]
    RT = Si["RT"]
    Y_em = Si["Y_em"]
    idx = Si["idx"]
    X = Si["X"]
    Y = Si["Y"]

    # Get number of bootstrap from Runtime and sample size N
    K = RT.shape[0]
    N = Y.shape[0]
    row = 2
    col = 5

    try:
        ax = Si._plot()
    except AttributeError:
        # basic bar plot not found or is not available
        pass

    # Plot emulator performance
    Y_p = np.linspace(np.min(Y), np.max(Y), N, endpoint=True)

    start = max(0, K - 10)
    for i in range(start, K):
        # Only showing the last 10
        if (i % (row * col) == 0) or (i == 0):
            fig = plt.figure(frameon=False)
            fig.suptitle("Showing last 10 Bootstrap Trials")
            it = 1

        title_str = f"Trial {K - i}"
        ax = fig.add_subplot(row, col, it, frameon=True, title=title_str)
        ax.plot(Em["Y_e"][:, i], Y[idx[:, i]], "r+", label="Emulator")
        ax.plot(Y_p, Y_p, "darkgray", label="1:1 Line")
        ax.axis("square")
        ax.set_xlim(np.min(Y), np.max(Y))
        ax.set_ylim(np.min(Y), np.max(Y))
        it += 1

        if i == (K - 1):
            handles, labels = ax.get_legend_handles_labels()
            fig.legend(handles, labels, loc="center right", bbox_to_anchor=(1.0, 0.5))

    fig.tight_layout()

    # Now plot regression lines of component functions
    row = 3
    col = 1
    last_bootstrap = idx[:, -1]
    for i in range(problem["num_vars"]):
        if (i % (row * col) == 0) or (i == 0):
            fig = plt.figure(frameon=False)
            it = 1
        title_str = (
            "Regression of parameter " + problem["names"][i] + r"$^{(Last Trial)}$"
        )
        ax = fig.add_subplot(row, col, it, frameon=True, title=title_str)
        ax.plot(X[last_bootstrap, i], Y[last_bootstrap], "r.")
        ax.plot(X[last_bootstrap, i], np.mean(Em["f0"]) + Y_em[:, i], "k.")
        ax.legend(
            [r"$\widetilde{Y}$", "$f_" + str(i + 1) + "$"],
            loc="upper left",
            bbox_to_anchor=(1.04, 1.0),
        )
        it += 1

    fig.tight_layout()
    plt.show()

    if "emulated" in Si:
        emulated = Si["emulated"]
        # Sum of squared residuals
        Y_test = Si["Y_test"]
        ssr = np.sum((emulated - Y_test) ** 2)

        # Plot testing results
        fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(6, 6))
        ax.plot(emulated, Y_test, "r+", label="Emulator")
        ax.plot(Y_test, Y_test, "darkgray", label="1:1 Line")
        ax.axis("square")
        ax.set_xlabel("Emulator")
        ax.set_ylabel("New Observation")
        ax.set_xlim(np.min(Y_test), np.max(Y_test))
        ax.set_ylim(np.min(Y_test), np.max(Y_test))
        ax.legend(loc="lower right")
        plt.title(f"Testing results\nSSR = {ssr:.2f}")
        fig.tight_layout()
        plt.show()

    return ax


if __name__ == "__main__":
    pass
