import matplotlib
import utils

if utils.hostname() != 'user':
    matplotlib.use('Agg')

import matplotlib.pyplot as plt
import warnings
import numpy as np
import matplotlib.animation as animation

warnings.simplefilter('ignore')
anim_running = True


def plot_slice_3d_2(image3d, mask, axis, pid, img_dir=None, idx=None):
    fig, ax = plt.subplots(2, 2, figsize=[8, 8])
    fig.canvas.set_window_title(pid)
    masked_image = image3d * mask
    if idx is None:
        roi_idxs = np.where(mask == 1.)
        if len(roi_idxs[0]) > 0:
            idx = (np.mean(roi_idxs[0]), np.mean(roi_idxs[1]), np.mean(roi_idxs[2]))
        else:
            print 'No nodules'
            idx = np.array(image3d.shape) / 2
    else:
        idx = idx.astype(int)
    if axis == 0:  # sax
        ax[0, 0].imshow(image3d[idx[0], :, :], cmap=plt.cm.gray)
        ax[0, 1].imshow(mask[idx[0], :, :], cmap=plt.cm.gray)
        ax[1, 0].imshow(masked_image[idx[0], :, :], cmap=plt.cm.gray)
    if axis == 1:  # 2 lungs
        ax[0, 0].imshow(image3d[:, idx[1], :], cmap=plt.cm.gray)
        ax[0, 1].imshow(mask[:, idx[1], :], cmap=plt.cm.gray)
        ax[1, 0].imshow(masked_image[:, idx[1], :], cmap=plt.cm.gray)
    if axis == 2:  # side view
        ax[0, 0].imshow(image3d[:, :, idx[2]], cmap=plt.cm.gray)
        ax[0, 1].imshow(mask[:, :, idx[2]], cmap=plt.cm.gray)
        ax[1, 0].imshow(masked_image[:, :, idx[2]], cmap=plt.cm.gray)

    if img_dir is not None:
        fig.savefig(img_dir + '/%s%s.png' % (pid, axis), bbox_inches='tight')
    else:
        plt.show()
    fig.clf()
    plt.close('all')


def plot_slice_3d_3(input, mask, prediction, axis, pid, img_dir=None, idx=None):
    # to convert cuda arrays to numpy array
    input = np.asarray(input)
    mask = np.asarray(mask)
    prediction = np.asarray(prediction)

    fig, ax = plt.subplots(2, 2, figsize=[8, 8])
    fig.canvas.set_window_title(pid)
    if idx is None:
        roi_idxs = np.where(mask > 0)
        if len(roi_idxs[0]) > 0:
            idx = (int(np.mean(roi_idxs[0])),
                   int(np.mean(roi_idxs[1])),
                   int(np.mean(roi_idxs[2])))
        else:
            print 'No nodules'
            idx = np.array(input.shape) / 2
    else:
        idx = idx.astype(int)
    if axis == 0:  # sax
        ax[0, 0].imshow(prediction[idx[0], :, :], cmap=plt.cm.gray)
        ax[1, 0].imshow(input[idx[0], :, :], cmap=plt.cm.gray)
        ax[0, 1].imshow(mask[idx[0], :, :], cmap=plt.cm.gray)
    if axis == 1:  # 2 lungs
        ax[0, 0].imshow(prediction[:, idx[1], :], cmap=plt.cm.gray)
        ax[1, 0].imshow(input[:, idx[1], :], cmap=plt.cm.gray)
        ax[0, 1].imshow(mask[:, idx[1], :], cmap=plt.cm.gray)
    if axis == 2:  # side view
        ax[0, 0].imshow(prediction[:, :, idx[2]], cmap=plt.cm.gray)
        ax[1, 0].imshow(input[:, :, idx[2]], cmap=plt.cm.gray)
        ax[0, 1].imshow(mask[:, :, idx[2]], cmap=plt.cm.gray)
    if img_dir is not None:
        fig.savefig(img_dir + '/%s-%s.png' % (pid, axis), bbox_inches='tight')
    else:
        plt.show()
    fig.clf()
    plt.close('all')


def plot_slice_3d_3axis(input, pid, img_dir=None, idx=None):
    # to convert cuda arrays to numpy array
    input = np.asarray(input)

    fig, ax = plt.subplots(2, 2, figsize=[8, 8])
    fig.canvas.set_window_title(pid)
    ax[0, 0].imshow(input[idx[0], :, :], cmap=plt.cm.gray)
    ax[1, 0].imshow(input[:, idx[1], :], cmap=plt.cm.gray)
    ax[0, 1].imshow(input[:, :, idx[2]], cmap=plt.cm.gray)

    if img_dir is not None:
        fig.savefig(img_dir + '/%s.png' % (pid), bbox_inches='tight')
    else:
        plt.show()
    fig.clf()
    plt.close('all')


def plot_slice_3d_4(input, mask, prediction, lung_mask, axis, pid, img_dir=None, idx=None):
    # to convert cuda arrays to numpy array
    input = np.asarray(input)
    mask = np.asarray(mask)
    prediction = np.asarray(prediction)

    fig, ax = plt.subplots(2, 2, figsize=[8, 8])
    fig.canvas.set_window_title(pid)
    if idx is None:
        roi_idxs = np.where(mask > 0)
        if len(roi_idxs[0]) > 0:
            idx = (int(np.mean(roi_idxs[0])),
                   int(np.mean(roi_idxs[1])),
                   int(np.mean(roi_idxs[2])))
        else:
            print 'No nodules'
            idx = np.array(input.shape) / 2
    else:
        idx = idx.astype(int)
    if axis == 0:  # sax
        ax[0, 0].imshow(prediction[idx[0], :, :], cmap=plt.cm.gray)
        ax[1, 0].imshow(input[idx[0], :, :], cmap=plt.cm.gray)
        ax[0, 1].imshow(mask[idx[0], :, :], cmap=plt.cm.gray)
        ax[1, 1].imshow(lung_mask[idx[0], :, :], cmap=plt.cm.gray)
    if axis == 1:  # 2 lungs
        ax[0, 0].imshow(prediction[:, idx[1], :], cmap=plt.cm.gray)
        ax[1, 0].imshow(input[:, idx[1], :], cmap=plt.cm.gray)
        ax[0, 1].imshow(mask[:, idx[1], :], cmap=plt.cm.gray)
        ax[1, 1].imshow(lung_mask[:, idx[1], :], cmap=plt.cm.gray)
    if axis == 2:  # side view
        ax[0, 0].imshow(prediction[:, :, idx[2]], cmap=plt.cm.gray)
        ax[1, 0].imshow(input[:, :, idx[2]], cmap=plt.cm.gray)
        ax[0, 1].imshow(mask[:, :, idx[2]], cmap=plt.cm.gray)
        ax[1, 1].imshow(lung_mask[:, :, idx[2]], cmap=plt.cm.gray)
    if img_dir is not None:
        fig.savefig(img_dir + '/%s-%s.png' % (pid, axis), bbox_inches='tight')
    else:
        plt.show()
    fig.clf()
    plt.close('all')


def plot_slice_3d_3_patch(input, mask, prediction, axis, pid, patch_size=64, img_dir=None, idx=None):
    # to convert cuda arrays to numpy array
    input = np.asarray(input)
    mask = np.asarray(mask)
    prediction = np.asarray(prediction)

    fig, ax = plt.subplots(2, 2, figsize=[8, 8])
    fig.canvas.set_window_title(pid)
    if idx is None:
        roi_idxs = np.where(mask > 0)
        if len(roi_idxs[0]) > 0:
            idx = (np.mean(roi_idxs[0]), np.mean(roi_idxs[1]), np.mean(roi_idxs[2]))
        else:
            print 'No nodules'
            idx = np.array(input.shape) / 2
    if axis == 0:  # sax
        sz, sy, sx = slice(idx[0], idx[0] + 1), slice(idx[1] - patch_size, idx[1] + patch_size), slice(
            idx[2] - patch_size, idx[2] + patch_size)
        ax[0, 0].imshow(prediction[sz, sy, sx], cmap=plt.cm.gray)
        ax[1, 0].imshow(input[sz, sy, sx], cmap=plt.cm.gray)
        ax[0, 1].imshow(mask[sz, sy, sx], cmap=plt.cm.gray)
    if axis == 1:  # 2 lungs
        ax[0, 0].imshow(prediction[:, idx[1], :], cmap=plt.cm.gray)
        ax[1, 0].imshow(input[:, idx[1], :], cmap=plt.cm.gray)
        ax[0, 1].imshow(mask[:, idx[1], :], cmap=plt.cm.gray)
    if axis == 2:  # side view
        ax[0, 0].imshow(prediction[:, :, idx[2]], cmap=plt.cm.gray)
        ax[1, 0].imshow(input[:, :, idx[2]], cmap=plt.cm.gray)
        ax[0, 1].imshow(mask[:, :, idx[2]], cmap=plt.cm.gray)
    if img_dir is not None:
        fig.savefig(img_dir + '/%s.png' % pid, bbox_inches='tight')
    else:
        plt.show()
    fig.clf()
    plt.close('all')


def plot_slice_3d_2_patch(ct_scan, mask, pid, img_dir=None, idx=None):
    # to convert cuda arrays to numpy array
    ct_scan = np.asarray(ct_scan)
    mask = np.asarray(mask)

    fig, ax = plt.subplots(2, 3, figsize=[8, 8])
    fig.canvas.set_window_title(pid)

    if idx == None:
        #just plot in the middle of the cube
        in_sh = ct_scan.shape
        idx = [in_sh[0]/2,in_sh[1]/2,in_sh[2]/2]
    print np.amin(ct_scan), np.amax(ct_scan)
    print np.amin(mask), np.amax(mask)


    ax[0, 0].imshow(ct_scan[idx[0], :, :], cmap=plt.cm.gray)
    ax[0, 1].imshow(ct_scan[:, idx[1], :], cmap=plt.cm.gray)
    ax[0, 2].imshow(ct_scan[:, :, idx[2]], cmap=plt.cm.gray)

    ax[1, 0].imshow(mask[idx[0], :, :], cmap=plt.cm.gray)
    ax[1, 1].imshow(mask[:, idx[1], :], cmap=plt.cm.gray)
    ax[1, 2].imshow(mask[:, :, idx[2]], cmap=plt.cm.gray)

    if img_dir is not None:
        fig.savefig(img_dir + '/%s.png' % pid, bbox_inches='tight')
    else:
        plt.show()
    fig.clf()
    plt.close('all')


def plot_2d(img, mask, pid, img_dir):
    # fig = plt.figure()
    fig, ax = plt.subplots(2, 2, figsize=[8, 8])
    fig.canvas.set_window_title(pid)
    ax[0, 0].imshow(img, cmap='gray')
    ax[0, 1].imshow(mask, cmap='gray')
    ax[1, 0].imshow(img * mask, cmap='gray')
    plt.show()
    fig.savefig(img_dir + '/%s.png' % pid, bbox_inches='tight')
    fig.clf()
    plt.close('all')


def plot_2d_4(img, img_prev, img_next, mask, pid, img_dir):
    fig, ax = plt.subplots(2, 2, figsize=[8, 8])
    fig.canvas.set_window_title(pid)
    ax[0, 0].imshow(img, cmap='gray')
    ax[0, 1].imshow(img_prev, cmap='gray')
    ax[1, 0].imshow(img_next, cmap='gray')
    ax[1, 1].imshow(img * mask, cmap='gray')
    plt.show()
    fig.savefig(img_dir + '/%s.png' % pid, bbox_inches='tight')
    fig.clf()
    plt.close('all')


def plot_2d_animation(input, mask, predictions):
    rgb_image = np.concatenate((input, input, input), axis=0)
    mask = np.concatenate((np.zeros_like(input), mask, predictions), axis=0)
    # green = targets
    # blue = predictions
    # red = overlap

    idxs = np.where(mask > 0.3)
    rgb_image[idxs] = mask[idxs]

    rgb_image = np.rollaxis(rgb_image, axis=0, start=4)
    print rgb_image.shape

    def get_data_step(step):
        return rgb_image[step, :, :, :]

    fig = plt.figure()
    im = fig.gca().imshow(get_data_step(0))

    def init():
        im.set_data(get_data_step(0))
        return im,

    def animate(i):
        im.set_data(get_data_step(i))
        return im,

    anim = animation.FuncAnimation(fig, animate, init_func=init,
                                   frames=rgb_image.shape[1],
                                   interval=20000 / rgb_image.shape[0],
                                   blit=True)

    def on_click(event):
        global anim_running
        if anim_running:
            anim.event_source.stop()
            anim_running = False
        else:
            anim.event_source.start()
            anim_running = True

    fig.canvas.mpl_connect('button_press_event', on_click)
    try:
        plt.show()
    except AttributeError:
        pass

def plot_all_slices(input, pid, img_dir=None):
    # to convert cuda arrays to numpy array
    input = np.asarray(input)

    for idx in range(0, input.shape[0]-3, 4):
        fig, ax = plt.subplots(2, 2, figsize=[8, 8])
        fig.canvas.set_window_title(pid)
        ax[0, 0].imshow(input[idx, :, :], cmap=plt.cm.gray)
        ax[1, 0].imshow(input[idx+1, :, :], cmap=plt.cm.gray)
        ax[0, 1].imshow(input[idx+2, :, :], cmap=plt.cm.gray)
        ax[1, 1].imshow(input[idx+3, :, :], cmap=plt.cm.gray)

        if img_dir is not None:
            fig.savefig(img_dir + '_' + str(pid) + '_' + str(idx) + '.png' , bbox_inches='tight')
        else:
            plt.show()
        fig.clf()
        plt.close('all')


def plot_all_slices(ct_scan, mask, pid, img_dir=None):
    # to convert cuda arrays to numpy array
    ct_scan = np.asarray(ct_scan)
    mask = np.asarray(mask)

    for idx in range(0, mask.shape[0]-3, 2):
        fig, ax = plt.subplots(2, 2, figsize=[8, 8])
        fig.canvas.set_window_title(pid)
        ax[0, 0].imshow(mask[idx, :, :], cmap=plt.cm.gray)
        ax[1, 0].imshow(ct_scan[idx+1, :, :], cmap=plt.cm.gray)
        ax[0, 1].imshow(mask[idx+2, :, :], cmap=plt.cm.gray)
        ax[1, 1].imshow(ct_scan[idx+3, :, :], cmap=plt.cm.gray)

        if img_dir is not None:
            fig.savefig(img_dir + '_' + str(pid) + '_' + str(idx) + '.png' , bbox_inches='tight')
        else:
            plt.show()
        fig.clf()
        plt.close('all')

def plot_4_slices(input, pid, img_dir=None, idx=None):
    # to convert cuda arrays to numpy array
    input = np.asarray(input)

    fig, ax = plt.subplots(2, 2, figsize=[8, 8])
    fig.canvas.set_window_title(pid)
    ax[0, 0].imshow(input[idx[0], :, :], cmap=plt.cm.gray)
    ax[1, 0].imshow(input[:, idx[1], :], cmap=plt.cm.gray)
    ax[0, 1].imshow(input[:, :, idx[2]], cmap=plt.cm.gray)
    ax[1, 1].imshow(input[:, :, idx[2]], cmap=plt.cm.gray)

    if img_dir is not None:
        fig.savefig(img_dir + '/%s.png' % (pid), bbox_inches='tight')
    else:
        plt.show()
    fig.clf()
    plt.close('all')


def plot_learning_curves(train_losses, valid_losses, expid, img_dir):
    fig = plt.figure()
    x_range = np.arange(len(train_losses)) + 1

    plt.plot(x_range, train_losses)
    plt.plot(x_range, valid_losses)

    if img_dir is not None:
        fig.savefig(img_dir + '/%s.png' % expid, bbox_inches='tight')
        print 'Saved plot'
    else:
        plt.show()
    fig.clf()
    plt.close('all')