import seaborn as sns
import matplotlib as mpl
import matplotlib.pyplot as plt
from data_handlers.dictionaries import create_adversarial_dictionary
from data_handlers.dictionaries import create_gaussian_dictionary
sns.reset_orig()
w, h = 7.0, 4.0
mpl.rcParams['figure.figsize'] = [w, h]


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser('Plot dictionary')
    parser.add_argument('--save', type=str, default=None,
                        help='')
    parser.add_argument('--gaussian', action="store_true",
                        help='Create a figure with a guassian dictionary')
    parser.add_argument('--noshow', action="store_true",
                        help='Do not display the graph')
    parser.add_argument('--seaborn', action="store_true",
                        help='Load the seaborn style')
    args = parser.parse_args()

    if args.gaussian:
        K, p = 100, 64
        D = create_gaussian_dictionary(K, p, seed=290890)
    else:
        K, p = 256, 50
        D = create_adversarial_dictionary(K, p, seed=420829)
    cmap = sns.cubehelix_palette(8, start=.5, rot=-.75, dark=.2, reverse=True,
                                 as_cmap=True)
    fig = plt.figure("Dictionary")
    ax = fig.add_subplot(1, 1, 1)
    fig.subplots_adjust(bottom=.12, top=.99, right=.99)

    ax.imshow(D.T, aspect=K/p*h/w, cmap=cmap, interpolation="none")
    ax.set_xlabel("$K$", fontsize='xx-large')
    ax.set_ylabel("$p$", fontsize='xx-large')
    ax.set_xticks([])
    ax.set_yticks([])

    if args.seaborn:
        import seaborn
        seaborn.set_color_codes(palette='deep')
        seaborn.set_style("darkgrid", {
            "axes.facecolor": ".9",
            "figures.facecolor": (1, 1, 0, 0.5)})
        seaborn.despine(left=True, bottom=True)
    mpl.rcParams['figure.figsize'] = [7.0, 4.0]

    if args.save is not None:
        plt.savefig("../../communications/thesis/figures/{}.pdf".format(args.save), dpi=150,
                    transparent=True)
    if not args.noshow:
        plt.show()
