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', ([], {'figsize': '(grid_y, grid_x)'}), True, 'import matplotlib.pyplot as plt\n'), (16, 'matplotlib.gridspec.GridSpec', 'gridspec.GridSpec', (['grid_y', 'grid_x'], {}), True, 'import matplotlib.gridspec as gridspec\n'), (21, 'matplotlib.pyplot.suptitle', 'plt.suptitle', (['title'], {}), 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', ([], {'figsize': '(8, 8)'}), True, 'import matplotlib.pyplot as plt\n'), (61, 'matplotlib.pyplot.title', 'plt.title', (['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', (['matrix.shape[0]'], {}), True, 'import numpy as np\n'), (69, 'numpy.arange', 'np.arange', (['matrix.shape[1]'], {}), True, 'import numpy as np\n'), (70, 'numpy.random.seed', 'np.random.seed', (['(3)'], {}), True, 'import numpy as np\n'), (71, 'numpy.random.shuffle', 'np.random.shuffle', (['rows'], {}), True, 'import numpy as np\n'), (72, 'numpy.random.shuffle', 'np.random.shuffle', (['cols'], {}), True, 'import numpy as np\n'), (78, 'numpy.min', 'np.min', (['matrix'], {}), True, 'import numpy as np\n'), (79, 'numpy.max', 'np.max', (['matrix'], {}), True, 'import numpy as np\n'), (84, 'matplotlib.pyplot.figure', 'plt.figure', ([], {'figsize': '(6, 12)'}), True, 'import matplotlib.pyplot as plt\n'), (89, 'numpy.linspace', 'np.linspace', (['(mat_min + 6)', '(mat_max - 6)', '(5)'], {}), True, 'import numpy as np\n'), (92, 'numpy.searchsorted', 'np.searchsorted', (['bounds', '(0)'], {}), True, 'import numpy as np\n'), (94, 'numpy.insert', 'np.insert', (['bounds', 'idx', '(0)'], {}), True, 'import numpy as np\n'), (95, 'matplotlib.colors.BoundaryNorm', 'BoundaryNorm', (['bounds', 'cmap.N'], {}), False, 'from matplotlib.colors import BoundaryNorm\n'), (97, 'matplotlib.pyplot.imshow', 'plt.imshow', (['matrix'], {'cmap': 'cmap', 'norm': 'norm'}), 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', (['"""off"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (28, 'matplotlib.pyplot.subplot', 'plt.subplot', (['gs1[j]'], {}), True, 'import matplotlib.pyplot as plt\n'), (36, 'matplotlib.pyplot.axis', 'plt.axis', (['"""off"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (55, 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['"""x_1"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (56, 'matplotlib.pyplot.ylabel', 'plt.ylabel', (['"""x_2"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (59, 'sys.exit', 'sys.exit', ([], {}), False, 'import sys\n'), (81, 'numpy.abs', 'np.abs', (['mat_min'], {}), True, 'import numpy as np\n'), (82, 'numpy.abs', 'np.abs', (['matrix'], {}), True, 'import numpy as np\n'), (50, 'numpy.array', 'np.array', (['[BLUE, ORANGE]'], {}), True, 'import numpy as np\n'), (80, 'numpy.abs', 'np.abs', (['mat_min'], {}), True, 'import numpy as np\n')]