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()

[((366, 402), 'matplotlib.pyplot.figure', 'plt.figure', ([], {'figsize': '(grid_y, grid_x)'}), '(figsize=(grid_y, grid_x))\n', (376, 402), True, 'import matplotlib.pyplot as plt\n'), ((413, 446), 'matplotlib.gridspec.GridSpec', 'gridspec.GridSpec', (['grid_y', 'grid_x'], {}), '(grid_y, grid_x)\n', (430, 446), True, 'import matplotlib.gridspec as gridspec\n'), ((544, 563), 'matplotlib.pyplot.suptitle', 'plt.suptitle', (['title'], {}), '(title)\n', (556, 563), True, 'import matplotlib.pyplot as plt\n'), ((1035, 1045), 'matplotlib.pyplot.show', 'plt.show', ([], {}), '()\n', (1043, 1045), True, 'import matplotlib.pyplot as plt\n'), ((1130, 1156), 'matplotlib.pyplot.figure', 'plt.figure', ([], {'figsize': '(8, 8)'}), '(figsize=(8, 8))\n', (1140, 1156), True, 'import matplotlib.pyplot as plt\n'), ((1523, 1539), 'matplotlib.pyplot.title', 'plt.title', (['title'], {}), '(title)\n', (1532, 1539), True, 'import matplotlib.pyplot as plt\n'), ((1544, 1554), 'matplotlib.pyplot.show', 'plt.show', ([], {}), '()\n', (1552, 1554), True, 'import matplotlib.pyplot as plt\n'), ((1643, 1669), 'numpy.arange', 'np.arange', (['matrix.shape[0]'], {}), '(matrix.shape[0])\n', (1652, 1669), True, 'import numpy as np\n'), ((1681, 1707), 'numpy.arange', 'np.arange', (['matrix.shape[1]'], {}), '(matrix.shape[1])\n', (1690, 1707), True, 'import numpy as np\n'), ((1712, 1729), 'numpy.random.seed', 'np.random.seed', (['(3)'], {}), '(3)\n', (1726, 1729), True, 'import numpy as np\n'), ((1734, 1757), 'numpy.random.shuffle', 'np.random.shuffle', (['rows'], {}), '(rows)\n', (1751, 1757), True, 'import numpy as np\n'), ((1762, 1785), 'numpy.random.shuffle', 'np.random.shuffle', (['cols'], {}), '(cols)\n', (1779, 1785), True, 'import numpy as np\n'), ((1929, 1943), 'numpy.min', 'np.min', (['matrix'], {}), '(matrix)\n', (1935, 1943), True, 'import numpy as np\n'), ((1958, 1972), 'numpy.max', 'np.max', (['matrix'], {}), '(matrix)\n', (1964, 1972), True, 'import numpy as np\n'), ((2140, 2167), 'matplotlib.pyplot.figure', 'plt.figure', ([], {'figsize': '(6, 12)'}), '(figsize=(6, 12))\n', (2150, 2167), True, 'import matplotlib.pyplot as plt\n'), ((2361, 2401), 'numpy.linspace', 'np.linspace', (['(mat_min + 6)', '(mat_max - 6)', '(5)'], {}), '(mat_min + 6, mat_max - 6, 5)\n', (2372, 2401), True, 'import numpy as np\n'), ((2470, 2496), 'numpy.searchsorted', 'np.searchsorted', (['bounds', '(0)'], {}), '(bounds, 0)\n', (2485, 2496), True, 'import numpy as np\n'), ((2511, 2536), 'numpy.insert', 'np.insert', (['bounds', 'idx', '(0)'], {}), '(bounds, idx, 0)\n', (2520, 2536), True, 'import numpy as np\n'), ((2548, 2576), 'matplotlib.colors.BoundaryNorm', 'BoundaryNorm', (['bounds', 'cmap.N'], {}), '(bounds, cmap.N)\n', (2560, 2576), False, 'from matplotlib.colors import BoundaryNorm\n'), ((2582, 2622), 'matplotlib.pyplot.imshow', 'plt.imshow', (['matrix'], {'cmap': 'cmap', 'norm': 'norm'}), '(matrix, cmap=cmap, norm=norm)\n', (2592, 2622), True, 'import matplotlib.pyplot as plt\n'), ((2712, 2721), 'matplotlib.pyplot.gca', 'plt.gca', ([], {}), '()\n', (2719, 2721), True, 'import matplotlib.pyplot as plt\n'), ((2902, 2920), 'matplotlib.pyplot.tight_layout', 'plt.tight_layout', ([], {}), '()\n', (2918, 2920), True, 'import matplotlib.pyplot as plt\n'), ((2931, 2940), 'matplotlib.pyplot.gcf', 'plt.gcf', ([], {}), '()\n', (2938, 2940), True, 'import matplotlib.pyplot as plt\n'), ((3037, 3047), 'matplotlib.pyplot.show', 'plt.show', ([], {}), '()\n', (3045, 3047), True, 'import matplotlib.pyplot as plt\n'), ((1333, 1348), 'matplotlib.pyplot.axis', 'plt.axis', (['"""off"""'], {}), "('off')\n", (1341, 1348), True, 'import matplotlib.pyplot as plt\n'), ((789, 808), 'matplotlib.pyplot.subplot', 'plt.subplot', (['gs1[j]'], {}), '(gs1[j])\n', (800, 808), True, 'import matplotlib.pyplot as plt\n'), ((995, 1010), 'matplotlib.pyplot.axis', 'plt.axis', (['"""off"""'], {}), "('off')\n", (1003, 1010), True, 'import matplotlib.pyplot as plt\n'), ((1380, 1397), 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['"""x_1"""'], {}), "('x_1')\n", (1390, 1397), True, 'import matplotlib.pyplot as plt\n'), ((1406, 1423), 'matplotlib.pyplot.ylabel', 'plt.ylabel', (['"""x_2"""'], {}), "('x_2')\n", (1416, 1423), True, 'import matplotlib.pyplot as plt\n'), ((1507, 1517), 'sys.exit', 'sys.exit', ([], {}), '()\n', (1515, 1517), False, 'import sys\n'), ((2045, 2060), 'numpy.abs', 'np.abs', (['mat_min'], {}), '(mat_min)\n', (2051, 2060), True, 'import numpy as np\n'), ((2104, 2118), 'numpy.abs', 'np.abs', (['matrix'], {}), '(matrix)\n', (2110, 2118), True, 'import numpy as np\n'), ((1268, 1292), 'numpy.array', 'np.array', (['[BLUE, ORANGE]'], {}), '([BLUE, ORANGE])\n', (1276, 1292), True, 'import numpy as np\n'), ((1996, 2011), 'numpy.abs', 'np.abs', (['mat_min'], {}), '(mat_min)\n', (2002, 2011), True, 'import numpy as np\n')]