import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns; sns.set()


def visualize(s, batch, prefix):
    if len(s.shape) == 5:
        x, b, m = batch['x'], batch['b'], batch['m']
        im_visualize(s, x, b, m, prefix)
    elif len(s.shape) == 3:
        x, b, m = batch['x'], batch['b'], batch['m']
        pc_visualize(s, x, b, m, prefix)
    elif len(s.shape) == 4:
        xc, yc, xt, yt = batch['xc'], batch['yc'], batch['xt'], batch['yt']
        fn_visualize(s, xc, yc, xt, yt, prefix)
    else:
        raise ValueError()

def im_visualize(s, x, b, m, prefix):
    B,N,H,W,C = s.shape
    for i in range(B):
        ss, xx, bb, mm = s[i], x[i], b[i], m[i]
        if ss.shape[-1] == 2: # kspace
            C = 1
            ss = np.expand_dims(np.absolute(np.fft.ifft2(np.fft.ifftshift(ss[...,0] + ss[...,1] * 1j, axes=(-2,-1)))), axis=-1)
            ss = np.array(ss*255, dtype=np.uint8)
            xx = np.expand_dims(np.absolute(np.fft.ifft2(np.fft.ifftshift(xx[...,0] + xx[...,1] * 1j, axes=(-2,-1)))), axis=-1)
            xx = np.array(xx*255, dtype=np.uint8)
            bb = bb[...,0:1]
            mm = mm[...,0:1]
        ss = np.transpose(ss, [1,0,2,3]).reshape(H,W*N,C).squeeze()
        xx = np.transpose(xx, [1,0,2,3]).reshape(H,W*N,C).squeeze()
        bb = np.transpose(bb, [1,0,2,3]).reshape(H,W*N,C).squeeze()
        mm = np.transpose(mm, [1,0,2,3]).reshape(H,W*N,C).squeeze()
        xm = xx * mm + (1-mm) * 128
        xo = xx * bb + (1-bb) * 128
        img = np.concatenate([xm, xo, ss]).astype(np.uint8)

        plt.imsave(f'{prefix}_{i}.png', img)

def pc_visualize(s, x, b, m, prefix):
    B,N,C = s.shape
    for i in range(B):
        ss, xx, bb = s[i], x[i], b[i]
        o = np.where(bb[:,0]==1)[0]
        fig = plt.figure(figsize=(7.5, 2.5))
        ax = fig.add_subplot(131, projection='3d')
        ax.scatter(xx[:,0], xx[:,1], xx[:,2], c='g', s=5)
        ax.axis('off')
        ax.grid(False)
        ax = fig.add_subplot(132, projection='3d')
        ax.scatter(xx[o,0], xx[o,1], xx[o,2], c='g', s=5)
        ax.axis('off')
        ax.grid(False)
        ax = fig.add_subplot(133, projection='3d')
        ax.scatter(ss[:,0], ss[:,1], ss[:,2], c='g', s=5)
        ax.axis('off')
        ax.grid(False)
        plt.savefig(f'{prefix}_{i}.png')
        plt.close('all')

def fn_visualize(s, xc, yc, xt, yt, prefix):
    B,K,N,C = s.shape
    for i in range(B):
        ss, xxc, yyc, xxt, yyt = s[i], xc[i], yc[i], xt[i], yt[i]
        fig = plt.figure(figsize=(4.0, 2.5*K))
        for k in range(K):
            ax = fig.add_subplot(K,1,k+1)
            ax.plot(xxc[k], yyc[k], 'rx', markersize=8)
            ax.plot(xxt[k], yyt[k], 'ko', markersize=3)
            ax.plot(xxt[k], ss[k], 'bo', markersize=3)
        plt.savefig(f'{prefix}_{i}.png')
        plt.close('all')


def plot_functions(m, s, batch, prefix):
    B,K,N,C = m.shape
    xc, yc, xt, yt = batch['xc'], batch['yc'], batch['xt'], batch['yt']
    for i in range(B):
        mm, ss, xxc, yyc, xxt, yyt = m[i,:,:,0], s[i,:,:,0], xc[i,:,:,0], yc[i,:,:,0], xt[i,:,:,0], yt[i,:,:,0]
        fig = plt.figure(figsize=(4.0, 2.5*K))
        for k in range(K):
            idx = np.argsort(xxt[k])
            ax = fig.add_subplot(K,1,k+1)
            ax.plot(xxc[k], yyc[k], 'rx', markersize=8)
            ax.plot(xxt[k], yyt[k], 'ko', markersize=3)
            ax.plot(xxt[k,idx], mm[k,idx], 'b', linewidth=2)
            plt.fill_between(
                xxt[k,idx],
                mm[k,idx] - ss[k,idx],
                mm[k,idx] + ss[k,idx],
                alpha=0.2,
                facecolor='#65c9f7',
                interpolate=True)
        plt.savefig(f'{prefix}_{i}.png')
        plt.close('all')

def plot_img_functions(m, s, batch, prefix):
    B,K,N,C = m.shape
    idx, xc, yc, xt, yt = batch['idx'], batch['xc'], batch['yc'], batch['xt'], batch['yt']
    yo = np.ones_like(yt) * 128
    yo[:,:,idx] = (yc + 0.5) * 255.
    yt =  (yt + 0.5) * 255.
    m = (m + 0.5) * 255.
    for i in range(B):
        yoi, yti, mi = yo[i], yt[i], m[i]
        yoi = np.reshape(yoi, [K,28,28]).astype(np.uint8)
        yoi = np.reshape(np.transpose(yoi, [1,0,2]), [28, K*28])
        yti = np.reshape(yti, [K,28,28]).astype(np.uint8)
        yti = np.reshape(np.transpose(yti, [1,0,2]), [28, K*28])
        mi = np.reshape(mi, [K,28,28]).astype(np.uint8)
        mi = np.reshape(np.transpose(mi, [1,0,2]), [28, K*28])
        img = np.concatenate([yoi, mi, yti], axis=0)

        plt.imsave(f'{prefix}_{i}.png', img)