# coding=utf-8
'''
Created on 2016-9-27

@author: dengdan
'''
import matplotlib as mpl
# mpl.use('TkAgg')
import matplotlib.pyplot as plt
import numpy as np

from pylib import util

mpl.use('Agg')


def hist(x, title=None, normed=False, show=True, save=False,
         save_path=None, bin_count=100, bins=None,
         cumulative=False):
    x = np.asarray(x)
    if len(np.shape(x)) > 1:
        #         x = np.reshape(x, np.prod(x.shape))
        x = util.np.flatten(x)
    if bins is None:
        bins = np.linspace(start=min(x), stop=max(x), num=bin_count, endpoint=True, retstep=False)
    plt.figure(num=title)
    plt.hist(x, bins, normed=normed, cumulative=cumulative)
    if save:
        if save_path is None:
            raise ValueError
        path = util.io.join_path(save_path, title + '.png')
        save_image(path)
    if show:
        plt.show()
        # util.img.imshow(title, path, block = block)


def plot_solver_data(solver_path):
    data = util.io.load(solver_path)
    training_losses = data.training_losses
    training_accuracies = data.training_accuracies
    val_losses = data.val_losses
    val_accuracies = data.val_accuracies
    plt.figure(solver_path)

    n = len(training_losses)
    x = range(n)

    plt.plot(x, training_losses, 'r-', label='Training Loss')

    if len(training_accuracies) > 0:
        plt.plot(x, training_accuracies, 'r--', label='Training Accuracy')

    if len(val_losses) > 0:
        n = len(val_losses)
        x = range(n)
        plt.plot(x, val_losses, 'g-', label='Validation Loss')

        if len(val_accuracies) > 0:
            plt.plot(x, val_accuracies, 'g--', label='Validation Accuracy')
    plt.legend()
    plt.show()


def rectangle(xy, width, height, color='red', linewidth=1, fill=False, alpha=None, axis=None):
    """
    draw a rectangle on plt axis
    """
    import matplotlib.patches as patches
    rect_ = patches.Rectangle(
        xy=xy,
        width=width,
        height=height,
        alpha=alpha,
        color=color,
        fill=fill,
        linewidth=linewidth
    )
    if axis is not None:
        axis.add_patch(rect_)
    return rect_


rect = rectangle


def set_pos(x, y):
    mgr = plt.get_current_fig_manager()
    backend = mpl.get_backend()
    if backend == 'TkAgg':
        mgr.window.wm_geometry("+%d+%d" % (x, y))
    elif backend == 'WXAgg':
        mgr.window.SetPosition((x, y))
    else:
        # This works for QT and GTK
        # You can also use window.setGeometry
        mgr.window.move(x, y)


def maximize_figure():
    mng = plt.get_current_fig_manager()
    mng.full_screen_toggle()


def line(xy_start, xy_end, color='red', linewidth=1, alpha=None, axis=None):
    """
    draw a line on plt axis
    """
    from matplotlib.lines import Line2D
    num = 100
    xdata = np.linspace(xy_start[0], xy_end[0], num=num)
    ydata = np.linspace(xy_start[1], xy_end[1], num=num)
    line = Line2D(
        alpha=alpha,
        color=color,
        linewidth=linewidth,
        xdata=xdata,
        ydata=ydata
    )
    if axis is not None:
        axis.add_line(line)
    return line


def imshow(title=None, img=None, gray=False):
    show_images([img], [title], gray=gray)


def show_images(images, titles=None, shape=None, share_axis=False,
                bgr2rgb=False, maximized=False,
                show=True, gray=False, save=False, color_bar=False,
                path=None, axis_off=False, vertical=False, subtitle=None):
    plt.close('all')
    if shape is None:
        if vertical:
            shape = (len(images), 1)
        else:
            shape = (1, len(images))

    shape = list(shape)
    if shape[0] < 0:
        shape[0] = (len(images) + shape[1]) / shape[1]
    elif shape[1] < 0:
        shape[1] = (len(images + shape[0])) / shape[0]
    ret_axes = []
    ax0 = None
    plt.figure()
    for idx, img in enumerate(images):
        if bgr2rgb:
            img = util.img.bgr2rgb(img)
        loc = (idx / shape[1], idx % shape[1])
        if idx == 0:
            ax = plt.subplot2grid(shape, loc)
            ax0 = ax
        else:
            if share_axis:
                ax = plt.subplot2grid(shape, loc, sharex=ax0, sharey=ax0)
            else:
                ax = plt.subplot2grid(shape, loc)
        if len(np.shape(img)) == 2 and gray:
            img_ax = ax.imshow(img, cmap='gray')
        else:
            img_ax = ax.imshow(img)

        if len(np.shape(img)) == 2 and color_bar:
            plt.colorbar(img_ax, ax=ax)
        if titles is not None:
            ax.set_title(titles[idx])

        if axis_off:
            plt.axis('off')
        #             plt.xticks([]), plt.yticks([])
        ret_axes.append(ax)

    if subtitle is not None:
        set_subtitle(subtitle)
    if maximized:
        maximize_figure()

    if save:
        if path is None:
            path = util.get_temp_path()
        #             raise ValueError('path can not be None when save is True')
        save_image(path)
    if show:
        plt.show()
    return path


def save_image(path, img=None, dpi=150):
    path = util.io.get_absolute_path(path)
    util.io.make_parent_dir(path)
    if img is None:
        plt.gcf().savefig(path, dpi=dpi)
    else:
        plt.imsave(path, img, format='png')


imwrite = save_image


def to_ROI(ax, ROI):
    xy1, xy2 = ROI
    xmin, ymin = xy1
    xmax, ymax = xy2
    ax.set_xlim(xmin, xmax)
    # ax.extent
    ax.set_ylim(ymax, ymin)


def set_subtitle(title, font_size=12):
    plt.gcf().suptitle(title, fontsize=font_size)


def show(maximized=False):
    if maximized:
        maximize_figure()
    plt.show()


def draw():
    plt.gcf().canvas.draw()


def get_random_line_style():
    colors = ['r', 'g', 'b']
    line_types = ['-']  # , '--', '-.', ':']
    idx = util.rand.randint(len(colors))
    color = colors[idx]
    idx = util.rand.randint(len(line_types))
    line_type = line_types[idx]
    return color + line_type
