import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import sys
import numpy as np
from matplotlib.colors import LinearSegmentedColormap
from matplotlib.colors import BoundaryNorm


def plot_images(
    num_sample_perclass=10, x=None, y=None, labels=None, title=None, cmap=None
):
    grid_x = num_sample_perclass + 1
    grid_y = len(labels)

    plt.figure(figsize=(grid_y, grid_x))
    gs1 = gridspec.GridSpec(grid_y, grid_x)
    gs1.update(wspace=0.025, hspace=0.05)

    font = {"family": "serif", "weight": "bold"}

    plt.suptitle(title)
    j = 0
    for i in range(grid_y):
        idxs = [0] + list(np.where(y == list(labels.keys())[i])[0][: grid_x - 1])
        label = labels[list(labels.keys())[i]]

        for k, idx in enumerate(idxs):
            ax1 = plt.subplot(gs1[j])

            if k == 0:
                ax1.text(0, 0.25, label, ha="right", wrap=True, fontdict=font)

            else:
                ax1.imshow(x[idx, ...], cmap=cmap)

            plt.axis("off")
            j += 1

    plt.show()


def plot_2D(x, y, title, axis="off"):
    BLUE, ORANGE = "#57B5E8", "#E69E00"
    plt.figure(figsize=(8, 8))
    plt.scatter(
        x[:, 0],
        x[:, 1],
        s=18,
        facecolors="none",
        edgecolors=np.array([BLUE, ORANGE])[y],
    )
    if axis == "off":
        plt.axis("off")
    elif axis == "on":
        plt.xlabel("x_1")
        plt.ylabel("x_2")
    else:
        print("incorrect values for arg: axis (on or off only)")
        sys.exit()

    plt.title(title)
    plt.show()


def plot_dna(df, label):
    matrix = df.values
    col_names = df.columns
    rows = np.arange(matrix.shape[0])
    cols = np.arange(matrix.shape[1])
    np.random.seed(3)
    np.random.shuffle(rows)
    np.random.shuffle(cols)

    matrix = matrix[:, cols[:100]].T
    matrix = matrix[:, rows]
    col_names = col_names[cols[:100]]
    label = label[rows]
    mat_min = np.min(matrix)
    mat_max = np.max(matrix)
    mat_min = -np.max([np.abs(mat_min), mat_max])
    mat_max = np.max([np.abs(mat_min), mat_max])
    matrix = np.ma.masked_where(np.abs(matrix) <= 0.3, matrix)

    plt.figure(figsize=(6, 12))
    cmap_list = ["red", "darkred", "green", "lime", "lightgreen"]
    cmap = LinearSegmentedColormap.from_list("Custom cmap", cmap_list, len(cmap_list))
    cmap.set_bad("black")

    bounds = np.linspace(
        mat_min + 6, mat_max - 6, 5
    )  # np.arange(mat_min + 6, mat_max - 6, 0.1)
    idx = np.searchsorted(bounds, 0)

    bounds = np.insert(bounds, idx, 0)
    norm = BoundaryNorm(bounds, cmap.N)

    plt.imshow(matrix, cmap=cmap, norm=norm)
    plt.xticks(np.arange(len(label)))
    plt.yticks(np.arange(len(col_names)))
    ax = plt.gca()
    ax.set_xticklabels(label, rotation=90)
    ax.set_yticklabels(col_names)
    ax.yaxis.tick_right()
    ax.tick_params(axis=u"both", which=u"both", labelsize=5, length=0.0)
    plt.tight_layout()
    fig = plt.gcf()
    # fig.set_size_inches((6, 12), forward=False)
    # fig.savefig("img/dna.png", dpi=200)
    plt.show()

[(15, 'matplotlib.pyplot.figure', 'plt.figure', True, 'import matplotlib.pyplot as plt\n'), (16, 'matplotlib.gridspec.GridSpec', 'gridspec.GridSpec', True, 'import matplotlib.gridspec as gridspec\n'), (21, 'matplotlib.pyplot.suptitle', 'plt.suptitle', True, 'import matplotlib.pyplot as plt\n'), (39, 'matplotlib.pyplot.show', 'plt.show', True, 'import matplotlib.pyplot as plt\n'), (44, 'matplotlib.pyplot.figure', 'plt.figure', True, 'import matplotlib.pyplot as plt\n'), (61, 'matplotlib.pyplot.title', 'plt.title', True, 'import matplotlib.pyplot as plt\n'), (62, 'matplotlib.pyplot.show', 'plt.show', True, 'import matplotlib.pyplot as plt\n'), (68, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (69, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (70, 'numpy.random.seed', 'np.random.seed', True, 'import numpy as np\n'), (71, 'numpy.random.shuffle', 'np.random.shuffle', True, 'import numpy as np\n'), (72, 'numpy.random.shuffle', 'np.random.shuffle', True, 'import numpy as np\n'), (78, 'numpy.min', 'np.min', True, 'import numpy as np\n'), (79, 'numpy.max', 'np.max', True, 'import numpy as np\n'), (84, 'matplotlib.pyplot.figure', 'plt.figure', True, 'import matplotlib.pyplot as plt\n'), (89, 'numpy.linspace', 'np.linspace', True, 'import numpy as np\n'), (92, 'numpy.searchsorted', 'np.searchsorted', True, 'import numpy as np\n'), (94, 'numpy.insert', 'np.insert', True, 'import numpy as np\n'), (95, 'matplotlib.colors.BoundaryNorm', 'BoundaryNorm', False, 'from matplotlib.colors import BoundaryNorm\n'), (97, 'matplotlib.pyplot.imshow', 'plt.imshow', True, 'import matplotlib.pyplot as plt\n'), (100, 'matplotlib.pyplot.gca', 'plt.gca', True, 'import matplotlib.pyplot as plt\n'), (105, 'matplotlib.pyplot.tight_layout', 'plt.tight_layout', True, 'import matplotlib.pyplot as plt\n'), (106, 'matplotlib.pyplot.gcf', 'plt.gcf', True, 'import matplotlib.pyplot as plt\n'), (109, 'matplotlib.pyplot.show', 'plt.show', True, 'import matplotlib.pyplot as plt\n'), (53, 'matplotlib.pyplot.axis', 'plt.axis', True, 'import matplotlib.pyplot as plt\n'), (28, 'matplotlib.pyplot.subplot', 'plt.subplot', True, 'import matplotlib.pyplot as plt\n'), (36, 'matplotlib.pyplot.axis', 'plt.axis', True, 'import matplotlib.pyplot as plt\n'), (55, 'matplotlib.pyplot.xlabel', 'plt.xlabel', True, 'import matplotlib.pyplot as plt\n'), (56, 'matplotlib.pyplot.ylabel', 'plt.ylabel', True, 'import matplotlib.pyplot as plt\n'), (59, 'sys.exit', 'sys.exit', False, 'import sys\n'), (81, 'numpy.abs', 'np.abs', True, 'import numpy as np\n'), (82, 'numpy.abs', 'np.abs', True, 'import numpy as np\n'), (50, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (80, 'numpy.abs', 'np.abs', True, 'import numpy as np\n')]