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, 4, 15, 40), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((16, 10, 16, 43), 'matplotlib.gridspec.GridSpec', 'gridspec.GridSpec', ({(16, 28, 16, 34): 'grid_y', (16, 36, 16, 42): 'grid_x'}, {}), '(grid_y, grid_x)', True, 'import matplotlib.gridspec as gridspec\n'), ((21, 4, 21, 23), 'matplotlib.pyplot.suptitle', 'plt.suptitle', ({(21, 17, 21, 22): 'title'}, {}), '(title)', True, 'import matplotlib.pyplot as plt\n'), ((39, 4, 39, 14), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((44, 4, 44, 30), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((61, 4, 61, 20), 'matplotlib.pyplot.title', 'plt.title', ({(61, 14, 61, 19): 'title'}, {}), '(title)', True, 'import matplotlib.pyplot as plt\n'), ((62, 4, 62, 14), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((68, 11, 68, 37), 'numpy.arange', 'np.arange', ({(68, 21, 68, 36): 'matrix.shape[0]'}, {}), '(matrix.shape[0])', True, 'import numpy as np\n'), ((69, 11, 69, 37), 'numpy.arange', 'np.arange', ({(69, 21, 69, 36): 'matrix.shape[1]'}, {}), '(matrix.shape[1])', True, 'import numpy as np\n'), ((70, 4, 70, 21), 'numpy.random.seed', 'np.random.seed', ({(70, 19, 70, 20): '(3)'}, {}), '(3)', True, 'import numpy as np\n'), ((71, 4, 71, 27), 'numpy.random.shuffle', 'np.random.shuffle', ({(71, 22, 71, 26): 'rows'}, {}), '(rows)', True, 'import numpy as np\n'), ((72, 4, 72, 27), 'numpy.random.shuffle', 'np.random.shuffle', ({(72, 22, 72, 26): 'cols'}, {}), '(cols)', True, 'import numpy as np\n'), ((78, 14, 78, 28), 'numpy.min', 'np.min', ({(78, 21, 78, 27): 'matrix'}, {}), '(matrix)', True, 'import numpy as np\n'), ((79, 14, 79, 28), 'numpy.max', 'np.max', ({(79, 21, 79, 27): 'matrix'}, {}), '(matrix)', True, 'import numpy as np\n'), ((84, 4, 84, 31), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((89, 13, 91, 5), 'numpy.linspace', 'np.linspace', ({(90, 8, 90, 19): 'mat_min + 6', (90, 21, 90, 32): 'mat_max - 6', (90, 34, 90, 35): '5'}, {}), '(mat_min + 6, mat_max - 6, 5)', True, 'import numpy as np\n'), ((92, 10, 92, 36), 'numpy.searchsorted', 'np.searchsorted', ({(92, 26, 92, 32): 'bounds', (92, 34, 92, 35): '0'}, {}), '(bounds, 0)', True, 'import numpy as np\n'), ((94, 13, 94, 38), 'numpy.insert', 'np.insert', ({(94, 23, 94, 29): 'bounds', (94, 31, 94, 34): 'idx', (94, 36, 94, 37): '0'}, {}), '(bounds, idx, 0)', True, 'import numpy as np\n'), ((95, 11, 95, 39), 'matplotlib.colors.BoundaryNorm', 'BoundaryNorm', ({(95, 24, 95, 30): 'bounds', (95, 32, 95, 38): 'cmap.N'}, {}), '(bounds, cmap.N)', False, 'from matplotlib.colors import BoundaryNorm\n'), ((97, 4, 97, 44), 'matplotlib.pyplot.imshow', 'plt.imshow', (), '', True, 'import matplotlib.pyplot as plt\n'), ((100, 9, 100, 18), 'matplotlib.pyplot.gca', 'plt.gca', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((105, 4, 105, 22), 'matplotlib.pyplot.tight_layout', 'plt.tight_layout', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((106, 10, 106, 19), 'matplotlib.pyplot.gcf', 'plt.gcf', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((109, 4, 109, 14), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((53, 8, 53, 23), 'matplotlib.pyplot.axis', 'plt.axis', ({(53, 17, 53, 22): '"""off"""'}, {}), "('off')", True, 'import matplotlib.pyplot as plt\n'), ((28, 18, 28, 37), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(28, 30, 28, 36): 'gs1[j]'}, {}), '(gs1[j])', True, 'import matplotlib.pyplot as plt\n'), ((36, 12, 36, 27), 'matplotlib.pyplot.axis', 'plt.axis', ({(36, 21, 36, 26): '"""off"""'}, {}), "('off')", True, 'import matplotlib.pyplot as plt\n'), ((55, 8, 55, 25), 'matplotlib.pyplot.xlabel', 'plt.xlabel', ({(55, 19, 55, 24): '"""x_1"""'}, {}), "('x_1')", True, 'import matplotlib.pyplot as plt\n'), ((56, 8, 56, 25), 'matplotlib.pyplot.ylabel', 'plt.ylabel', ({(56, 19, 56, 24): '"""x_2"""'}, {}), "('x_2')", True, 'import matplotlib.pyplot as plt\n'), ((59, 8, 59, 18), 'sys.exit', 'sys.exit', ({}, {}), '()', False, 'import sys\n'), ((81, 22, 81, 37), 'numpy.abs', 'np.abs', ({(81, 29, 81, 36): 'mat_min'}, {}), '(mat_min)', True, 'import numpy as np\n'), ((82, 32, 82, 46), 'numpy.abs', 'np.abs', ({(82, 39, 82, 45): 'matrix'}, {}), '(matrix)', True, 'import numpy as np\n'), ((50, 19, 50, 43), 'numpy.array', 'np.array', ({(50, 28, 50, 42): '[BLUE, ORANGE]'}, {}), '([BLUE, ORANGE])', True, 'import numpy as np\n'), ((80, 23, 80, 38), 'numpy.abs', 'np.abs', ({(80, 30, 80, 37): 'mat_min'}, {}), '(mat_min)', True, 'import numpy as np\n')]