# -*- coding: utf-8 -*-
"""
Plotims

Images mutliple ims/h5 arrays and allows for the addition of
    scale bars, color bars etc.
Based on plotims_gui from IDL
"""

import numpy as np
from scipy import stats
import h5py
import matplotlib.pyplot as plt
from matplotlib import gridspec

class ims():
    def __init__(self):
        self.data = np.zeros((2,2))
        self.names = ""

class Resize():
    def __init__(self):
        self.xstart = 0
        self.ystart = 0
        self.xend = 0
        self.yend = 0

class Binning():
    def __init__(self):
        self.binx = 1
        self.biny = 1

class Rotate():
    def __init__(self):
        self.angle = 0
        self.fliph = False
        self.flipv = False

class Colorbar_opt():
    def __init__(self, discr=False, dir='vertical', fs_num=14, fs_title=14, title=''):
        if discr:
            self.discr = discr
        else: self.discr = False
        if dir:
            self.dir = dir
        else: self.dir = 'vertical'
        if fs_num:
            self.fs_num = fs_num
        else: self.fs_num = 14
        if fs_title:
            self.fs_title = fs_title
        else: self.fs_title = 14
        if title:
            self.title = title
        else: self.title = ''

class Plot_opts():
    def __init__(self, aspect=1, interpol='nearest', title_fontsize=20, clim=None, ct='OrRd', n2z=None, frame=True):
        if aspect:
            self.aspect = aspect
        else: self.aspect = 1
        if interpol:
            self.interpol = interpol
        else: self.interpol = 'nearest'
        if title_fontsize:
            self.title_fontsize = title_fontsize
        else: self.title_fontsize = 20
        self.clim = clim
        if ct:
            self.ct = ct
        else: self.ct = 'OrRd'
        if n2z:
            self.n2z = True
        else: self.n2z = False
        if frame:
            self.frame = frame
        else: self.frame = True

class Scale_opts():
    def __init__(self, xscale=False, x_pix_size=0, x_scl_size=0, x_scl_text='', yscale=False, y_pix_size=0, y_scl_size=0, y_scl_text='', fontsize=16):
        if xscale:
            self.xscale = xscale
        else: self.xscale = False
        if yscale:
            self.yscale = yscale
        else: self.yscale = False
        if x_pix_size:
            self.x_pix_size = x_pix_size
        else: self.x_pix_size = 0
        if x_scl_size:
            self.x_scl_size = x_scl_size
        else: self.x_scl_size = 0
        if y_pix_size:
            self.y_pix_size = y_pix_size
        else: self.y_pix_size = 0
        if y_scl_size:
            self.y_scl_size = y_scl_size
        else: self.y_scl_size = 0
        if x_scl_text:
            self.x_scl_text = x_scl_text
        else: self.x_scl_text = ''
        if y_scl_text:
            self.y_scl_text = y_scl_text
        else: self.y_scl_text = ''
        if fontsize:
            self.fontsize = fontsize
        else: self.fontsize = 16

class Collated_image_opts():
    def __init__(self, ncol=False, nrow=False, cb=False):
        if ncol:
            self.ncol = ncol
        else: self.ncol = 3
        if nrow:
            self.nrow = nrow
        else: self.nrow = 3
        if cb:
            self.cb = cb
        else: self.cb = False

def read_ims(imsfile):
    element_array = ""
    line = ""
    with open(imsfile, "r") as f:
        f.readline()
        dim = [int(i) for i in f.readline().split(" ") if i.strip()] #should contain 3 or 4 elements, depending on ims dimensions (2 or 3D)
        if(len(dim) != 3):
            print("TODO") #TODO
        imsf = ims()
        imsf.data = np.zeros(dim)
        for j in range(0, dim[2]):
            for k in range(0,dim[1]):
                dim0count = 0
                while dim0count < dim[0]:
                    line = [float(i) for i in f.readline().split(" ") if i.strip()]
                    imsf.data[dim0count:dim0count+len(line),k,j] = line
                    dim0count += len(line)
        element_array = f.readlines()
    for i in range(0,dim[2]):
        element_array[i] = element_array[i].strip()
    print("Succesfully read "+imsfile)
    imsf.names = element_array
    return imsf

def save_as_tif(h5file, h5channel, el2plot, savefile_prefix):
    import tifffile
    imsdata = read_h5(h5file, h5channel)
    if el2plot == 'All':
        el2plot = np.asarray(imsdata.names)
    else:
        el2plot = np.array(el2plot)

    if el2plot.size == 1:
        if el2plot not in imsdata.names:
            print('ERROR: save_as_tif: '+el2plot+' not in imsdata.names')
            return False
        idx = imsdata.names.index(el2plot)
        data = imsdata.data[:,:,idx]
        data[np.isnan(data)] = 0.
        tifffile.imwrite(savefile_prefix+'_'+''.join(imsdata.names[idx].split(" "))+'.tif', data.astype('float32'), photometric='minisblack')
        print('Saved '+savefile_prefix+'_'+''.join(imsdata.names[idx].split(" "))+'.tif')
    else:
        for i in range(0, el2plot.size):
            if el2plot[i] not in imsdata.names:
                print('ERROR: save_as_tif: '+el2plot+' not in imsdata.names')
                return False
            idx = imsdata.names.index(el2plot[i])
            data = imsdata.data[:,:,idx]
            data[np.isnan(data)] = 0.
            tifffile.imwrite(savefile_prefix+'_'+''.join(imsdata.names[idx].split(" "))+'.tif', data.astype('float32'), photometric='minisblack')
            print('Saved '+savefile_prefix+'_'+''.join(imsdata.names[idx].split(" "))+'.tif')
    return True

def read_h5(h5file, datadir):
    if type(datadir) is not type(list()):
        datadir = [datadir]
    with h5py.File(h5file, 'r') as file:
        data = []
        names = []
        for path in datadir:
            try:
                imsdat = np.array(file[path])
                if imsdat.ndim == 2:
                    imsdat = imsdat.reshape((1,imsdat.shape[0], imsdat.shape[1]))
                try:
                    nms = file['/'.join(path.split('/')[0:-1])+'/names']
                except Exception:
                    nms = [''.encode('utf8')]*imsdat.shape[0]
                for nm in nms:
                    names.append(nm)
                data.append(imsdat)
            except Exception:
                print("Error: unknown data directory: "+path+" in "+h5file)
                return None

    imsdat = np.concatenate(data, axis=0)

    # rearrange ims array to match what plotims expects
    imsdata = np.moveaxis(imsdat, 0, -1)
    imsdata[np.isnan(imsdata)] = 0.

    rv = ims()
    rv.data = np.asarray(imsdata)
    rv.names = [n.decode('utf8') for n in names[:]]
    return rv

def write_ims(imsdata, names, filename):
    with open(filename, "w") as f:
        f.write(str(len(imsdata.shape)-1)+'\n')
        line = list('')
        for i in range(0, len(imsdata.shape)):
            line.append(str(imsdata.shape[i]))
        f.write(" ".join(line)+'\n')
        for i in range(0, len(names)):
            for j in range(0, imsdata.shape[1]):
                f.write(str(" ".join(str(k) for k in imsdata[:,j,i]))+'\n')
        for i in range(0, len(names)):
            f.write(names[i]+'\n')

def prepare_rgb_data(data, r_eoi, g_eoi, b_eoi, rmin, rmax, gmin, gmax, bmin, bmax):
    rgb_im = np.zeros([data.shape[0], data.shape[1], 3])
    r_im = data[:, :, r_eoi]
    g_im = data[:, :, g_eoi]
    b_im = data[:, :, b_eoi]
    # set lowest and highest value to selected min and max
    r_im[r_im < rmin] = rmin
    r_im[r_im > rmax] = rmax
    g_im[g_im < gmin] = gmin
    g_im[g_im > gmax] = gmax
    b_im[b_im < bmin] = bmin
    b_im[b_im > bmax] = bmax
    # fill rgb array
    rgb_im[:, :, 0] = np.squeeze(r_im)
    rgb_im[:, :, 1] = np.squeeze(g_im)
    rgb_im[:, :, 2] = np.squeeze(b_im)
    # normalise values between 0 and 255 for each RGB channel
    for i in range(0,3):
        if rgb_im[:, :, i].max() != rgb_im[:, :, i].min():
            rgb_im[:, :, i] = 255 * (rgb_im[:, :, i]-rgb_im[:, :, i].min())/(rgb_im[:, :, i].max()-rgb_im[:, :, i].min())
        else:
            rgb_im[:, :, i] = 0
    return rgb_im.astype(np.uint8)

def make_rgb_triangle():
    width = 296
    height = 296
    triangle = np.zeros([width,height,4])
    green = np.zeros([width,height])
    red = np.zeros([width,height])
    blue = np.zeros([width,height])
    alpha = np.zeros([width,height])
    alpha[:,:] = 255
    # make gradients
    for i in range(0,width):
        for j in range(0,height):
            green[i,j] = j
            if (green[i,j] > 255):
                green[i,j] = 255
            red[i,j] = 0.5*(255 - np.sqrt(3)*(i-147) - j)
            if(red[i,j] > 255):
                red[i,j] = 255
            if(red[i,j] < 0):
                red[i,j] = 0
            blue[i,j] = 0.5*(255 + np.sqrt(3)*(i-147) - j)
            if(blue[i,j] > 255):
                blue[i,j] = 255
            if(blue[i,j] < 0):
                blue[i,j] = 0
    # cut away corners to a triangle
    for i in range(0,width):
        for j in range(0,height):
            if(blue[i,j] == 0):
                red[i,j] = 0
            if(red[i,j] == 0):
                blue[i,j] = 0
            if(red[i,j] == 0 and blue[i,j] == 0):
                green[i,j] = 0
                alpha[i,j] = 0 # set background transparent
    triangle[:,:,0] = red
    triangle[:,:,1] = green
    triangle[:,:,2] = blue
    triangle[:,:,3] = alpha
    # scale max intensity to 255
    triangle2 = triangle
    for i in range(0, width):
        for j in range(0, height):
            if(triangle[i,j,0:3].max() > 0):
                triangle[i,j,0:3] = 255*triangle2[i,j,0:3]/triangle[i,j,0:3].max()
    rotate_triangle = np.zeros([height,width,4])
    for i in range(0, width):
        for j in range(0, height):
            rotate_triangle[height-j-1,i,:] = triangle[i,j,:]
    return rotate_triangle

def plot_rgb(imsdata, eoi_red, eoi_green, eoi_blue, filename, rmin=None, rmax=None, gmin=None, gmax=None, bmin=None, bmax=None, pix_size=None, scl_size=None, scl_unit=None):
    # we expect type of eoi_red/green/blue to be integer indices. If they are string, let's look in imsdata.names for their indices
    if type(eoi_red) is str:
        eoi_red = imsdata.names.index(eoi_red)
    if type(eoi_green) is str:
        eoi_green = imsdata.names.index(eoi_green)
    if type(eoi_blue) is str:
        eoi_blue = imsdata.names.index(eoi_blue)

    if rmin is None:
        rmin = np.min(imsdata.data[:,:,eoi_red])
    if rmax is None:
        rmax = np.max(imsdata.data[:,:,eoi_red])
    if gmin is None:
        gmin = np.min(imsdata.data[:,:,eoi_green])
    if gmax is None:
        gmax = np.max(imsdata.data[:,:,eoi_green])
    if bmin is None:
        bmin = np.min(imsdata.data[:,:,eoi_blue])
    if bmax is None:
        bmax = np.max(imsdata.data[:,:,eoi_blue])

    plt_opts = Plot_opts(aspect=1.)

    # prepare rgb array to plot (each colour channel must be scaled between 0 and 255)
    rgb_im = prepare_rgb_data(imsdata.data.copy(), eoi_red, eoi_green, eoi_blue, rmin, rmax, gmin, gmax, bmin, bmax)
    # matplotlib plotting, to save image
    fig = plt.figure(figsize=(10,10))
    gs = gridspec.GridSpec(2,1, height_ratios=[0.2,1])
    rgb0 = plt.subplot(gs[0], anchor='SW')
    rgb_triangle = make_rgb_triangle()
    rgb0.imshow(rgb_triangle.astype(np.uint8))
    rgb0.axis('off')
    # add labels on rgb triangle
    green_lbl = imsdata.names[eoi_green].split("-")
    green_lbl = green_lbl[0]
    red_lbl = imsdata.names[eoi_red].split("-")
    red_lbl = red_lbl[0]
    blue_lbl = imsdata.names[eoi_blue].split("-")
    blue_lbl = blue_lbl[0]
    rgb0.text(0.5, 0.56, green_lbl, size=16, ha="center", transform=rgb0.transAxes) #green label
    rgb0.text(0.25, 0.1, red_lbl, size=16, color='w', ha="center", transform=rgb0.transAxes) #red label
    rgb0.text(0.75, 0.1, blue_lbl, size=16, color='w', ha="center", transform=rgb0.transAxes) #blue label

    # add text with cutoff values
    red_cutoff = imsdata.names[eoi_red]+" cutoff value min: "+str(rmin)+" max:"+str(rmax)
    green_cutoff = imsdata.names[eoi_green]+" cutoff value min: "+str(gmin)+" max:"+str(gmax)
    blue_cutoff = imsdata.names[eoi_blue]+" cutoff value min: "+str(bmin)+" max:"+str(bmax)
    rgb0.text(1.2, 0.75, red_cutoff, size=12, ha="left", transform=rgb0.transAxes)
    rgb0.text(1.2, 0.625, green_cutoff, size=12, ha="left", transform=rgb0.transAxes)
    rgb0.text(1.2, 0.5, blue_cutoff, size=12, ha="left", transform=rgb0.transAxes)

    # plot actual rgb image
    rgb1 = plt.subplot(gs[1])
    rgb1.imshow(rgb_im.astype(np.uint8), aspect=1./plt_opts.aspect, interpolation=plt_opts.interpol)
    rgb1.axis('off')
    plt.tight_layout()
    # add scale_bar if requested
    if pix_size is not None and scl_size is not None:
        add_scalebar(rgb1, pix_size, scl_size, str(scl_size)+' '+scl_unit, dir='h')
    fig.savefig(filename, dpi=420)
    plt.close()

def ims_data_manip(imsdata, resize=None, binning=None, neg2zero=None, mathop=None, rotate=None):
    # resize
    if resize:
        imsdata = imsdata[resize.ystart:resize.yend, resize.xstart:resize.xend, :]
    # complete rest of binning
    if binning:
        binx = binning.binx
        biny = binning.biny
        newims = np.zeros( (int(np.floor(imsdata.shape[0]/biny)), int(np.floor(imsdata.shape[1]/binx)), imsdata.shape[2]) )
        xshift = int(round((imsdata.shape[1]/binx-newims.shape[1])*binx/2))
        yshift = int(round((imsdata.shape[0]/biny-newims.shape[0])*biny/2))
        for i in range(0, newims.shape[0]):
            for j in range(0, newims.shape[1]):
                for k in range(0, newims.shape[2]):
                    newims[i, j, k] = np.sum(imsdata[i*biny+yshift:(i+1)*biny+yshift, j*binx+xshift:(j+1)*binx+xshift, k])
        imsdata = newims
    # neg2zero
    if neg2zero:
        imsdata[imsdata < 0] = 0
    # sqrt or log
    if mathop:
        if(mathop == 'sqrt'):
            imsdata[imsdata < 0] = 0
            imsdata = np.sqrt(imsdata)
        elif(mathop == 'log'):
            imsdata[imsdata < 1] = 1
            imsdata = np.log10(imsdata)
    # rotate
    if rotate:
        # first rotate over chosen angle
        imsdata = np.rot90(imsdata, int(rotate.angle/90), (0,1))
        # then flip/invert axis
        if(rotate.fliph == True):
            imsdata = np.flip(imsdata, 1)
        if(rotate.flipv == True):
            imsdata = np.flip(imsdata, 0)
    return imsdata

def plot_correl(imsdata, imsnames, el_id=None, save=None):
    """
    Display correlation plots.

    Parameters
    ----------
    imsdata : float array
        imsdata is a N*M*Y float array containing the signal intensities of N*M datapoints for Y elements.
    imsnames : string
        imsnames is a string array of Y elements, containing the names of the corresponding elements.
    el_id : integer list, optional
        el_id should be a integer list containing the indices of the elements to include in the plot. The default is None.
    save : string, optional
        File path as which the image should be saved. The default is None.

    Returns
    -------
    None.

    """
    imsdata = np.array(imsdata)
    imsnames = np.array(imsnames)
    data =  imsdata.reshape(imsdata.shape[0]*imsdata.shape[1], imsdata.shape[2])
    if el_id is not None:
        data = data[:, el_id]
        imsnames = imsnames[el_id]

    plt.figure(figsize=(10,10))
    gs = gridspec.GridSpec(imsnames.shape[0], imsnames.shape[0])
    num_bins = 5
    scatterplot = []
    for i in range(imsnames.shape[0]):
        for j in range(imsnames.shape[0]):
            scatterplot.append(plt.subplot(gs[i,j]))
            if j>i:
                scatterplot[-1].plot(data[:,j], data[:,i], linewidth=0, marker='.', alpha=0.5)  # correlation plots
                scatterplot[-1].margins(0.05)
                # fit and plot regression line
                z = np.polyfit(data[:,j], data[:,i], 1)
                p = np.poly1d(z)
                y_model = p(data[:,j])
                n = data[:,i].size                                           # number of observations
                m = z.size                                                 # number of parameters
                dof = n - m                                                # degrees of freedom
                t = stats.t.ppf(0.975, n - m)                              # used for CI and PI bands 0.975=95%CI
                resid = data[:,i] - y_model
                # chi2 = np.sum((resid / y_model)**2)                        # chi-squared; estimates error in data
                # chi2_red = chi2 / dof                                      # reduced chi-squared; measures goodness of fit
                s_err = np.sqrt(np.sum(resid**2) / dof)                    # standard deviation of the error
                scatterplot[-1].plot(data[:,j], y_model, "--", color="black", linewidth=1.0, alpha=0.5, label="Fit")
                x2 = np.linspace(np.min(data[:,j]), np.max(data[:,j]), 100)
                y2 = p(x2)
                ci = t * s_err * np.sqrt(1/n + (x2 - np.mean(data[:,j]))**2 / np.sum((data[:,j] - np.mean(data[:,j]))**2))
                scatterplot[-1].fill_between(x2, y2 + ci, y2 - ci, color="gray", alpha=0.8)
                # calculate R² value
                rsq = np.sum((y_model - np.mean(data[:,i]))**2) / np.sum((data[:,i] - np.mean(data[:,i]))**2)
                scatterplot[-1].annotate('R² = {:.3f} '.format(rsq), xy=(0.05, 0.85), xycoords='axes fraction', fontsize=8, color='gray')

            elif j==i:
                ymin = np.min(data[:,j]) - (np.max(data[:,j])-np.min(data[:,j]))*0.05
                ymax = np.max(data[:,j]) + (np.max(data[:,j])-np.min(data[:,j]))*0.05
                hist, edges = np.histogram(data[:,j], bins=num_bins*2)
                hist = np.min(data[:,j]) + ((hist-np.min(hist))/(np.max(hist)-np.min(hist))) * (np.max(data[:,j])-np.min(data[:,j]))
                xhist = np.zeros(hist.shape[0]*2)
                yhist = np.zeros(hist.shape[0]*2)
                for k in range(hist.shape[0]):
                    if k == 0:
                        xhist[k] = np.min(data[:,j])
                        xhist[-1] = np.max(data[:,j])
                    else:
                        xhist[2*k-1:2*k+1] = edges[k]
                    yhist[2*k:2*k+2] = hist[k]
                scatterplot[-1].fill_between(xhist, yhist, np.zeros(hist.shape[0]*2)+ymin, color='tab:Blue', alpha=0.5)
                scatterplot[-1].set_ylim(ymin, ymax)
                # scatterplot[-1].hist(data[:,j], bins=num_bins*2, density=True, alpha=0.5) # density plots
                # scatterplot[-1].margins(0.05)
            else:
                # kernel density estimate plots
                kernel = stats.gaussian_kde(np.vstack([data[:,j], data[:,i]]))
                xmin = np.min(data[:,j]) - (np.max(data[:,j])-np.min(data[:,j]))*0.05
                xmax = np.max(data[:,j]) + (np.max(data[:,j])-np.min(data[:,j]))*0.05
                ymin = np.min(data[:,i]) - (np.max(data[:,i])-np.min(data[:,i]))*0.05
                ymax = np.max(data[:,i]) + (np.max(data[:,i])-np.min(data[:,i]))*0.05
                X, Y = np.mgrid[xmin:xmax:50j, ymin:ymax:50j]
                positions = np.vstack([X.ravel(), Y.ravel()])
                Z = np.reshape(kernel(positions).T, X.shape)
                scatterplot[-1].contour(X, Y, np.sqrt(Z), num_bins, cmap='Blues')
                # calculate Pearson correlation and give idea of confidence interval
                r, p = stats.pearsonr(data[:,j], data[:,i])
                p_stars = ''
                if p <= 0.05:
                    p_stars = '*'
                if p <= 0.01:
                    p_stars = '**'
                if p <= 0.001:
                    p_stars = '***'
                scatterplot[-1].annotate('r = {:.2f} '.format(r) + p_stars, xy=(0.05, 0.85), xycoords='axes fraction', fontsize=8, color='gray')

            if i == imsnames.shape[0]-1:
                scatterplot[-1].set_xlabel(imsnames[j], fontsize=12)
                scatterplot[-1].xaxis.set_major_locator(plt.MaxNLocator(3))
            if i != imsnames.shape[0]-1:
                scatterplot[-1].set_xticklabels('')
                scatterplot[-1].set_xticks([])

            if j == 0:
                scatterplot[-1].set_ylabel(imsnames[i], fontsize=12)
                scatterplot[-1].yaxis.set_major_locator(plt.MaxNLocator(3))
            if j != 0:
                scatterplot[-1].set_yticklabels('')
                scatterplot[-1].set_yticks([])

    plt.tight_layout()
    plt.show()
    if save is not None:
        plt.savefig(save, bbox_inches='tight', pad_inches=0, dpi=420)
        plt.close()

def add_scalebar(target, pix_size, scl_size, scale_text,scale_fontsize=16, dir=''):
    if(dir == 'horizontal' or dir == 'h'):
        target.plot([0,scl_size/pix_size], [-2,-2], lw=2, color='black', clip_on=False)
        target.text(scl_size/(2.*pix_size), -3, scale_text, ha='center', va='bottom', size=scale_fontsize, clip_on=False)
    if(dir == 'vertical' or dir == 'v' ):
        target.plot([-2,-2], [0,scl_size/pix_size], lw=2, color='black', clip_on=False)
        target.text(-3, scl_size/(2.*pix_size), scale_text, ha='right', va='center', size=scale_fontsize, rotation=90, clip_on=False)

def plot_image(imsdata, imsname, ctable, plt_opts=None, sb_opts=None, cb_opts=None, clim=None, save=None, subplot=None):
    # set option for discrete colorbar, only if 10 or less values are plotted
    if(cb_opts and cb_opts.discr and imsdata.max()-imsdata.min() <= 10):
        ctable = plt.cm.get_cmap(ctable, np.around(imsdata.max()-imsdata.min()+1).astype(int))
        clim = [imsdata.min()-0.5, imsdata.max()+0.5]
    if plt_opts:
        aspect = plt_opts.aspect
        interpol = plt_opts.interpol
        fs_im_tit = plt_opts.title_fontsize
        frame = plt_opts.frame
    else:
        aspect = 1
        interpol = 'nearest'
        fs_im_tit = 20
        frame = True
    # if this plot is part of a subplot, provide subplot axes. Otherwise extract plt axes
    if type(subplot) != type(None):
        ncols = subplot[0]
        nrows = subplot[1]
        n_el = imsdata.shape[2]
    else:
        ncols, nrows, n_el = 1, 1, 1

    dx, dy = imsdata.shape[1], imsdata.shape[0]
    img_high = 6 # height in inches of all image rows stacked, omitting whitespaces
    pad = 0.05 # padding between cb or title and image in inches

    ws_ver, ws_hor = 0.35, 0.01
    if cb_opts:
        if cb_opts.dir == 'vertical': #---vertical colorbar---#
            ws_ver = 0.35 # vertical whitespace between images [inch]
            ws_hor = 0.6 # horizontal whitespace between images [inch]
        elif cb_opts.dir == 'horizontal': #---horizontal colorbar---#
            ws_ver = 0.6 # vertical whitespace between images [inch]
            if cb_opts.title == '':
                ws_hor = 0.01 # horizontal whitespace between images [inch]
            elif cb_opts.title == "Int.;[$\sqrt{cts}$]" or cb_opts.title == "Int.;[cts]" or cb_opts.title == "Int.;[log(cts)]":
                ws_hor = 0.3
            else:
                ws_hor = 0.55
        cb_opts.title = '\n'.join(cb_opts.title.split(';'))

    height = img_high + (nrows+1)*(ws_ver+pad)  # width and height are in inches
    width = (ws_hor+pad)*(ncols+1) + ncols*(img_high/nrows)*(dx/dy)*aspect # each image has width dependent on aspect ratio and dx/dy ratio
    fig = plt.figure(figsize=(width, height))

    # now we know each image is ((img_high/nrows)*(dx/dy)*aspect, (img_high/nrows)) inches wide,high. Convert this to relative ratio of image and whitespace:
    im_rel_wide, im_rel_high = ((img_high/nrows)*(dx/dy)*aspect)/width, (img_high/nrows)/height
    ws_rel_wide, ws_rel_high = ws_hor/width, ws_ver/height
    col_id, row_id = 0, 0
    for j in range(n_el):
        if n_el == 1:
            data = imsdata[:,:]
            name = imsname
        else:
            data = imsdata[:,:,j]
            name = imsname[j]
        if type(clim) != type(None):
            # clim = clim
            if(clim[0] > np.min(data) and clim[1] < np.max(data)):
                extend = 'both'
            elif clim[0] > np.min(data):
                extend = 'min'
            elif clim[1] < np.max(data):
                extend = 'max'
            else:
                extend = 'neither'
            limit = clim
        else:
            limit = None
            extend = 'neither'
        im_xstart = ws_rel_wide+(im_rel_wide+ws_rel_wide+pad/width)*col_id
        im_ystart = 1.-(im_rel_high+ws_rel_high)*(row_id+1)
        ax_im = fig.add_axes([im_xstart, im_ystart, im_rel_wide, im_rel_high])
        img = ax_im.imshow(data, interpolation=interpol, cmap=ctable, aspect='auto', clim=limit)
        plt.text(dx/2, -1*(pad/(im_rel_high*height))*dy, name, ha='center', va='bottom', size=fs_im_tit, clip_on=False)
        if frame: # frame image or not
            ax_im.set(xticks=[], yticks=[])
        else:
            ax_im.axis('off')
        # add scale_bar if requested
        if(sb_opts and sb_opts.xscale):
            sb_ax = fig.add_axes([im_xstart, im_ystart+im_rel_high+pad/height, im_rel_wide, 0], sharex=ax_im)
            sb_ax.axis('off')
            add_scalebar(sb_ax, sb_opts.x_pix_size, sb_opts.x_scl_size, sb_opts.x_scl_text, scale_fontsize=sb_opts.fontsize, dir='h')
        if(sb_opts and sb_opts.yscale):
            sb_ax = fig.add_axes([im_xstart-pad/width, im_ystart, 0, im_rel_high], sharey=ax_im)
            sb_ax.axis('off')
            add_scalebar(sb_ax, sb_opts.y_pix_size, sb_opts.y_scl_size, sb_opts.y_scl_text, scale_fontsize=sb_opts.fontsize, dir='v')
        # draw colorbar on its own axis, within image dimensions
        if cb_opts:
            if cb_opts.dir == 'vertical':
                ax_cb = fig.add_axes([im_xstart+im_rel_wide+pad/width, im_ystart, 0.1/width, im_rel_high])
                ax_cb.text(2, 1+(pad/(im_rel_high*height)), cb_opts.title, ha='center', va='bottom', size=cb_opts.fs_title, clip_on=False, transform=ax_cb.transAxes)
                cbar = fig.colorbar(img, orientation=cb_opts.dir, extend=extend, cax=ax_cb)
            if cb_opts.dir == 'horizontal':
                ax_cb = fig.add_axes([im_xstart, im_ystart-(pad+0.1)/height, im_rel_wide, 0.1/height])
                ax_cb.text(1+(pad+ws_hor)/(2*im_rel_wide*width), 0, cb_opts.title, ha='center', va='center', size=cb_opts.fs_title, clip_on=False, transform=ax_cb.transAxes)
                cbar = fig.colorbar(img, orientation=cb_opts.dir, extend=extend, cax=ax_cb)
            cbar.ax.tick_params(labelsize=cb_opts.fs_num)

        col_id = col_id+1
        if col_id >= ncols:
            col_id = 0
            row_id = row_id+1
    # save the image (and close it)
    if save:
        plt.savefig(save, bbox_inches='tight', pad_inches=0, dpi=420)
        plt.close()

def plot_colim(imsdata, el_selection, colortable, plt_opts=None, sb_opts=None, cb_opts=None, colim_opts=None, save=None):
    if colim_opts:
        ncols = colim_opts.ncol
        nrows = colim_opts.nrow
        if colim_opts.cb == False:
            cb_opts = None # dont't draw colorbar if chosen (default is to draw)
    else:
        ncols = int(np.floor(np.sqrt(len(el_selection))))
        nrows = int(np.ceil(len(el_selection)/ncols))
    if plt_opts:
        if type(plt_opts.clim) != type(None):
            clim = plt_opts.clim
        else: clim = None
        plt_opts.title_fontsize = plt_opts.title_fontsize/2
    else:
        clim = None
    if sb_opts:
        sb_opts.fontsize = sb_opts.fontsize/2 #divide fontsize by 2 for collated imaging
    if cb_opts:
        cb_opts.fs_title = cb_opts.fs_title/2
        cb_opts.fs_num = cb_opts.fs_num/2

    # Make data array containing only the elements to plot.
    #   First let's check how many elements are common between el_selection and imsdata.names
    cnt = 0
    for i in range(len(el_selection)):
        if el_selection[i] in np.arange(len(imsdata.names)):
            for k in range(len(imsdata.names)):
                if k == el_selection[i]:
                    cnt = cnt+1
    #    Now make empty array, and repeat loop above to fill array appropiately
    datacube = np.zeros((imsdata.data.shape[0], imsdata.data.shape[1], cnt))
    names = list('')
    if clim:
        cb_lim = np.zeros((cnt,2))
    else:
        cb_lim = None
    cnt = 0
    for i in range(len(el_selection)):
        if el_selection[i] in np.arange(len(imsdata.names)):
            for k in range(len(imsdata.names)):
                if k == el_selection[i]:
                    datacube[:,:,cnt] = imsdata.data[:,:,k]
                    names.append(imsdata.names[k])
                    if clim:
                        cb_lim[cnt,0] = np.min(imsdata.data[:,:,k])+clim[0]*(np.max(imsdata.data[:,:,k])-np.min(imsdata.data[:,:,k]))
                        cb_lim[cnt,1] = np.min(imsdata.data[:,:,k])+clim[1]*(np.max(imsdata.data[:,:,k])-np.min(imsdata.data[:,:,k]))
                    cnt = cnt+1

    # Perform plotting. plot_image will know how to handle collated images based on amount of rows and columns.
    plot_image(datacube, names, colortable, plt_opts=plt_opts, sb_opts=sb_opts, cb_opts=cb_opts, clim=cb_lim, save=save, subplot=(ncols, nrows))
    # set fonts back to normal value for further imaging; probably useless as sb_opts will be local variable, but meh
    if sb_opts:
        sb_opts.fontsize = sb_opts.fontsize*2
    if plt_opts:
        plt_opts.title_fontsize = plt_opts.title_fontsize*2
    if cb_opts:
        cb_opts.fs_title = cb_opts.fs_title*2
        cb_opts.fs_num = cb_opts.fs_num*2


"""
Utility functions for data loading, downloading, display
"""
__author__ = "Marcel Kurovski"
__copyright__ = "Marcel Kurovski"
__license__ = "mit"

import json
import logging
import os
import requests
import shutil
import sys
import time
import zipfile

import matplotlib.pyplot as plt

from graphpipe import remote
import numpy as np
from scipy.io import loadmat

EMNIST_MATLAB_URL = 'http://www.itl.nist.gov/iaui/vip/cs_links/EMNIST/matlab.zip'
EMNIST_FILENAME = 'emnist-byclass.mat'

_logger = logging.getLogger(__name__)

def setup_logging(loglevel):
    """Setup basic logging

    Args:
      loglevel (int): minimum loglevel for emitting messages
    """
    logformat = "[%(asctime)s] %(levelname)s:%(name)s:%(message)s"
    logging.basicConfig(level=loglevel, stream=sys.stdout,
                        format=logformat, datefmt="%Y-%m-%d %H:%M:%S")

setup_logging(logging.INFO)

def load_emnist(emnist_folder_path, download=False):
    """
    Loads EMNIST data from folder (eventually downloads first)
    using the Matlab format
    See https://www.nist.gov/itl/iad/image-group/emnist-dataset

    Args:
        emnist_folder_path (str): folder that contains EMNIST Matlab files
        download (bool): in case `emnist_folder_path` does not exist or contain
                         the Matlab files, download the data and extract them

    Returns:
        x_train (:obj:`np.array`): (n_train, 28, 28) training images
        y_train (:obj:`np.array`): (n_train, 1) training labels
        x_test (:obj:`np.array`): (n_test, 28, 28) training images
        y_test (:obj:`np.array`y): (n_test, 1) training labels
        mapping (:obj:`np.array`): (62, 2) maps training labels (first column)
            to ascii codes (second column)
    """
    filepath = os.path.join(emnist_folder_path, EMNIST_FILENAME)

    if not (os.path.isdir(emnist_folder_path) and os.path.isfile(filepath)):
        if download:
            download_emnist(emnist_folder_path)
        else:
            error_msg = ("Folder {} or file {} does not exist "
                         "and download is deactivated").format(
                                emnist_folder_path, EMNIST_FILENAME)
            raise FileNotFoundError(error_msg)

    raw_emnist_data = loadmat(filepath)

    _logger.info("Loading train and test data from %s", filepath)

    # load train data
    x_train = raw_emnist_data["dataset"][0][0][0][0][0][0].astype(np.float32)
    x_train = x_train.reshape((x_train.shape[0], 28, 28), order='A')
    y_train = raw_emnist_data["dataset"][0][0][0][0][0][1]

    # load test data
    x_test = raw_emnist_data["dataset"][0][0][1][0][0][0].astype(np.float32)
    x_test = x_test.reshape((x_test.shape[0], 28, 28), order='A')
    y_test = raw_emnist_data["dataset"][0][0][1][0][0][1]

    # load mapping from label to chr
    mapping = raw_emnist_data["dataset"][0][0][2]

    return x_train, y_train, x_test, y_test, mapping

def download_emnist(emnist_folder_path):
    """
    Downloads, extracts and moves files to desired path

    Args:
        emnist_folder_path (str): folder for EMNIST data download and extraction
    """
    _logger.info("Data not found. Starting EMNIST Download from %s",
                 EMNIST_MATLAB_URL)
    if not os.path.isdir(emnist_folder_path):
        os.mkdir(emnist_folder_path)
        _logger.info("Target folder not found. Created %s", emnist_folder_path)

    emnist_matlab_zip_filepath = os.path.join(emnist_folder_path,
                                              'emnist_matlab.zip')
    req = requests.get(EMNIST_MATLAB_URL, stream=True)

    with open(emnist_matlab_zip_filepath, 'wb') as fp:
        for chunk in req.iter_content(chunk_size=1024):
            if chunk:
                fp.write(chunk)

    _logger.info("Download successful. Extracting and moving files...")
    with zipfile.ZipFile(emnist_matlab_zip_filepath, 'r') as zip_ref:
        zip_ref.extractall(emnist_folder_path)

    matlab_folder_path = os.path.join(emnist_folder_path, 'matlab')
    for matlab_file in os.listdir(matlab_folder_path):
        shutil.move(os.path.join(matlab_folder_path, matlab_file),
                    os.path.join(emnist_folder_path, matlab_file))

    shutil.rmtree(matlab_folder_path)

def get_emnist_mapping():
    """
    Create Mapping for 62 classes from ASCII code to digit or character

    Returns:
        mapping (dict): maps ASCII integer to digit / character string
    """
    ascii_digits = list(range(48, 58))
    ascii_uppercase_letters = list(range(65, 65+26))
    ascii_lowercase_letters = list(range(97, 97+26))

    digits = [chr(value) for value in ascii_digits]
    uppercase_letters = [chr(value) for value in ascii_uppercase_letters]
    lowercase_letters = [chr(value) for value in ascii_lowercase_letters]

    mapping = dict(zip(range(10+2*26),
                       digits+uppercase_letters+lowercase_letters))

    return mapping

def show_img(idx, x_train, y_train, x_test, y_test, mapping, mode='train'):
    """
    Shows the selected EMNIST image and its label

    Args:
        idx(int):
        x_train (:obj:`np.array`): (n_train, 28, 28) training images
        y_train (:obj:`np.array`): (n_train, 1) training labels
        x_test (:obj:`np.array`): (n_test, 28, 28) training images
        y_test (:obj:`np.array`y): (n_test, 1) training labels
        mapping (:obj:`np.array`): (62, 2) maps training labels (first column)
            to ascii codes (second column)
        mode (str): select image either from `train` or `test` dataset
    """
    if mode == 'train':
        img = x_train[idx]
        label = mapping[y_train[idx, 0]]
    else:
        img = x_test[idx]
        label = mapping[y_test[idx, 0]]
    plt.figure()
    plt.gray()
    plt.imshow(img)
    plt.show()
    print("Label: {}".format(label))

def show_train_progress(iteration, train_loss, test_loss, train_acc, test_acc):
    print(("Iteration: {:06d}, "
           "Loss (Train/Test): {:.3f} / {:.3f}, "
           "Accuracy: {:.3%} / {:.3%}").format(iteration,
                                               train_loss, test_loss,
                                               train_acc, test_acc))

def eval_serving_performance(n_examples, n_print_examples, request_url, seed=42,
                             dataset='test', use_graphpipe=False,
                             emnist_folder_path='emnist_data/'):
    """
    Fires queries against a model service, evaluates the result in terms
    of classification accuracy and processing time, and prints examples

    Args:
        n_examples (int): number of queries
        n_print_examples (int): number of examples to print img and classification
        request_url (str): URL of model service used for classification
        seed (int): seed for NumPy random number generator
        dataset (str): `train` or `test` data to pick evaluation examples from
        use_graphpipe (bool): use graphpipe client to execute queries (uses
                                flatbuffers instead)
        emnist_folder_path (str): folder containing EMNIST data

    Returns:
        durations ([int]): list of ms every request took
    """
    x_train, y_train, x_test, y_test, _ = load_emnist(emnist_folder_path)
    mapping = get_emnist_mapping()

    acc = 0
    durations = []

    if dataset == 'train':
        x_eval, y_eval = x_train, y_train
    else:
        x_eval, y_eval = x_test, y_test

    np.random.seed(seed)
    eval_img_indices = np.random.choice(np.arange(x_eval.shape[0]),
                                        n_examples,
                                        replace=False)

    for idx, test_img_idx in enumerate(eval_img_indices):
        test_img_flatten = x_eval[test_img_idx].reshape(1, 784) / 255

        if not use_graphpipe:
            start = time.time()
            test_img_payload = {'instances': test_img_flatten.tolist()}
            test_img_payload = json.dumps(test_img_payload)
            test_img_softmax_pred = requests.post(request_url,
                                                  data=test_img_payload)
            test_img_softmax_pred = test_img_softmax_pred.json()['predictions']
        else:
            start = time.time()
            test_img_softmax_pred = remote.execute(request_url, test_img_flatten)

        durations.append(int((time.time() - start) * 1000000)/1000)
        test_img_class_pred = np.argmax(test_img_softmax_pred)
        acc += (test_img_class_pred == y_eval[test_img_idx][0])

        # print the first 10 images with their true and predicted label
        if idx < n_print_examples:
            show_img(test_img_idx, x_train, y_train, x_test, y_test, mapping,
                     mode=dataset)
            print("Predicted Label: {}".format(mapping[test_img_class_pred]))

    print("Accuracy on {} test images: {:.2%}".format(n_examples,
                                                      acc / n_examples))

    return durations

def eval_throughput(duration, request_url, batch_size=1, dataset='test',
                    use_graphpipe=False, emnist_folder_path='emnist_data/'):
    """
    Performs a Throughput test running sequential queries against webservice
    counting the number of requests performed within duration

    Args:
        duration (int): number of seconds for running test queries
        request_url (str): URL of model service used for classification
        batch_size (int): no. of examples per batch
        dataset (str): `train` or `test` data to pick evaluation examples from
        use_graphpipe (bool): use graphpipe client to execute queries (uses
                                flatbuffers instead)
        emnist_folder_path (str): folder containing EMNIST data

    Returns:
        num_reqs (int): total number of requests performed
        reqs_per_second (float): thoughput as the number of requests performed
                                 per second

    """
    x_train, _, x_test, _, _ = load_emnist(emnist_folder_path)
    if dataset == 'train':
        data = x_train
    else:
        data = x_test
    data = data.reshape(-1, 1, 28 * 28) / 255

    num_reqs = 0
    _logger.info("Throughput Evaluation on URL {} for {} seconds ...".format(
            request_url, duration
    ))

    if not use_graphpipe:
        end_time = time.time() + duration
        while time.time() <= end_time:
            batch = data[num_reqs:(num_reqs+batch_size)].reshape(-1, 28*28)
            test_img_payload = \
                {"instances": batch.tolist()}
            test_img_payload = json.dumps(test_img_payload)
            test_img_softmax_pred = requests.post(request_url,
                                                  data=test_img_payload)
            test_img_softmax_pred.json()['predictions']
            num_reqs += batch_size
    else:
        end_time = time.time() + duration
        while time.time() <= end_time:
            remote.execute(request_url, data[num_reqs:(num_reqs+batch_size)])
            num_reqs += batch_size

    reqs_per_second = num_reqs / duration

    _logger.info(("Throughput Summary:\nURL: {}\nDuration: {} [s]\n"
                  + "Total Requests: {}\nRequests/Second: {:.2f}").format(
            request_url, duration, num_reqs, reqs_per_second
    ))

    return num_reqs, reqs_per_second


from pathlib import Path

import h5py
import numpy as np

from .polarview import PolarView

from hexrdgui.calibration.utils.maud_headers import header0, header, block_hdr
from hexrdgui.constants import PolarXAxisType, ViewType
from hexrdgui.create_hedm_instrument import create_hedm_instrument
from hexrdgui.hexrd_config import HexrdConfig
from hexrdgui.overlays import update_overlay_data
from hexrdgui.utils.conversions import tth_to_q

def polar_viewer():
    return InstrumentViewer()

class InstrumentViewer:

    def __init__(self):
        self.type = ViewType.polar
        self.instr = create_hedm_instrument()

        self.draw_polar()

    @property
    def all_detector_borders(self):
        return self.pv.all_detector_borders

    @property
    def angular_grid(self):
        return self.pv.angular_grid

    @property
    def raw_img(self):
        return self.pv.raw_img

    @property
    def raw_mask(self):
        return self.pv.raw_mask

    @property
    def snipped_img(self):
        return self.pv.snipped_img

    @property
    def img(self):
        return self.pv.img

    @property
    def snip_background(self):
        return self.pv.snip_background

    def update_angular_grid(self):
        self.pv.update_angular_grid()

    def update_image(self):
        self.pv.generate_image()

    def reapply_masks(self):
        self.pv.reapply_masks()

    def draw_polar(self):
        """show polar view of rings"""
        self.pv = PolarView(self.instr)
        self.pv.warp_all_images()

        '''
        tth_min = HexrdConfig().polar_res_tth_min
        tth_max = HexrdConfig().polar_res_tth_max
        eta_min = HexrdConfig().polar_res_eta_min
        eta_max = HexrdConfig().polar_res_eta_max

        self._extent = [tth_min, tth_max, eta_max, eta_min]  # l, r, b, t
        '''
        self._extent = np.degrees(self.pv.extent)  # l, r, b, t

    def update_overlay_data(self):
        update_overlay_data(self.instr, self.type)

    def update_detector(self, det):
        self.pv.update_detector(det)

    def write_image(self, filename='polar_image.npz'):
        filename = Path(filename)

        azimuthal_integration = np.array(
            HexrdConfig().last_unscaled_azimuthal_integral_data)

        if HexrdConfig().polar_x_axis_type == PolarXAxisType.q:
            # Convert to Q
            azimuthal_integration[0] = tth_to_q(azimuthal_integration[0],
                                                self.instr.beam_energy)

        # Re-format the data so that it is in 2 columns
        azimuthal_integration = azimuthal_integration.T

        # Lineout suffixes and their delimiters
        lineout_suffixes = {
            '.csv': ',',
            '.xy': ' ',
        }
        if filename.suffix in lineout_suffixes:
            # Just save the lineout
            delimiter = lineout_suffixes[filename.suffix]

            # Delete the file if it already exists
            if filename.exists():
                filename.unlink()

            np.savetxt(filename, azimuthal_integration, delimiter=delimiter)
            return

        intensities = self.raw_img.data
        intensities[self.raw_mask] = np.nan

        eta, tth = np.degrees(self.angular_grid)

        # Prepare the data to write out
        data = {
            'tth_coordinates': tth,
            'eta_coordinates': eta,
            'q_coordinates': tth_to_q(tth, self.instr.beam_energy),
            'intensities': intensities,
            'extent': self._extent,
            'azimuthal_integration': azimuthal_integration,
        }

        if self.snip_background is not None:
            # Add the snip background if it was used
            data['snip_background'] = self.snip_background

        # Add visible polar mask data if we have any
        for name, mask in HexrdConfig().masks.items():
            if name not in HexrdConfig().visible_masks:
                continue

            data[f'mask_{name}'] = mask

        # Delete the file if it already exists
        if filename.exists():
            filename.unlink()

        # Check the file extension
        if filename.suffix.lower() == '.npz':
            # If it looks like npz, save as npz
            np.savez(filename, **data)
        else:
            # Default to HDF5 format
            with h5py.File(filename, 'w') as f:
                for key, value in data.items():
                    f.create_dataset(key, data=value)

    def write_maud(self, filename='polar_to_maud.esg'):
        filename = Path(filename)

        with open(filename, 'w') as fid:
            eta_vec = np.degrees(self.angular_grid[0])
            intensities = self.img
            first_block = True
            for i, eta in enumerate(np.average(eta_vec, axis=1).flatten()):
                if np.all(np.isnan(intensities[i])):
                    # Skip this block
                    continue

                if first_block:
                    hstr = header0 % (i, np.linalg.norm(self.pv.tvec_s), eta)
                    first_block = False
                else:
                    hstr = header % (i, eta)

                fid.write(hstr)
                fid.write(block_hdr)
                tth = HexrdConfig().last_unscaled_azimuthal_integral_data[0]
                for rho, inten in zip(tth, intensities[i]):
                    if np.isnan(inten):
                        continue

                    vals = f' {rho}  {inten}\n'
                    fid.write(vals)


import requests
from bs4 import BeautifulSoup

'''
get job about keyward(ex: python, java, ...) from remoteok.com

all_job=[job1_dict,job2_dict]
job1_dict={
  "title":title,
  "company":company,
  "location":location,
  "introduction":introduction,
  "date":date,
  "link":link
}
'''

HEADERS_DEFAULT = {
    'User-Agent':
    'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_9_3) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/35.0.1916.47 Safari/537.36'
}

ROK_URL="https://remoteok.io"

def get_jobs(keyward,headers):
  req=requests.get(f"{ROK_URL}/remote-dev+{keyward}-jobs",headers=headers)
  soup=BeautifulSoup(req.text,"html.parser")
  jobcard_list=soup.find_all("tr",{"class":"job"})
  job_list=[]
  for jobcard in jobcard_list:
    #remove closed
    if jobcard.find("td",{"class":"company"}).find("span",{"class":"closed"}) is not None:
      continue
    #introduction
    introduction_list=jobcard.find("td",{"class":"tags"}).find_all('a')
    introduction="develop tool:"
    for info in introduction_list:
      tag=info.text.replace("\xa0","")
      introduction+=f" [{tag}]"
    #location
    try:
      location=jobcard.find("td",{"class":"company"}).find("div",{"class":"location"}).text
    except:
      location=None
    job={
      "title":jobcard.find("td",{"class":"company"}).find("h2",{"itemprop":"title"}).text,
      "company":jobcard.find("td",{"class":"company"}).find("h3",{"itemprop":"name"}).text,
      "location":location,
      "introduction":introduction,
      "date":jobcard.find("td",{"class":"time"}).find("time").text,
      "link":ROK_URL+jobcard.find("td",{"class":"source"}).find("a")["href"]
    }
    job_list.append(job)
  return job_list

def ROK_jobs(keyward,headers=HEADERS_DEFAULT):
  all_jobs=get_jobs(keyward,headers)
  print(f"keyward: {keyward}, website: remoteok, page: 1, extracted")
  return all_jobs


import discord
from discord.ext import commands
import os
from __main__ import settings, send_cmd_help
import logging
from cogs.utils.dataIO import dataIO
from .utils import checks
import re

# Some magic values
ALL_CHANNELS = 'all'
MODE_INCLUSIVE = 'incl'  # deletes all messages matching
MODE_EXCLUSIVE = 'excl'  # delete all messages not matching
MODE_DISABLED = 'none'   # pattern disabled
MODES = [MODE_DISABLED, MODE_EXCLUSIVE, MODE_INCLUSIVE]

log = logging.getLogger('red.recensor')

DATA_PATH = "data/recensor/"
JSON_PATH = DATA_PATH + "regexen.json"

class ReCensor:

    # Data format:
    # {
    #  serverid (str): {
    #   channelid (str): {
    #     regex (str): mode (str from MODES)
    #     }
    #   }
    # }

    def __init__(self, bot):
        self.bot = bot
        self.regexen = dataIO.load_json(JSON_PATH)
        self.recache = {}
        bot.loop.create_task(self.compile_regexen())

    def _re_present(self, obj):
        """Determines if any patterns are set for a server or channel"""
        if type(obj) is discord.Server:
            server = obj
            if server.id in self.regexen:
                for relist in self.regexen[server.id].values():
                    if bool(relist):  # nonempty list
                        return True
                return False
            else:
                return False

        elif type(obj) is discord.Channel:
            server = obj.server
            channel = obj
            if channel.id in self.regexen[server.id]:
                return bool(self.regexen[server.id][channel.id])
            else:
                return False

        elif type(obj) is str:  # won't work with ALL_CHANNELS
            channel = self.bot.get_channel(obj)
            server = channel.server
            if channel.id in self.regexen[server.id]:
                return bool(self.regexen[server.id][channel.id])
            else:
                return False

    def _ls_excl(self, server):
        """returns a list of channel IDs with exclusive filters"""
        clist = []
        if type(server) is discord.Server:
            server = server.id
        if server in self.regexen:
            for c, relist in self.regexen[server].items():
                if MODE_EXCLUSIVE in relist.values():
                    clist.append(c)
        return clist

    # Background cache regexen for speed
    async def compile_regexen(self):
        for s, channels in self.regexen.items():
            for regex in channels:
                self.recache[regex] = re.compile(regex)

    @commands.group(name='recensor', pass_context=True)
    @checks.mod_or_permissions(manage_messages=True)
    async def recensor(self, ctx):
        """Configure regular expression censorship"""
        if ctx.invoked_subcommand is None:
            await send_cmd_help(ctx)

    @recensor.command(pass_context=True, name='list')
    async def _list(self, ctx, channel: discord.Channel=None):
        """Lists regexes used to filter messages.
        Channel listing includes global patterns."""
        server = ctx.message.server
        self.regexen = dataIO.load_json(JSON_PATH)
        if not self._re_present(server):
            await self.bot.say('There are no filter patterns set for this server.')
            return
        table = ' | '.join(['mode', 'pattern']) + '\n'  # header
        for c in self.regexen[server.id]:
            if c == ALL_CHANNELS and self._re_present(server):
                    table += '\nServer-wide:\n'
            elif (channel and channel.id == c) or not channel:
                if channel:
                    ch_obj = channel
                else:
                    ch_obj = self.bot.get_channel(c)
                if ch_obj is None:
                    table += '\n' + 'Channel ID %s (deleted):' % c + '\n'
                if self._re_present(ch_obj):
                    table += '\n#' + ch_obj.name + '\n'

            for regex, mode in self.regexen[server.id][c].items():
                table += ' | '.join([mode, regex]) + '\n'
        await self.bot.say('```py\n' + table + '```')

    @recensor.command(pass_context=True, name='add')
    async def _add(self, ctx, pattern: str, mode: str=MODE_INCLUSIVE, channel: discord.Channel=None):
        """Adds a pattern to filter messages. Mods, bot admins, and the bot's
        owner are not subjected to the filter.
        If the pattern contains spaces, it must be put in double quotes. Single quotes will not work.

        mode is one of:
        incl: Default, filter messages that match the pattern
        excl: filter non-matching, only one allowed per channel or server
        none: adds pattern to storage but doesn't apply filtering. Use recensor set to enable.

        To use channel, mode must also be specified. If channel is not specified,
        the filter is used across the entire server."""
        server = ctx.message.server

        # initialize
        self.regexen = dataIO.load_json(JSON_PATH)
        if server.id not in self.regexen:
            self.regexen[server.id] = {}

        if pattern.startswith("'"):
            await self.bot.say("Patterns cannot be specified within single quotes.")
            return

        if mode not in MODES:
            await self.bot.say('"%s" is not a valid mode. You must specify one of `%s`.' % (mode, '`, `'.join(MODES)))
            return
        if mode == MODE_EXCLUSIVE:
            if ALL_CHANNELS in self._ls_excl(server):
                await self.bot.say("There is already a server-wide exclusive filter. Remove or disable it first.")
                return
            if channel and channel.id in self._ls_excl(server):
                await self.bot.say("That channel already has an exclusive filter. Remove or disable it first.")
                return
        cid = channel.id if channel else ALL_CHANNELS
        if cid not in self.regexen[server.id]:
            self.regexen[server.id][cid] = {}
        self.regexen[server.id][cid][pattern] = mode
        await self.bot.say('Pattern added.')
        dataIO.save_json(JSON_PATH, self.regexen)

    @recensor.command(pass_context=True, name='set')
    async def _set(self, ctx, mode: str, channel: discord.Channel=None):
        """Lists regexes used to filter messages"""
        server = ctx.message.server
        self.regexen = dataIO.load_json(JSON_PATH)
        if not self._re_present(server):
            await self.bot.say('There are no patterns in the server to modify.')
            return
        if mode not in MODES:
            self.bot.reply('"%s" is not a valid mode. You must specify one of `%s`.') % (mode, '`, `'.join(MODES))

        if mode == MODE_EXCLUSIVE:
            if ALL_CHANNELS in self._ls_excl(server):
                await self.bot.say("There is already a server-wide exclusive filter. Remove or disable it first.")
                return
            if channel and channel.id in self._ls_excl(server):
                await self.bot.say("That channel already has an exclusive filter. Remove or disable it first.")
                return

        re_list = {}
        i = 1
        table = ' | '.join(['#'.ljust(4), 'mode', 'pattern']) + '\n'  # header
        for c in self.regexen[server.id]:
            if c == ALL_CHANNELS and self._re_present(server):
                    table += '\nServer-wide:\n'
            elif (channel and channel.id == c) or not channel:
                if channel:
                    ch_obj = channel
                else:
                    ch_obj = self.bot.get_channel(c)
                if ch_obj is None:
                    table += '\n' + 'Channel ID %s (deleted):' % c + '\n'
                if self._re_present(ch_obj):
                    table += '\n#' + ch_obj.name + '\n'

            for regex, oldmode in self.regexen[server.id][c].items():
                table += ' | '.join([str(i).ljust(4), oldmode, regex]) + '\n'
                re_list[str(i)] = (server.id, c, regex, oldmode)
                i += 1
        prompt = 'Choose the number of the pattern to set to `%s`:\n' % mode
        await self.bot.say(prompt + '```py\n' + table + '```')

        msg = await self.bot.wait_for_message(author=ctx.message.author, timeout=15)
        if msg is None:
            return
        msg = msg.content.strip()
        if msg in re_list:
            sid, cid, regex, _ = re_list[msg]
            self.regexen[sid][cid][regex] = mode
            await self.bot.say('Mode set.')
            dataIO.save_json(JSON_PATH, self.regexen)

    @recensor.command(pass_context=True, name='del')
    async def _del(self, ctx, channel: discord.Channel=None):
        """Lists regexes used to filter messages"""
        server = ctx.message.server
        self.regexen = dataIO.load_json(JSON_PATH)
        if not self._re_present(server):
            await self.bot.say('There are no filter patterns set for this server.')
            return
        re_list = {}
        i = 1
        table = ' | '.join(['#'.ljust(4), 'mode', 'pattern']) + '\n'  # header
        for c in self.regexen[server.id]:
            if c == ALL_CHANNELS and self._re_present(server):
                    table += '\nServer-wide:\n'
            elif (channel and channel.id == c) or not channel:
                if channel:
                    ch_obj = channel
                else:
                    ch_obj = self.bot.get_channel(c)
                if ch_obj is None:
                    table += '\n' + 'Channel ID %s (deleted):' % c + '\n'
                if self._re_present(ch_obj):
                    table += '\n#' + ch_obj.name + '\n'

            for regex, mode in self.regexen[server.id][c].items():
                table += ' | '.join([str(i).ljust(4), mode, regex]) + '\n'
                re_list[str(i)] = (server.id, c, regex)
                i += 1
        prompt = 'Choose the number of the pattern to delete:\n'
        await self.bot.say(prompt + '```py\n' + table + '```')
        msg = await self.bot.wait_for_message(author=ctx.message.author, timeout=15)
        if msg is None:
            return
        msg = msg.content.strip()
        if msg in re_list:
            sid, cid, regex = re_list[msg]
            del(self.regexen[sid][cid][regex])
            await self.bot.say('Pattern removed.')
        dataIO.save_json(JSON_PATH, self.regexen)

    def immune_from_filter(self, message):
        """Tests message to see if it is exempt from filter. Taken from mod.py"""
        user = message.author
        server = message.server
        admin_role = settings.get_server_admin(server)
        mod_role = settings.get_server_mod(server)

        if user.id == settings.owner:
            return True
        elif discord.utils.get(user.roles, name=admin_role):
            return True
        elif discord.utils.get(user.roles, name=mod_role):
            return True
        else:
            return False

    async def on_message(self, message):
        # Fast checks
        if message.channel.is_private:
            return
        if message.author.id == self.bot.user.id:
            return

        server = message.server
        sid = server.id
        can_delete = message.channel.permissions_for(server.me).manage_messages

        # Owner, admins and mods are immune to the filter
        if self.immune_from_filter(message) or not can_delete:
            return

        if sid in self.regexen:
            patterns = {}
            # compile list of patterns from global and channel
            for key in [ALL_CHANNELS, message.channel.id]:
                if key in self.regexen[sid]:
                    patterns.update(self.regexen[sid][key])
            # Iterate through patterns
            for regex, mode in patterns.items():
                # Skip disabled patterns
                if mode == MODE_DISABLED:
                    continue
                regex = self.recache[regex] if regex in self.recache else re.compile(regex)
                if (mode == MODE_EXCLUSIVE) != bool(regex.match(message.content)):  # xor
                    await self.bot.delete_message(message)

def check_folder():
    if not os.path.exists(DATA_PATH):
        log.debug('Creating folder: %s' % DATA_PATH)
        os.makedirs(DATA_PATH)

def check_file():
    if dataIO.is_valid_json(JSON_PATH) is False:
        log.debug('Creating json: %s' % JSON_PATH)
        dataIO.save_json(JSON_PATH, {})

def setup(bot):
    check_folder()
    check_file()
    n = ReCensor(bot)
    bot.add_cog(n)


from firebase import firebase
import cv2

FBConn = firebase.FirebaseApplication("https://pikachut-8f775-default-rtdb.firebaseio.com/",None)
johnsmith = cv2.imread('johnsmith.jpg')

while True:
    temperature = int(input("What is the temp?"))
    data_to_upload = {
        'Temp':temperature,
        'Name': johnsmith
    }

    result=FBConn.post('/MyTestData/',data_to_upload)

    print(result)


from collections import Counter
def oneOrTwoEditsAway(str1, str2):
	if abs(len(str1) - len(str2)) > 1:
		return False
	elif len(str1) == len(str2):
		cntr = 0
		for s1, s2 in zip(str1, str2):
			if s1 != s2:
				cntr += 1
				if cntr > 1:
					return False
		return True

	short_str = str1
	long_str = str2
	if len(str1) > len(str2):
		short_str = str2
		long_str = str1

	cntr1 = Counter(short_str)
	cntr2 = Counter(long_str)
	unique_char = ""
	unique_counter = 0
	for char in long_str:
		if cntr1[char] != cntr2[char]:
			unique_char = char
			unique_counter += 1
			if unique_counter > 1:
				return False
	if unique_char + short_str == long_str or short_str + unique_char == long_str:
		return True

	for i in range(len(short_str)):
		fixed_str = short_str[0:i+1] + unique_char + short_str[i+1:]
		if fixed_str == long_str:
			return True
	return False

print(oneOrTwoEditsAway("pale", "ple"))
print(oneOrTwoEditsAway("pales", "pale"))
print(oneOrTwoEditsAway("pale", "bale"))
print(oneOrTwoEditsAway("pale", "bake"))
print(oneOrTwoEditsAway("paler", "balre"))
print(oneOrTwoEditsAway("pale", "pake"))


from torch.utils.data import random_split
import matplotlib.pyplot as plt
from data import get_device, to_device, SepsisDataset, RandomDataset
from net import FederatedNet
from client import Client
from args import parse_args
import plotext as plt

def main():
    args = parse_args()
    num_clients = args.clients

    if args.sepsis:
        classes = 2
        input_dim = 3

        train_dataset = SepsisDataset(train=True, num_clients=num_clients)
        test_dataset = SepsisDataset(train=False, num_clients=num_clients)

        train_dataset, dev_dataset = random_split(train_dataset, [int(len(train_dataset) * 0.83), int(len(train_dataset) * 0.17)])

        print(f'{len(train_dataset)} {len(dev_dataset)} {len(test_dataset)}')

    elif args.random:
        classes = 3
        input_dim = 4

        train_dataset = RandomDataset(train=True)
        test_dataset = RandomDataset(train=False)

        train_dataset, dev_dataset = random_split(train_dataset, [int(len(train_dataset) * 0.83), int(len(train_dataset) * 0.17)])

        print(f'{len(train_dataset)} {len(dev_dataset)} {len(test_dataset)}')

    else:
        print("Error: please specify data set argument value")

    total_train_size = len(train_dataset)
    total_test_size = len(test_dataset)
    total_dev_size = len(dev_dataset)

    # num_clients = args.clients
    rounds = args.epochs
    batch_size = args.batch
    epochs_per_client = 3
    learning_rate = 2e-2

    device = get_device()

    examples_per_client = total_train_size // num_clients
    client_datasets = random_split(train_dataset, [min(i + examples_per_client,
            total_train_size) - i for i in range(0, total_train_size, examples_per_client)])
    clients = [Client('client_' + str(i), client_datasets[i]) for i in range(num_clients)]

    global_net = to_device(FederatedNet(input_dim, classes), device)

    history = []
    round_acc = []

    for i in range(rounds):
        print('Start Round {} ...'.format(i + 1))
        curr_parameters = global_net.get_parameters()
        new_parameters = dict([(layer_name, {'weight': 0, 'bias': 0}) for layer_name in curr_parameters])
        for client in clients:
            client_parameters = client.train(curr_parameters, device, input_dim, classes, epochs_per_client, learning_rate, batch_size)
            fraction = client.get_dataset_size() / total_train_size
            for layer_name in client_parameters:
                new_parameters[layer_name]['weight'] += fraction * client_parameters[layer_name]['weight']
                new_parameters[layer_name]['bias'] += fraction * client_parameters[layer_name]['bias']
        global_net.apply_parameters(new_parameters)

        train_loss, train_acc = global_net.evaluate(train_dataset, device, batch_size)
        dev_loss, dev_acc = global_net.evaluate(dev_dataset, device, batch_size)
        print('After round {}, train_loss = {}, dev_loss = {}, dev_acc = {}\n'.format(i + 1, round(train_loss, 4),
                round(dev_loss, 4), round(dev_acc, 4)))
        history.append((train_loss, dev_loss))
        round_acc.append(train_acc)

    plt.plot(round_acc)
    plt.theme('matrix')
    plt.plotsize(50, 15)
    plt.show()

if __name__ == '__main__':
    main()


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import re
import os.path
from setuptools import setup, find_packages

def get_version_string():
    with open(os.path.join(package_path, 'src/QuantWorkshop/__init__.py'), 'r') as f:
        version_line = re.search(
            r'__version__\s+=\s+(.*)', f.read()
        ).group(1)
    return version_line

package_path = os.path.dirname(os.path.abspath(__file__))

setup(
    name='QuantWorkshop',
    version=get_version_string(),
    description='This is a test of the setup',
    long_description='long_description',
    long_description_content_type="text/markdown",
    url="https://github.com/kenblikylee/imgkernel",
    author='Bruce Frank Wong',
    author_email='Bruce.Frank.Wong@gmail.com',
    packages=find_packages(where='src'),
    package_dir={'': 'src'},
    install_requires=open('requirements.txt').read().splitlines(),
    setup_requires=[
        'pytest-runner',
    ],
    tests_require=[
        'pytest',
    ],
)


'''
receipt.py
======================
order management module of POS_packages
'''

import pandas as pd
import openpyxl as xl
from datetime import datetime
from openpyxl.utils.dataframe import dataframe_to_rows
from calculator import POS_calculator, cal

# receive order and save data(excel)
class POS_receipt:
    '''POS의 주문 기능, 주문내역 저장(excel) 기능을 하는 클래스입니다.

    '''
    def __init__(self):
        '''생성자
        :explain: order_history.xlsx 파일이 경로에 이미 있다면 불러오고, 없다면 새로 생성합니다.
        receipt/summary worksheet를 불러오거나, 생성해서 지정합니다.
        '''
        self.details_cols_list = ['주문일시', '메뉴', '수량', '단가', '금액', '지불수단']
        self.receipt_df = pd.DataFrame(columns=self.details_cols_list)

        self.summary_cols_list = ['주문일시', '금액']
        self.summary_df = pd.DataFrame(columns=self.summary_cols_list)

        self.wb_filename = r'C:\Users\JH\POS_packages\order_history.xlsx'

        try:
            self.wb = xl.load_workbook(filename=self.wb_filename)
            self.details_ws = self.wb['receipt']
            self.summary_ws = self.wb['summary']
        except FileNotFoundError:
            self.wb = xl.Workbook(self.wb_filename)
            self.details_ws = self.wb.active
            self.details_ws.title = 'receipt'
            self.details_ws.append(self.details_cols_list)
            self.summary_ws = self.wb.create_sheet()
            self.summary_ws.title = 'summary'
            self.summary_ws.append(self.summary_cols_list)
            self.wb.save(filename=self.wb_filename)

    def receive_order(self, payment_method_str):
        '''
        :explain: 입력한 지불수단에 따른 주문내역 dataframe을 생성하고 반환합니다. 장바구니 비우기, 요약내역 생성, excel에 저장하는 함수를 호출합니다.

        :param string payment_method_str: 지불수단
        :return dataframe: receipt_df 주문내역
        '''
        self.receipt_df = pd.DataFrame(columns=self.details_cols_list)

        datetime_now_str = datetime.now().strftime('%Y-%m-%d %H:%M:%S')

        for i in range(len(cal.cart_list)):
            if cal.cart_list[i] not in cal.cart_list[:i]:
                details_list = [datetime_now_str, cal.cart_list[i], cal.return_menu_qty(cal.cart_list[i]),
                                cal.return_menu_cost(cal.cart_list[i]), cal.return_menu_amount(cal.cart_list[i]),
                                payment_method_str]
                self.receipt_df = self.receipt_df.append(pd.DataFrame([details_list], columns=self.details_cols_list),
                                                         ignore_index=True)

        cal.cancel_all_qty()
        self.create_summary_df()
        self.save_receipt_df_to_excel()
        self.save_summary_df_to_excel()

        return self.receipt_df

    def create_summary_df(self):
        '''
        :explain: 요약내역 dataframe을 생성하고 반환합니다.

        :return dataframe: summary_df 요약내역
        '''
        self.summary_df = pd.DataFrame(columns=self.summary_cols_list)

        summary_list = [self.receipt_df['주문일시'][0], sum(self.receipt_df['금액'])]
        self.summary_df = self.summary_df.append(pd.DataFrame([summary_list], columns=self.summary_cols_list),
                                                 ignore_index=True)
        return self.summary_df

    def save_receipt_df_to_excel(self):
        '''
        :explain: 주문내역 dataframe을 excel로 저장합니다.

        '''
        for row in dataframe_to_rows(self.receipt_df, index=False, header=False):
            self.details_ws.append(row)
            self.wb.save(filename=self.wb_filename)

    def save_summary_df_to_excel(self):
        '''
        :explain: 요약내역 dataframe을 excel로 저장합니다.

        '''
        for row in dataframe_to_rows(self.summary_df, index=False, header=False):
            self.summary_ws.append(row)
            self.wb.save(filename=self.wb_filename)

rec = POS_receipt()


from django.db import models

# Create your models here.
from datetime import timezone

from django.db import models

# Create your models here.

class Order(models.Model):
    order_status=(
        (1, "交易中"),
        (2, "已完成"),
        (3, "已取消")
    )
    p_id = models.AutoField(primary_key=True, verbose_name="订单序号")
    id = models.CharField(max_length=1000,default="000",verbose_name="订单流水号")
    #   entid = models.ForeignKey("user.Enterprise",on_delete=models.CASCADE,verbose_name="此订单所属企业") --从所属需求中读取
    farmers = models.ManyToManyField("user.Farmers",verbose_name="组")  # 组
    money = models.FloatField(default=0,verbose_name="总交易金额") #交易金额
    moneyToFarmers = models.FloatField(default=0, verbose_name="给农工金额")
    moneyToApp = models.FloatField(default=0, verbose_name="给软件金额")
    needId = models.ForeignKey("needs.Needs",on_delete=models.CASCADE,verbose_name="此订单所属需求",null=True)
    beginTime = models.DateTimeField(auto_now_add=True) #订单创建时间
    lastModified = models.DateTimeField(auto_now=True) #最后一次更改时间
    status = models.CharField(max_length=200, default="交易中")


from os import path
from tempfile import TemporaryDirectory

from pytest import raises, mark

from duplicate_file_removal.file_record import RecordStatus
from duplicate_file_removal.processor import RecordsProcessor
from tests.utils.util import ignore_exceptions, count, no_stdout

def test_remove_duplicates():
    with raises(RuntimeError):  # make sure paths provided have to be absolute
        RecordsProcessor.remove_duplicates([], ["..\\SomeFolder"])

def test_scanner_results_to_groups(scanner, gen_files_by_specification):
    num_of_dup_files = 10
    num_of_unique_files = 5

    base_dir = gen_files_by_specification(num_of_dup_files, num_of_unique_files)
    res = scanner.scan(base_dir.name)
    res = RecordsProcessor.scanner_results_to_groups(res)
    assert len(res) == 1  # occurrences for unique files filtered, and 1 occurrence for the 10 duplicate files

    dup_files_list = next(iter(res.values()))  # get the value of the first element
    assert len(dup_files_list) == num_of_dup_files

def test_remove_duplicate_normalized(binary_data, scanner):
    """
    Test to insure that given a more specific priority in a less specificity,
    then specificity > priority
    """
    temp_dir_1 = TemporaryDirectory()
    temp_file_name_1 = path.join(temp_dir_1.name, "file_1")
    with open(temp_file_name_1, "wb") as f:
        f.write(binary_data)

    temp_dir_2 = TemporaryDirectory(dir=temp_dir_1.name)
    temp_file_name_2 = path.join(temp_dir_2.name, "file_2")
    with open(temp_file_name_2, "wb") as f:
        f.write(binary_data)

    records = scanner.scan(temp_dir_1.name)
    assert len(records) == 2  # sanity check
    priority = ["C:\\", temp_dir_2.name]

    RecordsProcessor._remove_duplicate_normalized(records, priority)

    assert records[0].status == RecordStatus.deleted
    assert records[1].status == RecordStatus.exists

    # TemporaryDirectory cleanup warning ignored
    ignore_exceptions(temp_dir_1.cleanup)
    ignore_exceptions(temp_dir_2.cleanup)

# noinspection PyTypeChecker
def test_highest_specificity():
    chosen_one = object()

    candidates = [("C:\\", object()), ("C:\\specific", object()), ("C:\\specific\\choose_me", chosen_one)]
    assert chosen_one is RecordsProcessor._highest_specificity(candidates)

@mark.parametrize("is_simulation,expected", [(False, 2), (True, 3)])
def test_delete_records(is_simulation, expected, scanner, gen_files_by_specification):
    # If changed, adjust parametrize value accordingly
    num_of_dup_files = 2
    num_of_unique_files = 5

    base_dir = gen_files_by_specification(num_of_dup_files, num_of_unique_files)
    res = RecordsProcessor.scanner_results_to_groups(scanner.scan(base_dir.name))
    duplicate_records = next(iter(res.values()))

    res.add(duplicate_records[0])  # Try to 'cheat' the scanner and have two FileRecords to the same physical file
    with no_stdout():
        RecordsProcessor.delete_records(duplicate_records[0], duplicate_records, is_simulation)

    assert count(filter(lambda x: x.status == RecordStatus.exists, duplicate_records)) == expected


import torch
import pytorch_lightning as pl
from tokenizers import Tokenizer
from transformers import AdamW

from utils.token_utils import get_vocab_size, get_token_id, get_weights
from utils.utils import fussed_lasso, calc_acc, calculate_mask_percentage, get_pad_id, calc_perplexity, \
    calc_cross_entropy_batch_wise, calc_policy_loss

class LightingReinforceRationalizedLanguageModel(pl.LightningModule):
    '''
    PL wrapper for training a language model together with a rational extractor.
    '''

    def __init__(self, language_model, rational_extractor, tokenizer, hparams=None,
                 sparsity_weight=0.0001,
                 fussed_lasso_weight=0.1):
        super().__init__()
        self.hparams = hparams
        self.language_model = language_model
        self.rational_extractor = rational_extractor
        self.tokenizer = tokenizer

        self.sparsity_weight = sparsity_weight
        self.fussed_lasso_weight = fussed_lasso_weight

        self.log_list = [
            "total_loss", "acc", "total_mask_loss", "mask_mean", "mask_fussed_lasso", "cross_entropy_loss",
            "perplexity",
        ]

        self.pad_token_id = get_pad_id(tokenizer)

        self.freeze_language_model = hparams["freeze_language_model"]

    def forward(self, x, targets, ):

        rational = self.get_rational(x)

        masked_input = rational["masked_input"]

        prediction = self.forward_masked_input(masked_input, targets)
        return {"logits": prediction, **rational, "x": x}

    def get_rational(self, x, greedy=False):
        if greedy:
            print(greedy)
            return self.rational_extractor(x, greedy=greedy)
        else:
            return self.rational_extractor(x)

    def forward_masked_input(self, masked_input, targets):
        ## Concatenate the two together and put through the lstm
        lm_in = torch.cat([masked_input, targets])

        prediction = self.language_model(lm_in)
        return prediction

    def get_scores(self, forward_dict, targets):
        """
        Gets a dict with multiply scores of the give prediction
        prediction
        forward_dict: {logits, mask, chosen_policy}
        """
        targets = targets.long()

        n_targets = targets.shape[0]
        predictions = forward_dict["logits"][-(n_targets + 1):-1]

        total_h_loss = 0
        h_mean = 0
        fussed_lasso_loss = 0

        if "mask" in forward_dict.keys():
            ### Make sure batch first.
            mask = forward_dict["mask"].permute(1, 0).float()
            tokens = forward_dict["x"].permute(1, 0).float()
            h_mean = calculate_mask_percentage(tokens, mask, reduce=False,
                                               pad_id=get_token_id(self.tokenizer, "pad_token"))
            fussed_lasso_loss = fussed_lasso(tokens, mask, reduce=False,
                                             pad_id=get_token_id(self.tokenizer, "pad_token"))
            total_h_loss = self.sparsity_weight * h_mean + self.fussed_lasso_weight * fussed_lasso_loss

        cross_entropy_loss = calc_cross_entropy_batch_wise(predictions, targets, self.tokenizer)
        ##Make sure it is mean per batch

        rewards = cross_entropy_loss + total_h_loss
        vocab_size = get_vocab_size(self.tokenizer)
        perplexity = calc_perplexity(predictions,
                                     targets, self.tokenizer)

        # Get the policy loss. (old one)
        # total_loss = -torch.mean(rewards.detach() * torch.log(forward_dict["chosen_policy"]))
        total_loss = calc_policy_loss(rewards, forward_dict["chosen_policy"])

        if not self.freeze_language_model:
            total_loss += torch.mean(rewards)

        acc = calc_acc(predictions.reshape(-1, vocab_size),
                       targets.flatten(), exclude=self.pad_token_id)

        return {"total_loss": total_loss, "acc": acc, "total_mask_loss": total_h_loss.mean(),
                "mask_mean": h_mean.mean(),
                "mask_fussed_lasso": fussed_lasso_loss.mean(), "cross_entropy_loss": cross_entropy_loss.mean(),
                "perplexity": perplexity}

    def batch_to_out(self, batch):
        rational_in = batch[0].permute(1, 0)
        targets = batch[1].permute(1, 0)

        out = self.forward(rational_in, targets)

        scores = self.get_scores(out, targets)

        return scores

    def training_step(self, batch, batch_idx):

        batch_out = self.batch_to_out(batch)

        self.log_results(batch_out)

        return batch_out["total_loss"]

    @torch.no_grad()
    def validation_step(self, batch, batch_idx):

        batch_out = self.batch_to_out(batch)

        self.log_results(batch_out, prepend="val ")

        return batch_out["total_loss"]

    def complete_dialogues(self, sentences, total_length, with_rational=True, greedy_rationals=True):
        return [self.complete_dialogue(sentence, total_length=total_length, with_rational=with_rational,
                                       greedy_rationals=greedy_rationals) for sentence in
                sentences]

    def complete_dialogue(self, completed_dialogue, total_length=100, with_rational=True,
                          greedy_rationals=True):

        if type(self.tokenizer) == Tokenizer:
            encoding = self.tokenizer.encode(completed_dialogue).ids
        else:
            encoding = self.tokenizer.encode(completed_dialogue)

        all_tokens = torch.tensor(encoding).to(self.device).unsqueeze(1)

        ids_tensor = all_tokens
        rationals = []
        responses = []
        rationalized_input = []
        while (len(ids_tensor)) < total_length:

            if with_rational:
                # Get the rational
                rational = self.rational_extractor(ids_tensor, greedy=greedy_rationals)

                # Map back to tokens
                rational_input = self.tokenizer.decode(rational["masked_input"].long().view(-1).detach().cpu().numpy(),
                                                       skip_special_tokens=False).replace(" #",
                                                                                "").replace(
                    "#", "")
                next_input = rational["masked_input"]
                # The mask
                rationals.append(rational["mask"].flatten())

            else:
                next_input = ids_tensor
                rational_input = self.tokenizer.decode(ids_tensor.long().view(-1).detach().cpu().numpy(),
                                                       skip_special_tokens=False).replace(" #",
                                                                                "").replace(
                    "#", "")
            rationalized_input.append(rational_input)

            # Generate next ids based on the masked input
            # next_ids = self.language_model.generate_next_tokens(next_input, n_tokens=n_rational)
            next_ids = self.language_model.next_utterance(next_input.flatten(),
                                                          get_token_id(self.tokenizer, "sep_token"),
                                                          max_length=20).reshape(-1, 1)
            # next_ids = self.language_model.lm.generate(next_input.reshape(1,-1),
            #                                             eos_token_id=self.tokenizer.sep_token_id,
            #                                             num_beams=5,
            #                                             min_length=5,
            #                                             max_length=1000,#(20+len(next_input)),
            #                                             forced_eos_token_id=self.tokenizer.sep_token_id,
            #                                             ).reshape(-1,1) #TODO keep it to check if it will work in the future
            # Add to all tokens
            all_tokens = torch.cat([all_tokens, next_ids])

            # Map back to the sentence
            responses.append(
                # self.tokenizer.decode(next_ids.reshape(-1).detach().cpu().numpy(), skip_special_tokens=False).replace(" #", "").replace("#", "")
                self.tokenizer.decode(next_ids.reshape(-1).detach().cpu().numpy(), skip_special_tokens=False)
            )

            # Map back to tensor
            ids_tensor = all_tokens

        completed_dialogue = self.tokenizer.decode(all_tokens.reshape(-1).detach().cpu().numpy(),
                                                   skip_special_tokens=False).replace(" #", "").replace("#", "")

        result = {"completed_dialogue": completed_dialogue, "rationals": rationals,
                  "rationalized_input": rationalized_input,
                  "response": responses}

        return result

    def configure_optimizers(
            self,
    ):
        if not self.freeze_language_model:
            parameters = list(self.language_model.parameters()) + list(self.rational_extractor.parameters())
        else:
            parameters = list(self.rational_extractor.parameters())

        optimizer = AdamW(
            parameters,
            lr=self.hparams['learning_rate'])
        return optimizer

    def log_results(self, result, prepend=""):

        for k in self.log_list:
            self.log(prepend + k, result[k], on_step=True, on_epoch=True)


import numpy as np
import cv2
import h5py
import pandas as pd
from scipy.stats import entropy

class FeatureExtractor:
    def __init__(self, batch_size=10000, img_size=(224,224)):
        (self.image_w, self.image_h) = img_size
        self.batch_size = batch_size

        # opencv body detection model
        self.opencv_path = './opencv/'
        self.full_body_detector = cv2.CascadeClassifier(self.opencv_path + 'haarcascade_fullbody.xml')
        self.upper_body_detector = cv2.CascadeClassifier(self.opencv_path + 'haarcascade_upperbody.xml')
        self.lower_body_detector = cv2.CascadeClassifier(self.opencv_path + 'haarcascade_lowerbody.xml')
        self.face_detector = cv2.CascadeClassifier(self.opencv_path + 'haarcascade_frontalface_alt2.xml')

        # size(big, medium, small)
        self.size_degree = (0.55, 0.35)
        # classify color(HSV not RGB)
        self.colors = ((0, 0, 255, "white"),
                      (0, 0, 0, "black"),
                      (0, 0, 120, "grey"),
                      (176, 215, 190, "red"),
                      (23, 151, 220, "yellow"),
                      (35, 132, 139, "khaki"),
                      (105, 85, 180, "blue"),
                      (120, 64, 80, "navy"),
                      (13, 89, 100, "brown"),
                      (15, 113, 180, "beige"))

    # Main
    def extract_feature(self, image_file='temp/images.hdf5', masks_file='temp/masks.hdf5'):
        # check detector loaded
        if self.full_body_detector.empty() or self.upper_body_detector.empty() or self.lower_body_detector.empty() or self.face_detector.empty():
            print('Classifier XML not Loaded!')
            return None

        # load data
        try :
            with h5py.File(image_file, 'r') as f, h5py.File(masks_file, 'r') as m:
                num_X = len(f['image'])

                ##########################################################
                #num_X = 1000
                #self.batch_size = 1000
                self.a = np.zeros((num_X, 3), dtype=np.float32)

                # initialization
                color = np.empty(num_X, dtype=np.dtype("U10"))
                bg_color = np.empty(num_X, dtype=np.dtype("U10"))
                comp = np.zeros(num_X, dtype=np.uint8)
                size = np.zeros(num_X, dtype=np.uint8)
                model_type = np.zeros(num_X, dtype=np.uint8)

                for i in range(0, num_X, self.batch_size):
                    # lazy load data(hdf5)
                    X = f['image'][i:i+self.batch_size]
                    masks = m['mask'][i:i+self.batch_size]
                    print('{}~{} Data Loaded'.format(i, i+len(masks)))

                    # 상품, 배경 색감 추출
                    color[i:i+self.batch_size], bg_color[i:i+self.batch_size] = self.extract_color(X, masks)
                    # 배경의 복잡도
                    comp[i:i+self.batch_size] = self.image_complexity(X, masks)
                    # 상품 크기 비율
                    size[i:i+self.batch_size] = self.estimate_size(masks)
                    # 모델의 노출도
                    model_type[i:i+self.batch_size] = self.detect_body(X)

                df = pd.DataFrame({'id': f['index'], 'color': color, 'bg_color': bg_color, 'bg_complexity': comp, 'size': size, 'model_type': model_type})
                print('Feature Extracted.')
        except :
            df = None
        return df

    def nearest_color(self, x):
        return min(self.colors, key=lambda color: sum((c - x_) ** 2 for c, x_ in zip(color, x)))[-1]

    def extract_color(self, X, masks):
        # reshape
        masks_ = np.reshape(masks, masks.shape[:-1]).astype(np.bool)
        # extract color
        num_X = len(X)
        color = np.zeros(num_X, dtype=np.dtype("U10"))
        bg_color = np.zeros(num_X, dtype=np.dtype("U10"))

        # convert BGR(opencv default) to HSV
        X_ = np.reshape(X, (num_X * self.image_w, self.image_h, 3))
        X_ = cv2.cvtColor(X_, cv2.COLOR_BGR2HSV)
        X_ = np.reshape(X_, X.shape)

        # test
        for i, (x, m) in enumerate(zip(X_, masks_)):
            #print(self.nearest_color(colors, np.mean(X[i][masks[i]], axis=0).astype(int)))
            color[i] = self.nearest_color(np.mean(x[m], axis=0).astype(np.uint8))
            #print(self.nearest_color(colors, np.mean(X[i][np.logical_not(masks[i])], axis=0).astype(int)))
            bg_color[i] = self.nearest_color(np.mean(x[np.logical_not(m)], axis=0).astype(np.uint8))

        print('extract_color Done.')
        return color, bg_color

    # 배경 복잡도 (1: 복잡함 0:안복잡함)
    def image_complexity(self, X, masks):
        # 복잡도 기준
        threshold = 25.0
        # reshape
        # masks = np.reshape(masks, masks.shape[:-1])

        R, G, B = X[:, :, :, 0:1], X[:, :, :, 1:2], X[:, :, :, 2:]

        # compute rg = R - G
        rg = np.absolute(R - G)

        # compute yb = 0.5 * (R + G) - B
        yb = np.absolute(0.5 * (R + G) - B)

        _masks = np.logical_not(masks)
        # compute the mean and standard deviation of both `rg` and `yb` ()
        rb_masked = np.ma.masked_where(_masks, rg)
        (rb_mean, rb_std) = (np.mean(rb_masked, axis=(1, 2, 3)), np.std(rb_masked, axis=(1, 2, 3)))

        yb_masked = np.ma.masked_where(_masks, yb)
        (yb_mean, yb_std) = (np.mean(yb_masked, axis=(1, 2, 3)), np.std(yb_masked, axis=(1, 2, 3)))

        # combine the mean and standard deviations
        std_root = np.sqrt((rb_std ** 2) + (yb_std ** 2))
        mean_root = np.sqrt((rb_mean ** 2) + (yb_mean ** 2))

        # derive the "colorfulness" metric and return it
        complexity = std_root + (0.3 * mean_root)

        # 수치 수정
        print('image_complexity Done.')
        return (complexity >= threshold).astype(np.uint8)

    def estimate_size(self, masks):
        d = np.sqrt(np.sum(masks, axis=(1, 2, 3), dtype=np.uint32) / (masks.shape[1] * masks.shape[2]))
        # big: 2, medium: 1, small: 0
        size = np.zeros_like(d)
        size = np.where(d >= self.size_degree[1], 1, size)
        size = np.where(d >= self.size_degree[0], 2, size)

        print('estimate_size Done.')
        return size

    # 5: 전신 4: 얼굴빼고 전신 3: 상반신 2: 얼굴빼고 상반신 1: 하반신 0: 모델 없음
    def detect_body(self, X):
        # init
        num_X = len(X)
        model_type = np.zeros(num_X, dtype=np.uint8)

        # detect each image
        for i, x in enumerate(X):
            img = cv2.cvtColor(x, cv2.COLOR_RGB2GRAY)
            # full_body_detected
            full_body_detected = len(self.full_body_detector.detectMultiScale(img, 1.01, 3)) > 0
            #upper_body_detected
            upper_body_detected = len(self.upper_body_detector.detectMultiScale(img, 1.01, 3)) > 0
            #lower_body_detected
            lower_body_detected = len(self.lower_body_detector.detectMultiScale(img, 1.01, 3)) > 0

            if full_body_detected and (upper_body_detected or lower_body_detected):
                # 5: 전신(with face)
                if len(self.face_detector.detectMultiScale(img, 1.01, 3)) > 0:
                    model_type[i] = 5
                # 4: 얼굴빼고 전신
                else:
                    model_type[i] = 4
            elif upper_body_detected:
                # 3: 상반신(with face)
                if len(self.face_detector.detectMultiScale(img, 1.01, 3)) > 0:
                    model_type[i] = 3
                # 2: 얼굴빼고 상반신
                else:
                    model_type[i] = 2
            # 1: 하반신
            elif lower_body_detected:
                model_type[i] = 1
            # 0: 모델 없음
            else :
                model_type[i] = 0

        print('detect_body Done.')
        return model_type

if __name__=='__main__':
    item_dir = 'temp/'
    extractor = FeatureExtractor()
    # extract feature
    df = extractor.extract_feature(image_file=item_dir+'images.hdf5', masks_file=item_dir+'masks.hdf5')
    # save csv
    df.to_csv(item_dir + 'features.csv', sep=',', na_rep='NaN', encoding='utf-8')