#!/usr/bin/env python
u"""
plot_AIS_grid_4maps.py
Written by Tyler Sutterley (05/2023)
Creates 3 GMT-like plots of the Antarctic Ice Sheet
on a polar stereographic south (3031) projection

Grounded ice and islands from Eric Rignot grounded ice image
Glacial drainage basins updated from Rignot et al. (2012) and IMBIE-2 (2016)
2003-2004 MODIS mosaic of Antarctica (MOA) from NSIDC (Haran et al., 2013)
    https://nsidc.org/data/moa/

PYTHON DEPENDENCIES:
    numpy: Scientific Computing Tools For Python
        http://www.numpy.org
        http://www.scipy.org/NumPy_for_Matlab_Users
    scipy: Scientific Tools for Python
        http://www.scipy.org/
    netCDF4: Python interface to the netCDF C library
        https://unidata.github.io/netcdf4-python/netCDF4/index.html
    h5py: Python interface for Hierarchal Data Format 5 (HDF5)
        http://h5py.org
    matplotlib: Python 2D plotting library
        http://matplotlib.org/
        https://github.com/matplotlib/matplotlib
    cartopy: Python package designed for geospatial data processing
        https://scitools.org.uk/cartopy
    pyshp: Python read/write support for ESRI Shapefile format
        https://github.com/GeospatialPython/pyshp
    gdal: Pythonic interface to the Geospatial Data Abstraction Library (GDAL)
        https://pypi.python.org/pypi/GDAL/

UPDATE HISTORY:
    Updated 05/2023: use pathlib to define and operate on paths
        added option to set the input variable names or column order
    Updated 03/2023: switch from parameter files to argparse arguments
        updated inputs to spatial from_ascii function
    Updated 07/2022: place some imports behind try/except statements
    Updated 05/2022: use argparse descriptions within documentation
        use mask, shift grid and interpolation functions from tools
    Updated 12/2021: added color palette table (cpt) file reader from tools
    Updated 05/2021: define int/float precision to prevent deprecation warning
    Updated 10/2020: using spatial utilities for reading
        using argparse to set command-line parameters
    Forked 01/2020 from plot_AIS_grid_3maps.py
    Updated 10/2019: changing Y/N flags to True/False
    Updated 09/2019: added parameter for specifying if netCDF4 or HDF5
    Written 09/2019
"""
from __future__ import print_function

import sys
import os
import copy
import logging
import pathlib
import argparse
import warnings
import traceback
import numpy as np
import gravity_toolkit as gravtk

# attempt imports
try:
    import cartopy.crs as ccrs
except ModuleNotFoundError:
    warnings.filterwarnings("module")
    warnings.warn("cartopy not available", ImportWarning)
try:
    import matplotlib
    import matplotlib.pyplot as plt
    import matplotlib.cm as cm
    import matplotlib.colors as colors
    import matplotlib.ticker as ticker
    import matplotlib.offsetbox as offsetbox
    matplotlib.rcParams['axes.linewidth'] = 2.0
    matplotlib.rcParams['font.family'] = 'sans-serif'
    matplotlib.rcParams['font.sans-serif'] = ['Helvetica']
    matplotlib.rcParams['mathtext.default'] = 'regular'
except ModuleNotFoundError:
    warnings.filterwarnings("module")
    warnings.warn("matplotlib not available", ImportWarning)
try:
    import osgeo.gdal
except ModuleNotFoundError:
    warnings.filterwarnings("module")
    warnings.warn("GDAL not available", ImportWarning)
try:
    import shapefile
except ModuleNotFoundError:
    warnings.filterwarnings("module")
    warnings.warn("shapefile not available", ImportWarning)
# ignore warnings
warnings.filterwarnings("ignore")

# Antarctic 2012 basins
# region directory, filename, title and data type
region_dir = ['masks','Rignot_ANT']
region_title = ['AAp','ApB','BC','CCp','CpD','DDp','DpE','EEp','EpFp',
    'FpG','GH','HHp','HpI','IIpp','IppJ','JJpp','JppK','KKp','KpA']
# regional filenames
region_filename = 'basin_{0}_index.ascii'
# regional datatypes
region_dtype = {'names':('lat','lon'),'formats':('f','f')}

# IMBIE-2 Drainage basins
IMBIE_basin_file = ['masks','ANT_Basins_IMBIE2_v1.6','ANT_Basins_IMBIE2_v1.6.shp']
# basin titles within shapefile to extract
IMBIE_title = ('A-Ap','Ap-B','B-C','C-Cp','Cp-D','D-Dp','Dp-E','E-Ep','Ep-F',
    'F-Fp','F-G','G-H','H-Hp','Hp-I','I-Ipp','Ipp-J','J-Jpp','Jpp-K','K-A')

# background image mosaics
# MODIS mosaic of Antarctica
image_file = ['MOA','moa750_2004_hp1_v1.1.tif']

# Coastlines for antarctica (islands)
coast_file = ['masks','IceBoundaries_Antarctica_v02',
    'ant_ice_sheet_islands_v2.shp']

# Antarctica (AIS)
# x and y limit (modified from Bamber 1km DEM)
xlimits = np.array([-3100000,3100000])
ylimits = np.array([-2600000,2600000])
# cartopy transform for polar stereographic south
try:
    projection = ccrs.Stereographic(central_longitude=0.0,
        central_latitude=-90.0,true_scale_latitude=-71.0)
except (NameError,ValueError) as exc:
    pass

# PURPOSE: keep track of threads
def info(args):
    logging.info(pathlib.Path(sys.argv[0]).name)
    logging.info(args)
    logging.info(f'module name: {__name__}')
    if hasattr(os, 'getppid'):
        logging.info(f'parent process: {os.getppid():d}')
    logging.info(f'process id: {os.getpid():d}')

# PURPOSE: plot Rignot 2012 drainage basin polylines
def plot_rignot_basins(ax, base_dir):
    region_directory = base_dir.joinpath(*region_dir)
    # for each region
    for reg in region_title:
        # read the regional polylines
        region_file = region_directory.joinpath(region_filename.format(reg))
        region_ll = np.loadtxt(region_file, dtype=region_dtype)
        # converting region lat/lon into plot coordinates
        points = projection.transform_points(ccrs.PlateCarree(),
            region_ll['lon'], region_ll['lat'])
        ax.plot(points[:,0], points[:,1], color='k', transform=projection)

# PURPOSE: plot Antarctic drainage basins from IMBIE2 (Mouginot)
def plot_IMBIE2_basins(ax, base_dir):
    # read drainage basin polylines from shapefile (using splat operator)
    basin_shapefile = base_dir.joinpath(*IMBIE_basin_file)
    logging.debug(str(basin_shapefile))
    shape_input = shapefile.Reader(str(basin_shapefile))
    shape_entities = shape_input.shapes()
    shape_attributes = shape_input.records()
    # find record index for region by iterating through shape attributes
    # no islands or large regions
    i=[i for i,a in enumerate(shape_attributes) if a[1] in IMBIE_title]
    # for each valid shape entity
    for indice in i:
        # extract Polar-Stereographic coordinates for record
        points = np.array(shape_entities[indice].points)
        # IMBIE-2 basins can have multiple parts
        parts = shape_entities[indice].parts
        parts.append(len(points))
        for p1,p2 in zip(parts[:-1],parts[1:]):
            ax.plot(points[p1:p2,0], points[p1:p2,1], c='k',
                transform=projection)

# PURPOSE: plot Antarctic drainage sub-basins from IMBIE-2 (Mouginot)
def plot_IMBIE2_subbasins(ax, base_dir):
    # read drainage basin polylines from shapefile (using splat operator)
    IMBIE_subbasin_file = ['Basins_20Oct2016_v1.7','Basins_v1.7.shp']
    basin_shapefile = base_dir.joinpath('masks',*IMBIE_subbasin_file)
    logging.debug(str(basin_shapefile))
    shape_input = shapefile.Reader(str(basin_shapefile))
    shape_entities = shape_input.shapes()
    shape_attributes = shape_input.records()
    # iterate through shape entities and attributes
    indices = [i for i,a in enumerate(shape_attributes) if (a[1] != 'Islands')]
    for i in indices:
        # extract Polar-Stereographic coordinates for record
        points = np.array(shape_entities[i].points)
        # IMBIE-2 basins can have multiple parts
        parts = shape_entities[i].parts
        parts.append(len(points))
        for p1,p2 in zip(parts[:-1],parts[1:]):
            ax.plot(points[p1:p2,0], points[p1:p2,1], c='k',
                transform=projection)

# PURPOSE: plot Antarctic grounded ice delineation
def plot_grounded_ice(ax, base_dir, START=1):
    coast_shapefile = base_dir.joinpath(*coast_file)
    logging.debug(str(coast_shapefile))
    shape_input = shapefile.Reader(str(coast_shapefile))
    shape_entities = shape_input.shapes()
    shape_attributes = shape_input.records()
    i = [i for i,e in enumerate(shape_entities) if (np.ndim(e.points) > 1)]
    for indice in i[START:]:
        # extract Polar-Stereographic coordinates for record
        points = np.array(shape_entities[indice].points)
        ax.plot(points[:,0], points[:,1], c='k',
            transform=projection)

# PURPOSE: plot MODIS mosaic of Antarctica as background image
def plot_image_mosaic(ax, base_dir, MASKED=True):
    # read MODIS mosaic of Antarctica
    image_geotiff_file = base_dir.joinpath(*image_file)
    logging.debug(str(image_geotiff_file))
    ds = osgeo.gdal.Open(str(image_geotiff_file))
    # get dimensions
    xsize = ds.RasterXSize
    ysize = ds.RasterYSize
    # get geotiff info
    info_geotiff = ds.GetGeoTransform()
    # calculate image extents
    xmin = info_geotiff[0]
    ymax = info_geotiff[3]
    xmax = xmin + (xsize-1)*info_geotiff[1]
    ymin = ymax + (ysize-1)*info_geotiff[5]
    # read as grayscale image
    mosaic = np.ma.array(ds.ReadAsArray())
    # mask image mosaic
    if MASKED:
        # mask invalid values
        mosaic.fill_value = 0
        # create mask array for bad values
        mosaic.mask = (mosaic.data == mosaic.fill_value)
    # image extents
    extents=(xmin,xmax,ymin,ymax)
    # dataset range
    vmin, vmax = (0, 16386)
    # create color map with transparent bad points
    image_cmap = copy.copy(cm.gist_gray)
    image_cmap.set_bad(alpha=0.0)
    # nearest to not interpolate image
    im = ax.imshow(mosaic, interpolation='nearest', extent=extents,
        cmap=image_cmap, vmin=vmin, vmax=vmax, origin='upper',
        transform=projection)
    im.set_rasterized(True)
    # close the dataset
    ds = None

# PURPOSE: add a plot scale
def add_plot_scale(ax,X,Y,dx,dy,masked,fc1='w',fc2='k'):
    if masked:
        x1,x2,y1,y2 = [X-0.1*dx,X+1.2*dx,Y-2.5*dy,Y+3.2*dy]
        ax.fill([x1,x2,x2,x1,x1], [y1,y1,y2,y2,y1], fc1, zorder=4)
    for i,c in enumerate([fc1,fc2,fc1,fc2]):
        x1,x2,y1,y2 = [X+0.25*i*dx,X+0.25*(i+1)*dx,Y,Y+dy]
        ax.fill([x1,x2,x2,x1,x1], [y1,y1,y2,y2,y1], c, zorder=5)
    ax.plot([X,X+dx,X+dx,X,X], [Y,Y,Y+dy,Y+dy,Y], fc2, zorder=6)
    for i in range(3):
        ax.plot([X+0.5*i*dx,X+0.5*i*dx], [Y,Y-0.5*dy], fc2, zorder=6)
        ax.text(X+0.5*i*dx, Y-0.9*dy, '{0:0.0f}'.format(0.5*i*dx/1e3),
            ha='center', va='top', fontsize=10, color=fc2, zorder=6)
    ax.text(X+0.5*dx, Y+1.3*dy, 'km', ha='center', va='bottom',
        fontsize=10, color=fc2, zorder=6)

# plot grid program
def plot_grid(base_dir, FILENAMES,
    DATAFORM=None,
    VARIABLES=[],
    MASK=None,
    INTERPOLATION=None,
    DDEG=None,
    INTERVAL=None,
    SCALE_FACTOR=1.0,
    REMOVE_FILE=None,
    COLOR_MAP=None,
    CPT_FILE=None,
    PLOT_RANGE=None,
    BOUNDARY=None,
    ALPHA=1.0,
    CONTOURS=False,
    CONTOUR_RANGE=None,
    MEAN_CONTOUR=False,
    TITLES=None,
    LABELS=None,
    CBEXTEND=None,
    CBTITLE=None,
    CBUNITS=None,
    CBFORMAT=None,
    BASEMAP=False,
    BASIN_TYPE=None,
    DRAW_GRID_LINES=False,
    GRID=None,
    DRAW_SCALE=False,
    FIGURE_FILE=None,
    FIGURE_FORMAT=None,
    FIGURE_DPI=None,
    MODE=0o775):

    # extend list if a single format was entered for all files
    if len(DATAFORM) < len(FILENAMES):
        DATAFORM = DATAFORM*len(FILENAMES)

    # read CPT or use color map
    if CPT_FILE is not None:
        # cpt file
        cmap = gravtk.tools.from_cpt(CPT_FILE)
    else:
        # colormap
        cmap = copy.copy(cm.get_cmap(COLOR_MAP))

    # if using MODIS mosaic of Antarctica as basemap
    if BASEMAP:
        # transparent color map for bad values
        cmap.set_bad(alpha=0.0)
    else:
        # grey color map for bad values
        cmap.set_bad('w',0.5)

    # set transparency ALPHA
    if BOUNDARY is None:
        # contours
        levels = np.arange(PLOT_RANGE[0], PLOT_RANGE[1]+PLOT_RANGE[2],
            PLOT_RANGE[2])
        norm = colors.Normalize(vmin=PLOT_RANGE[0], vmax=PLOT_RANGE[1])
    else:
        # boundary between contours
        levels = np.array(BOUNDARY, dtype=np.float64)
        norm = colors.BoundaryNorm(BOUNDARY, ncolors=256)

    # convert degree spacing and interval parameters
    # Grid spacing
    dlon,dlat = (DDEG[0],DDEG[0]) if (len(DDEG) == 1) else (DDEG[0],DDEG[1])
    # Grid dimensions
    if (INTERVAL == 1):# (0:360, 90:-90)
        nlon = np.int64((360.0/dlon)+1.0)
        nlat = np.int64((180.0/dlat)+1.0)
    elif (INTERVAL == 2):# degree spacing/2
        nlon = np.int64((360.0/dlon))
        nlat = np.int64((180.0/dlat))

    # interpolation method for image background using transform_scalar
    if (INTERPOLATION == 'nearest'):
        order = 0
    elif (INTERPOLATION == 'bilinear'):
        order = 1
    elif (INTERPOLATION == 'cubic'):
        order = 3

    # create masked array if missing values
    if MASK is not None:
        # Read Land-Sea Mask of specified input file
        # 0=Ocean, 1=Land, 2=Lake, 3=Small Island, 4=Ice Shelf
        # Open the land-sea NetCDF file for reading
        landsea = gravtk.spatial().from_netCDF4(MASK,
            date=False, varname='LSMASK')
        # create land function
        nth,nphi = landsea.shape
        mask = np.zeros((nth,nphi),dtype=bool)
        # combine land and island levels for land function
        indx,indy = np.nonzero((landsea.data >= 1) & (landsea.data <= 3))
        mask[indx,indy] = True

    # remove a spatial field from each input map
    if REMOVE_FILE is not None:
        REMOVE = gravtk.spatial().from_netCDF4(REMOVE_FILE,
            date=False).data[:,:]
    else:
        REMOVE = 0.0

    # image extents
    ax = {}
    # setup polar stereographic maps
    fig, ((ax[0],ax[1]),(ax[2],ax[3])) = plt.subplots(num=1, nrows=2, ncols=2,
        figsize=(6,6.3), subplot_kw=dict(projection=projection))
    # WGS84 Ellipsoid parameters
    a_axis = 6378137.0# [m] semimajor axis of the ellipsoid
    flat = 1.0/298.257223563# flattening of the ellipsoid
    # (4pi/3)R^3 = (4pi/3)(a^2)b = (4pi/3)(a^3)(1 -f)
    rad_e = a_axis*(1.0 - flat)**(1.0/3.0)

    for i,ax1 in ax.items():
       # plot image of MODIS mosaic of Antarctica as base layer
        if BASEMAP:
            # plot MODIS mosaic of Antarctica
            plot_image_mosaic(ax1, base_dir)

        # input ascii/netCDF4/HDF5 file
        if (DATAFORM[i] == 'ascii'):
            # ascii (.txt)
            dinput = gravtk.spatial().from_ascii(FILENAMES[i], date=False,
                columns=VARIABLES, spacing=[dlon,dlat], nlat=nlat, nlon=nlon)
        elif (DATAFORM[i] == 'netCDF4'):
            # netCDF4 (.nc)
            field_mapping = gravtk.spatial().default_field_mapping(VARIABLES)
            dinput = gravtk.spatial().from_netCDF4(FILENAMES[i], date=False,
                field_mapping=field_mapping)
        elif (DATAFORM[i] == 'HDF5'):
            # HDF5 (.H5)
            field_mapping = gravtk.spatial().default_field_mapping(VARIABLES)
            dinput = gravtk.spatial().from_HDF5(FILENAMES[i], date=False,
                field_mapping=field_mapping)

        # remove offset and scale to units
        if (REMOVE != 0.0) or (SCALE_FACTOR != 1.0):
            dinput = dinput.offset(-REMOVE).scale(SCALE_FACTOR)

        # update masked values
        if MASK is not None:
            dinput.replace_invalid(fill_value=dinput.fill_value, mask=mask)

        # if dlat is negative
        if (np.sign(dlat) == -1):
            dinput = dinput.flip(axis=0)

        # calculate image coordinates
        mx = np.int64((xlimits[1]-xlimits[0])/1000.)+1
        my = np.int64((ylimits[1]-ylimits[0])/1000.)+1
        X = np.linspace(xlimits[0],xlimits[1],mx)
        Y = np.linspace(ylimits[0],ylimits[1],my)
        gridx,gridy = np.meshgrid(X,Y)
        # create mesh lon/lat
        points = ccrs.PlateCarree().transform_points(projection,
            gridx.flatten(), gridy.flatten())
        lonsin = points[:,0].reshape(my,mx)
        latsin = points[:,1].reshape(my,mx)

        # interpolate to image coordinates
        if (INTERVAL == 1) and (np.max(dinput.lon) > 180):# (0:360, 90:-90)
            shift_data,lon180 = gravtk.tools.shift_grid(180.0,
                dinput.data,dinput.lon)
            shift_mask,lon180 = gravtk.tools.shift_grid(180.0,
                dinput.mask,dinput.lon)
            img = gravtk.tools.interp_grid(shift_data,lon180,
                dinput.lat,lonsin,latsin,order)
            msk = gravtk.tools.interp_grid(shift_mask,lon180,
                dinput.lat,lonsin,latsin,order)
        elif (INTERVAL == 2) and (np.max(dinput.lon) > 180):# DDEG/2
            shift_data,lon180 = gravtk.tools.shift_grid(180.0+dlon,
                dinput.data,dinput.lon)
            shift_mask,lon180 = gravtk.tools.shift_grid(180.0+dlon,
                dinput.mask,dinput.lon)
            img = gravtk.tools.interp_grid(shift_data,
                lon180,dinput.lat,lonsin,latsin,order)
            msk = gravtk.tools.interp_grid(shift_mask,
                lon180,dinput.lat,lonsin,latsin,order)
        else:# -180:180 or modification of there of
            img = gravtk.tools.interp_grid(dinput.data,
                dinput.lon,dinput.lat,lonsin,latsin,order)
            msk = gravtk.tools.interp_grid(dinput.mask,
                dinput.lon,dinput.lat,lonsin,latsin,order)
        # create masked array of image
        img = np.ma.array(img, mask=msk.astype(bool))

        # plot only grounded points
        if MASK is not None:
            img = gravtk.tools.mask_oceans(lonsin,latsin,
                data=img,order=order,iceshelves=False)
        # plot image with transparency using normalization
        im = ax1.imshow(img, interpolation='nearest', cmap=cmap,
            extent=(xlimits[0],xlimits[1],ylimits[0],ylimits[1]),
            norm=norm, alpha=ALPHA, origin='lower', transform=projection)

        # create mesh lon/lat
        lon, lat = np.meshgrid(dinput.lon,dinput.lat)
        data = dinput.to_masked_array()
        # plot line contours
        if CONTOURS:
            clevs = np.arange(CONTOUR_RANGE[0],
                CONTOUR_RANGE[1] + CONTOUR_RANGE[2],
                CONTOUR_RANGE[2])
            # remove 0 (will plot in red)
            reduce_clevs = clevs[np.nonzero(clevs)]
            # plot contours
            ax1.contour(lon,lat,data,reduce_clevs,colors='0.2',
                linestyles='solid',transform=ccrs.PlateCarree())
            ax1.contour(lon,lat,data,[0],colors='red',linestyles='solid',
                linewidths=1.5,transform=ccrs.PlateCarree())

        # plot line contour for global average
        if MEAN_CONTOUR and CONTOURS:
            # calculate areas of each grid cell
            dphi,dth = (dlon*np.pi/180.0,dlat*np.pi/180.0)
            indy,indx = np.nonzero(np.logical_not(data.mask))
            area = (rad_e**2)*dth*dphi*np.cos(lat[indy,indx]*np.pi/180.0)
            # calculate average
            ave=np.sum(area*data[indy,indx])/np.sum(area)
            # plot line contour of global average
            ax1.contour(lon,lat,data,[ave],colors='blue',linestyles='solid',
                linewidths=1.5,transform=ccrs.PlateCarree())

        # add basins based on BASIN_TYPE (Rignot 2012, IMBIE-2, IMBIE-2 subbasins)
        if (BASIN_TYPE == 'Rignot'):
            plot_rignot_basins(ax1, base_dir)
            start_indice = 1
        elif (BASIN_TYPE == 'IMBIE-2'):
            plot_IMBIE2_basins(ax1, base_dir)
            start_indice = 1
        elif (BASIN_TYPE == 'IMBIE-2_subbasin'):
            plot_IMBIE2_subbasins(ax1, base_dir)
            start_indice = 1
        else:
            start_indice = 0
        # add glaciers/ice caps/islands
        plot_grounded_ice(ax1, base_dir, START=start_indice)

        # draw lat/lon grid lines
        if DRAW_GRID_LINES:
            # meridian and parallel grid spacing
            llx,lly = (GRID[0],GRID[0]) if (len(GRID) == 1) else (GRID[0],GRID[1])
            grid_meridians = np.arange(-180, 180 + llx, llx)
            grid_parallels = np.arange(-90, 90 + lly, lly)
            gl = ax1.gridlines(crs=ccrs.PlateCarree(), draw_labels=False,
                linewidth=0.1, color='0.25', linestyle='-')
            gl.xlocator = ticker.FixedLocator(grid_meridians)
            gl.ylocator = ticker.FixedLocator(grid_parallels)

        # add title for each subplot
        if TITLES is not None:
            TITLE = ' '.join(TITLES[i].split('_'))
            ax1.set_title(TITLE.replace('-',u'\u2013'), fontsize=14)
            ax1.title.set_y(0.995)
        # Add figure label for each subplot
        if LABELS is not None:
            at = offsetbox.AnchoredText(LABELS[i],
                loc=2, pad=0, borderpad=0.25, frameon=True,
                prop=dict(size=18,weight='bold',color='k'))
            at.patch.set_boxstyle("Square,pad=0.1")
            at.patch.set_edgecolor("white")
            ax1.axes.add_artist(at)

        # x and y limits, axis = equal
        ax1.set_xlim(xlimits)
        ax1.set_ylim(ylimits)
        ax1.set_aspect('equal', adjustable='box')
        # no ticks on the x and y axes
        ax1.get_xaxis().set_ticks([])
        ax1.get_yaxis().set_ticks([])
        # stronger linewidth on frame
        ax1.spines['geo'].set_linewidth(2.0)
        ax1.spines['geo'].set_zorder(10)
        ax1.spines['geo'].set_capstyle('projecting')

    # draw map scale to corners of axis
    if DRAW_SCALE:
        add_plot_scale(ax[2],-292e4,-215e4,1600e3,140e3,False)

    # Add colorbar
    # Add an axes at position rect [left, bottom, width, height]
    cbar_ax = fig.add_axes([0.085, 0.095, 0.83, 0.035])
    # extend = add extension triangles to upper and lower bounds
    # options: neither, both, min, max
    cbar = fig.colorbar(im, cax=cbar_ax, extend=CBEXTEND,
        extendfrac=0.0375, drawedges=False, orientation='horizontal')
    # rasterized colorbar to remove lines
    cbar.solids.set_rasterized(True)
    # Add label to the colorbar
    cbar.ax.set_title(CBTITLE, fontsize=14, rotation=0, y=-1.65, va='top')
    cbar.ax.set_xlabel(CBUNITS, fontsize=14, rotation=0, va='center')
    cbar.ax.xaxis.set_label_coords(1.08, 0.5)
    # Set the tick levels for the colorbar
    cbar.set_ticks(levels)
    cbar.set_ticklabels([CBFORMAT.format(ct) for ct in levels])
    # ticks lines all the way across
    cbar.ax.tick_params(which='both', width=1, length=16, labelsize=14,
        direction='in')

    # adjust subplots within figure
    fig.subplots_adjust(left=0.01, right=0.99, bottom=0.14, top=0.97,
        hspace=0.05, wspace=0.05)
    # create output directory if non-existent
    FIGURE_FILE.parent.mkdir(mode=MODE, parents=True, exist_ok=True)
    # save to file
    logging.info(str(FIGURE_FILE))
    plt.savefig(FIGURE_FILE,
        metadata={'Title':pathlib.Path(sys.argv[0]).name},
        dpi=FIGURE_DPI, format=FIGURE_FORMAT)
    plt.clf()
    # change the permissions mode
    FIGURE_FILE.chmod(mode=MODE)

# PURPOSE: create argument parser
def arguments():
    parser = argparse.ArgumentParser(
        description="""Creates 4 GMT-like plots of the Antarctic ice sheet
            on a polar stereographic south (EPSG 3031) projection
            """,
        fromfile_prefix_chars="@"
    )
    parser.convert_arg_line_to_args = gravtk.utilities.convert_arg_line_to_args
    # command line parameters
    parser.add_argument('infile', nargs=4,
        type=pathlib.Path,
        help='Input grid files')
    # working data directory
    parser.add_argument('--directory','-D',
        type=pathlib.Path, default=pathlib.Path.cwd(),
        help='Working data directory')
    # Input data format (ascii, netCDF4, HDF5)
    parser.add_argument('--format','-F',
        type=str, nargs='+',
        default='netCDF4', choices=['ascii','netCDF4','HDF5'],
        help='Input data format')
    # variable names (for ascii names of columns)
    parser.add_argument('--variables','-v',
        type=str, nargs='+', default=['lon','lat','z'],
        help='Variable names of data in input file')
    # land-sea mask
    lsmask = gravtk.utilities.get_data_path(['data','landsea_hd.nc'])
    parser.add_argument('--mask',
        type=pathlib.Path, default=lsmask,
        help='Land-sea mask')
    # output grid parameters
    parser.add_argument('--spacing',
        type=float, nargs='+', default=[0.5,0.5], metavar=('dlon','dlat'),
        help='Spatial resolution of input data')
    parser.add_argument('--interval',
        type=int, default=2, choices=[1,2],
        help=('Input grid interval (1: global, 2: centered global)'))
    # Interpolation method
    parser.add_argument('--interpolation','-I',
        type=str, default='bilinear', choices=['nearest','bilinear','cubic'],
        help='Interpolation method')
    # scale factor
    parser.add_argument('--scale-factor','-s',
        type=float, default=1.0,
        help='Multiplicative scale factor for converting to plot units')
    # plot range
    parser.add_argument('--plot-range','-R',
        type=float, nargs=3, metavar=('MIN','MAX','STEP'),
        help='Plot range and step size for normalization')
    parser.add_argument('--boundary','-B',
        type=float, nargs='+',
        help='Plot boundary for normalization')
    # color palette table or named color map
    try:
        cmap_set = set(cm.datad.keys()) | set(cm.cmaps_listed.keys())
    except (ValueError, NameError) as exc:
        cmap_set = []
    parser.add_argument('--colormap','-m',
        metavar='COLORMAP', type=str, default='viridis',
        choices=sorted(cmap_set),
        help='Named Matplotlib colormap')
    parser.add_argument('--cpt-file','-c',
        type=pathlib.Path,
        help='Input Color Palette Table (.cpt) file')
    # color map alpha
    parser.add_argument('--alpha','-a',
        type=float, default=1.0,
        help='Named Matplotlib colormap')
    # plot contour parameters
    parser.add_argument('--plot-contours',
        default=False, action='store_true',
        help='Plot contours')
    parser.add_argument('--contour-range',
        type=float, nargs=3, metavar=('MIN','MAX','STEP'),
        help='Contour range and step size')
    parser.add_argument('--mean-contour',
        default=False, action='store_true',
        help='Plot contours for mean of dataset')
    # title and label
    parser.add_argument('--plot-title', nargs=4,
        type=str, help='Plot title')
    parser.add_argument('--plot-label', nargs=4,
        type=str, help='Plot label')
    # colorbar parameters
    parser.add_argument('--cbextend',
        type=str, default='both',
        choices=['neither', 'both', 'min', 'max'],
        help='Add extension triangles to colorbar')
    parser.add_argument('--cbtitle',
        type=str, default='',
        help='Title label for colorbar')
    parser.add_argument('--cbunits',
        type=str, default='',
        help='Units label for colorbar')
    parser.add_argument('--cbformat',
        type=str, default='{0:3.0f}',
        help='Tick format for colorbar')
    # additional parameters
    parser.add_argument('--basemap',
        default=False, action='store_true',
        help='Add background basemap image')
    parser.add_argument('--basin-type',
        type=str, default='',
        help='Add delineations for glacier drainage basins')
    parser.add_argument('--draw-grid-lines',
        default=False, action='store_true',
        help='Add map grid lines')
    parser.add_argument('--grid-lines',
        type=float, nargs='+', default=(15,15),
        help='Input grid spacing for meridians and parallels')
    parser.add_argument('--draw-scale',
        default=False, action='store_true',
        help='Add map scale bar')
    # output file, format and dpi
    parser.add_argument('--figure-file','-O',
        type=pathlib.Path,
        help='Output figure file')
    parser.add_argument('--figure-format','-f',
        type=str, default='png', choices=('pdf','png','jpg','svg'),
        help='Output figure format')
    parser.add_argument('--figure-dpi','-d',
        type=int, default=180,
        help='Output figure resolution in dots per inch (dpi)')
    # print information about each input and output file
    parser.add_argument('--verbose','-V',
        action='count', default=0,
        help='Verbose output of run')
    # permissions mode of the local directories and files (number in octal)
    parser.add_argument('--mode','-M',
        type=lambda x: int(x,base=8), default=0o775,
        help='Permissions mode of output files')
    # return the parser
    return parser

# This is the main part of the program that calls the individual functions
def main():
    # Read the system arguments listed after the program
    parser = arguments()
    args,_ = parser.parse_known_args()

    # create logger
    loglevels = [logging.CRITICAL, logging.INFO, logging.DEBUG]
    logging.basicConfig(level=loglevels[args.verbose])

    # try to run the analysis with listed parameters
    try:
        info(args)
        # run plot program with parameters
        plot_grid(args.directory, args.infile,
            DATAFORM=args.format,
            VARIABLES=args.variables,
            DDEG=args.spacing,
            INTERVAL=args.interval,
            INTERPOLATION=args.interpolation,
            SCALE_FACTOR=args.scale_factor,
            CPT_FILE=args.cpt_file,
            COLOR_MAP=args.colormap,
            PLOT_RANGE=args.plot_range,
            BOUNDARY=args.boundary,
            ALPHA=args.alpha,
            CONTOURS=args.plot_contours,
            CONTOUR_RANGE=args.contour_range,
            MEAN_CONTOUR=args.mean_contour,
            TITLES=args.plot_title,
            LABELS=args.plot_label,
            CBEXTEND=args.cbextend,
            CBTITLE=args.cbtitle,
            CBUNITS=args.cbunits,
            CBFORMAT=args.cbformat,
            BASEMAP=args.basemap,
            BASIN_TYPE=args.basin_type,
            DRAW_GRID_LINES=args.draw_grid_lines,
            GRID=args.grid_lines,
            DRAW_SCALE=args.draw_scale,
            FIGURE_FILE=args.figure_file,
            FIGURE_FORMAT=args.figure_format,
            FIGURE_DPI=args.figure_dpi,
            MODE=args.mode)
    except Exception as exc:
        # if there has been an error exception
        # print the type, value, and stack trace of the
        # current exception being handled
        logging.critical(f'process id {os.getpid():d} failed')
        logging.error(traceback.format_exc())

# run main program
if __name__ == '__main__':
    main()


from neurolight.match.costs import get_costs
from neurolight.match.preprocess import mouselight_preprocessing
from funlib.match.helper_functions import match

import networkx as nx

from obj_to_npz import parse_npy_graph
import pickle

from pathlib import Path
import copy

penalty_attr = "penalty"
location_attr = "location"

trimmed = Path.cwd() / "invalid_trimmed"
trimmed.mkdir(parents=True, exist_ok=True)

for i, valid in enumerate(Path("invalid").iterdir()):
    print(valid.name)
    graph = pickle.load((valid / "graph.obj").open("rb"))
    tree = pickle.load((valid / "tree.obj").open("rb"))

    temp = copy.deepcopy(graph)

    max_edge_len = 48
    mouselight_preprocessing(
        temp,
        max_dist=150,
        voxel_size=[10, 3, 3],
        penalty_attr=penalty_attr,
        location_attr=location_attr,
    )

    node_match_threshold = 76
    edge_match_threshold = 76
    node_balance = 10

    node_costs, edge_costs = get_costs(
        temp,
        tree,
        location_attr=location_attr,
        penalty_attr=penalty_attr,
        node_match_threshold=100,
        edge_match_threshold=100,
        node_balance=node_balance,
    )

    try:
        matched = match(temp, tree, node_costs, edge_costs)
    except:
        print("failed to match!")
        matched = graph

    untouched_ccs = list(set(x) for x in nx.connected_components(graph))

    for node in matched.nodes():
        for cc in untouched_ccs:
            if node in cc:
                untouched_ccs.remove(cc)

    (trimmed / f"{i:03}").mkdir(exist_ok=True, parents=True)
    for untouched_cc in untouched_ccs:
        graph.remove_nodes_from(untouched_cc)
    pickle.dump(graph, open(f"invalid_trimmed/{i:03}/graph.obj", "wb"))
    pickle.dump(tree, open(f"invalid_trimmed/{i:03}/tree.obj", "wb"))


import uuid
from flask import Flask, request, jsonify
from marshmallow import Schema, fields, ValidationError
from celery import Celery
from celery.result import AsyncResult

class PayloadSchema(Schema):
    fab = fields.Str()
    area = fields.Str()
    tool_name = fields.Str()
    start_date = fields.Str()
    end_date = fields.Str()

app = Flask(__name__)
app.config['CELERY_BROKER_URL'] = 'redis://localhost:5000'
app.config['CELERY_RESULT_BACKEND'] = 'redis://localhost:5555'
celery = Celery(
    app.name,
    broker=app.config['CELERY_BROKER_URL'],
    backend=app.config['CELERY_RESULT_BACKEND'])

@celery.task(bind=True, name="task.tool_health_spc_query", max_retries=3)
def tool_health_spc_query(self, final_payload):
    # SPC query code...
    data = 'results from SPC query'
    print('hello world')

    return {
        'status': 'task completed',
        'data': data
    }

@celery.task(bind=True, name="task.tool_health_fd_query", max_retries=3)
def tool_health_fd_query(self, final_payload):
    # FD query code...
    data = 'results from FD query'
    print('hello world')

    return {
        'status': 'task completed',
        'data': data
    }

class ToolHealthFactory:
    def __init__(self, payload):
        self.payload = payload.copy()   # payload is the request in dictionary

    def queue_task(self):
        pass

class SPCQuery(ToolHealthFactory):
    def __init__(self, payload, spc_db, spc_query):
        self.spc_db = spc_db
        self.spc_query = spc_query
        super().__init(payload)

    def queue_task(self):
        # add additional key value pairs to payload dictionary

        self.payload['spc_db'] = self.spc_db
        self.payload['spc_query'] = self.spc_query

        # add task to queue
        task = tool_health_spc_query.delay(self.payload)
        return {'id': task.id}

class FDQuery(ToolHealthFactory):
    def __init__(self, payload, endpoint, param):
        self.endpoint = endpoint
        self.param = param
        super().__init(payload)

    def queue_task(self):
        # add additional key value pairs to payload dictionary

        self.payload['endpoint'] = self.endpoint
        self.payload['param'] = self.param

        # add task to queue
        task = tool_health_fd_query.delay(self.payload)
        return {'id': task.id}

@app.route('/toolhealth', methods=['POST', 'GET'])
def tool_health():
    payload_json = request.body()
    try:
        payload = PayloadSchema().load(payload_json)

    except ValidationError as e:
        return jsonify({'error': e}), 400

    # create SPC and FD objects
    spc_query = SPCQuery(payload, 'spc db', 'spc query')
    fd_query = FDQuery(payload, 'endpoint', 'param')

    # pushing tasks to queue
    spc_task = spc_query.queue_task()
    fd_task = fd_query.queue_task()

    ids_dict = {
        'tool_health_id': uuid.uuid1(),
        'task_ids': {
            'spc_query_id': spc_task['id'],
            'fd_task_id': fd_task['id']
        }
    }

    return jsonify({
        'message': 'queue tasks successful',
        'ids': ids_dict
        }), 202

@app.route('/tool_health/<tool_health_id>', methods=['POST', 'GET'])
def tool_health_results(tool_health_id):
    payload = request.get_json()    # request body is ids_dict

    if payload['tool_health_id'] == tool_health_id:
        spc_result = AsyncResult(payload['task_ids']['spc_query_id'])
        fd_result = AsyncResult(payload['task_ids']['fd_query_id'])
        return jsonify({
            'spc status': spc_result.status,
            'spc results': spc_result.get(),
            'fd status': fd_result.status,
            'fd results': fd_result.get(),
            })

    return jsonify({
        'message': 'tool_health_id not found. Please try again.'
    })


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

import sys
import os
import argparse
from workflow.util.useful import const
from workflow.util.useful import parse_group_file,mkdir,Rparser,get_name,image_trans,cutcol_dataFrame

def read_params(args):
    parser = argparse.ArgumentParser(description='a wrapper for LEfSe | v1.0 at 2017/1/23 by huangy')
    parser.add_argument('-i', '--summarize_all', dest='infile', metavar='FILE', type=str, required=True,
                        help="set hte otu_table_all.txt produced by 02_summarize_trans.py")
    parser.add_argument('-l', '--LEfSe_path', dest='LEfSe_path', metavar='DIR', type=str, default=None,
                        help="set the LEfSe path, default find in env")
    parser.add_argument('-g', '--group_file', dest='group', metavar='FILE', type=str, required=True,
                        help="set the group file")
    parser.add_argument('-o', '--out_dir', dest='out_dir', metavar='DIR', type=str, required=True,
                        help="set the output dir")
    parser.add_argument('--LDA', dest='LDA', metavar='FLOAT', type=float, default=2,
                        help="set the LDA cutoff, [default is 2]")
    parser.add_argument('--class_row', dest='c', metavar='INT', type=int, default=1,
                        help="set the class row, [default is 1]")
    parser.add_argument('--subject_row', dest='u', metavar='INT', type=int, default=2,
                        help="set the subject row, [default is 2]")
    parser.add_argument('--quantile',dest="quantile",metavar='FLOAT',type=float,default=0)
    parser.add_argument('--cut_off',dest='cut_off',metavar='FLOAT',type=float,default=1e-20)
    args = parser.parse_args()
    params = vars(args)
    params['groupDir'] = params['group']
    params['group'] = parse_group_file(params['group'])
    if params['LEfSe_path'] is None:
        params['LEfSe_path'] = ''
    else:
        params['LEfSe_path'] += '/'
    return params

def do_format(infile, outfile, groupDir,quantile,cut_off):
    data,group = cutcol_dataFrame(infile,groupDir)
    groups=group.values()
    samples = group.keys()

    data.index=[i.replace('[','') for i in data.index.tolist()]
    data.index=[i.replace(']','') for i in data.index.tolist()]
    for temp in data.index:
        if temp.endswith('Other'):
            data.drop(temp)
        if temp.endswith('norank'):
            data.drop(temp)
        if temp.endswith('unclassfied'):
            data.drop(temp)
    sample_num = len(data.columns)
    data["sum"] = data.sum(axis=1)
    data = data.sort_values(by="sum",ascending=False)
    quantile_value = float(data["sum"].quantile(quantile))
    dd = data[data>cut_off]
    data["num_True"] = dd.count(axis=1)
    index_list = []
    for i in range(len(data.index)):
        if data.ix[i,sample_num] < quantile_value:
            if data.ix[i,sample_num+1] <2:
                index_list.append(i)
    data = data.drop(data.index[index_list])
    del data["sum"]
    del data["num_True"]
    with open(outfile, 'w') as out_fp:
        out_fp.write('class\t%s\n' % '\t'.join(groups))
        out_fp.write('Taxon\t%s\n' % '\t'.join(samples))
        data.to_csv(out_fp,header=False,sep="\t")

def get_commands(infile, LEfSe_path, out_dir, LDA, c, u):
    commands = []
    command = '%sformat_input.py %s %s/LDA.in -c %s -u %s -o 1000000' % (LEfSe_path, infile, out_dir, c, u)
    commands.append(command)
    command = '%srun_lefse.py %s/LDA.in %s/LDA.res -l %s' % (LEfSe_path, out_dir, out_dir, LDA)
    commands.append(command)
    command = '%splot_res.py %s/LDA.res %s/LDA.pdf --format pdf --dpi 150' % (LEfSe_path, out_dir, out_dir)
    commands.append(command)
    command = '%splot_res.py %s/LDA.res %s/LDA.png --format png --dpi 150' % (LEfSe_path, out_dir, out_dir)
    commands.append(command)
#    command = '%splot_cladogram.py %s/LDA.res %s/LDA.cladogram.pdf --format pdf --dpi 150' % (
 #       LEfSe_path, out_dir, out_dir)
 #   commands.append(command)
 #   command = '%splot_cladogram.py %s/LDA.res %s/LDA.cladogram.png --format png --dpi 150' % (
     #   LEfSe_path, out_dir, out_dir)
   # commands.append(command)
  #  if not os.path.isdir('%s/biomarkers_raw_images' % out_dir):
   #     os.mkdir('%s/biomarkers_raw_images' % out_dir)
   # command = '%(LEfSe_path)splot_features.py %(out_dir)s/LDA.in %(out_dir)s/LDA.res %(out_dir)s/biomarkers_raw_images/ --format pdf --dpi 200' % {
    #    'LEfSe_path': LEfSe_path, 'out_dir': out_dir}
    #commands.append(command)
    return commands

if __name__ == '__main__':
    params = read_params(sys.argv)
    mkdir(params['out_dir'])
    for_analysis = '%s/otu_table_for_lefse.txt' % params['out_dir']
    do_format(params['infile'], for_analysis, params['groupDir'],params['quantile'],params['cut_off'])
    commands = get_commands(for_analysis, params['LEfSe_path'], params['out_dir'],
                            params['LDA'], params['c'], params['u'])
    with open(params['out_dir'] + '/commands.sh', 'w') as fp:
        fp.write('\n'.join(commands))
    for command in commands:
        os.system(command)
    os.system("awk '$4>%s' %s/LDA.res >%s/diff.marker.filter.txt.tmp"%(params['LDA'],params['out_dir'],params['out_dir']))
    with open("%s/diff.marker.filter.txt.tmp"%params['out_dir']) as fqin,open("%s/diff.marker.filter.tsv"%params['out_dir'],"w") as fqout:
        for line in fqin:
            tabs = line.strip().split("\t")
            fqout.write("%s\t%s\t\t%s\t\t%s\n"%(tabs[0],tabs[1],tabs[4],tabs[2]))
    os.system("rm %s/diff.marker.filter.txt.tmp"%(params['out_dir']))
    os.system("cut -f 1 %s/diff.marker.filter.tsv |while read line;do echo \"grep \'$line\' %s >> %s/diff.marker.filter.profile.tsv\";done>%s/do.sh"%(params['out_dir'],params['infile'],params['out_dir'],params['out_dir']))
    os.system("head -1 %s > %s/diff.marker.filter.profile.tsv"%(params['infile'],params['out_dir']))
    os.system("sh %s/do.sh"%params['out_dir'])
    os.system("rm %s/do.sh"%params['out_dir'])


from sqlalchemy.orm import sessionmaker
from sqlalchemy import create_engine

from models import PublicTeam, SkillList, Mentors

import settings

engine = create_engine(settings.DB_NAME)

Session = sessionmaker(bind=engine)

session = Session()

def insert_skill(name):
    s = SkillList(name=name)
    session.add(s)
    session.commit()

def get_skills():
    skillies = [s.name for s in session.query(SkillList).all()]
    return skillies

def get_skill(skill_name):
    skill = session.query(SkillList).filter(skill_name == SkillList.name).one()
    return skill

def get_mentor(mentor_name):
    mentor = session.query(Mentors).filter(mentor_name == Mentors.name).one()
    return mentor

def get_team(team_name):
    team = session.query(PublicTeam).filter(team_name == PublicTeam.name).one()
    return team

def get_all_teams():
    teams = [(team.name, team.idea_description) for team in session.query(
        PublicTeam).all()]
    return teams

def insert_team(name, idea_desc, repo, need_members, members_needed_desc,
                room, techs, place="null"):
    team = PublicTeam(name=name, idea_description=idea_desc, repository=repo,
                      need_more_members=need_members,
                      members_needed_desc=members_needed_desc, room=room,
                      place=place)
    for tech in techs:
        skill = get_skill(tech['name'])
        team.skills.append(skill)
    session.add(team)
    session.commit()

def insert_mentor(name, description, picture, teams):
    mentor = Mentors(name=name, description=description, picture=picture)
    for team in teams:
        t = get_team(team['name'])
        mentor.teams.append(t)
    session.add(mentor)
    session.commit()

def get_teams_in_room(room):
    teams = session.query(PublicTeam).filter(room == PublicTeam.room).count()
    return teams

def get_teams_for_technology(tech):
    skill = get_skill(tech)
    return [t.name for t in skill.teams]

def get_mentor_teams(mentor_name):
    mentor = get_mentor(mentor_name)
    for team in mentor.teams:
        print("name: {}, \nidea_description: {}, \nrepository: {}, \
              \nmore members?: {}, \nwhat members needed?: {},\
              \nroom: {}, \nskills: {}".format(team.name,
                                               team.idea_description,
                                               team.repository,
                                               team.need_more_members,
                                               team.members_needed_desc,
                                               team.room,
                                               [skill.name for skill in
                                                team.skills]))

def get_mentor_rooms(mentor_name):
    mentor = get_mentor(mentor_name)
    rooms = []
    for team in mentor.teams:
        rooms.append(team.room)

    sorted_rooms = sorted([int(room) for room in rooms if room != 'mazata'])
    if 'mazata' in rooms:
        sorted_rooms.append('mazata')

    return sorted_rooms

def remove_skill(team_name, skill_name):
    team = get_team(team_name)
    for skill in team.skills:
        if skill.name == skill_name:
            team.skills.remove(skill)

def add_skill_to_team(team_name, skills):
    team = get_team(team_name)
    for skill in skills:
        s = get_skill(skill)
        if s not in team.skills:
            team.skills.append(s)

def update_team_attr(team_name, attr, new_value):
    team = get_team(team_name)
    setattr(team, attr, new_value)


from sklearn.neighbors import KNeighborsClassifier
from sklearn.neighbors import NearestNeighbors
from csv import reader
import sys
import numpy as np
from pyspark import SparkContext
from csv import reader
from operator import add
import pandas as pd
import heapq
from pyspark.ml.feature import OneHotEncoder, StringIndexer
from pyspark.sql.types import BooleanType
from pyspark.sql.functions import udf,col,when
from pyspark.sql.types import IntegerType,StringType
from pyspark.sql.functions import isnan

data = sc.textFile("/FileStore/tables/h9gi_nx95-fba74.tsv")
data = data.mapPartitions(lambda x: reader(x, delimiter='\t'))
header = data.first()
data = data.filter(lambda row : row != header)

def isfloat(value):
  try:
    float(value)
    return True
  except ValueError:
    return False
def dofloat(entry):
  k = []
  for i in range(len(entry)):
    if isfloat(entry[i]) != True:
      k.append(i)
  return k
df = spark.createDataFrame(data, header)

m = df.count()
for i in range(len(header)):
  k = df.filter((df[header[i]] == "") | df[header[i]].isNull() | isnan(df[header[i]])).count()
  null_percent = float(k)/m
  if null_percent > 0.1:
    df = df.drop(header[i])
    continue

  categories = df.select(header[i]).sample(withReplacement=False, fraction=0.1).distinct().replace('', '0').rdd.flatMap(lambda x : x).collect()
  tem = dofloat(categories)
  if len(tem) == 0:
    continue
  if len(tem) >10:
    df = df.drop(header[i])
    continue
  categories = df.select(header[i]).distinct().rdd.flatMap(lambda x : x).collect()
  for category in categories:
      function = udf(lambda item: 1 if item == category else 0, StringType())
      new_column_name = header[i]+'_'+category
      df = df.withColumn(new_column_name, function(col(header[i])))
  df = df.drop(header[i])

df = df.replace('', '0')
data = df.rdd.map(list)
data_collect = df.rdd.map(list).collect()

k_neighbors = 25
knnobj = NearestNeighbors(n_neighbors = k_neighbors).fit(data_collect)
bc_knnobj = sc.broadcast(knnobj)
results = data.map(lambda x: bc_knnobj.value.kneighbors(x))

def cleanresult(entry,k):
  k_neighbor = entry[0][0][k]
  neighbors = entry[1][0]
  neighboures_distance = entry[0][0]
  return [k_neighbor, neighbors,neighboures_distance]

kdistanceres= results.map(lambda x:cleanresult(x,k_neighbors-1)).zipWithIndex().collect()

def add_kdiatanceofO(kdistanceres):
  for i in range(len(kdistanceres)):
    kdiatanceofO = []
    for j in kdistanceres[i][0][1]:
      kdiatanceofO.append(kdistanceres[j][0][0])
    kdistanceres[i][0].append(kdiatanceofO)

add_kdiatanceofO(kdistanceres)

kdistanceres_rdd = sc.parallelize(kdistanceres)
def reach_distance(d1,d2):
  return max(d1,d2)

def lrd(entry,k):
  sum_reach_distance = 0
  for i in range(1,k):
    sum_reach_distance += reach_distance(entry[0][2][i], entry[0][3][i])
  return ([k/sum_reach_distance,entry[0][1] ], entry[1])

lrd_rdd = kdistanceres_rdd.map(lambda x:lrd(x,k_neighbors))
lrd = lrd_rdd.collect()

def add_lrdOfO(lrd):
  for i in range(len(lrd)):
    sum_lrdOfO = 0
    for j in lrd[i][0][1]:
      sum_lrdOfO += lrd[j][0][0]
    lrd[i][0].append(sum_lrdOfO)
add_lrdOfO(lrd)

lrd_rdd = sc.parallelize(lrd)
def lof(entry,k):
  lof = entry[0][2]/k/entry[0][0]
  return [lof,entry[1]]
lrd = lrd_rdd.map(lambda x:lof(x,k_neighbors)).collect()

def n_outlier(entry,n):
  entry = sorted(entry, key=lambda x: x[0] ,reverse = True)
  return entry[0:n]
outlier = n_outlier(lrd,100)


import json # we need to use the JSON package to load the data, since the data is stored in JSON format
import matplotlib.pyplot as plt # visualization
import numpy as np
from numpy import linalg as LA # calculate norm
import pandas as pd
import time # time different methods
from collections import Counter # select most frequent elements in a list

# =============================================================================
# MAIN function is at the end of this file !!!
# =============================================================================

def text_modification(data):
    for x in data:
        x['text'] = x['text'].lower()
        x['text'] = x['text'].split()
    return data

# Top frequent words
# "data" is the input list of dictionaries, and "num" is the number of top frequent words needed
def top_frequent_words(data, num):
    words = []
    for x in data:
        words = words + x['text']

    # Find 160 most frequent words
    Count = Counter(words)
    most_occur = Count.most_common(num)

    # Get rid of the tuple in the list most_occur1 and most_occur2
    words1 = []
    for x in most_occur:
        words1.append(x[0])

#    print(words1)
    return words1

# Create the "word count features"
def text_features(data, wordList):
    word_matrix = [] # a matrix shows the frequency of the top words occurring in comments
    for x in wordList:
        temp = []
        for y in data:
            num = y['text'].count(x)
            temp.append(num)
        word_matrix.append(temp)

    return word_matrix

# Feature and label extraction
def information_extraction(data):
    text = []
    is_root = []
    controversiality = []
    children = []
    log_children = []
    sqrt_children = []
    popularity_score = []

    for x in data:
        text.append(x['text'])
        if x['is_root']==False:
            is_root.append(0)
        else:
            is_root.append(1)
        controversiality.append(x['controversiality'])
        children.append(x['children'])
        log_children.append(np.log10(x['children']+1))
        sqrt_children.append(np.sqrt(x['children']))

        popularity_score.append(x['popularity_score'])

    return text, is_root, controversiality, children, log_children, sqrt_children, popularity_score

# Create dummy variable list
def onelist(n):
    listofones = [1] * n
    return listofones

# Create a new feature: number of words in each comment
def comment_words(X_text):
    X_numwords = []

    for x in X_text:
        X_numwords.append(len(x))

    return X_numwords

# Sort the comments according to propularity scores
# select out the top frequent 150 words in top 50 comments
# remove the first 100 words
# Then count the number of presence of these words in each comment
def topfreq_highscore(data, wordList):
    sorted_list = sorted(data, key=lambda k: k['popularity_score'], reverse=True) # records sorted by popularity scores

    # From the 50 comments with top scores, select out 150 words with top frequency
    words_150 = top_frequent_words(sorted_list[0:49], 150)
    # Remove the first 100 words since they probably have low correlations with popularity scores
    words_50 = words_150[101:]

    # Count the number of presence of these of words in each comment
    count = []
    for x in data:
        temp = 0
        for words in words_50:
            if words in x['text']:
                temp += 1
        count.append(temp)

    return count

## Using closed-form
def error_function(W_est, X, y):
    # Calculate MSE
    totalError = np.dot((y-np.dot(X, W_est)).T, y-np.dot(X, W_est))
    return totalError[0][0]/len(y)

def closed_form_runner(X, y):
    XTX = X.T.dot(X)

    if np.linalg.det(XTX) == 0.0:
        # If the matrix is singular, this shows that some of the features
        # are not linearly independent
        print("Singular matrix error")
        return

    XTX_inv = pd.DataFrame(np.linalg.inv(XTX.values), XTX.columns, XTX.index)

    # get estimation of weight matrix
    W_est = XTX_inv.dot(X.T.dot(y))

    err = error_function(W_est, X, y)
    return err, W_est

## Using gradient descent
def gradient_descent_setup(X, y, eta_0, beta, epsilon):
    # Initialize W with zeros
    W_0 = np.zeros(shape=(X.shape[1],1))
    W_est, err_record = gradient_descent_runner(X, y, W_0, beta, eta_0, epsilon)
    err = error_function(W_est, X, y)
    return W_est, err_record, err

def gradient_descent_runner(X, y, W_0, beta, eta_0, epsilon):
    W = [W_0, 1+W_0]

    err_record = []

    # Calculate XTX outside the loop the reduce running time
    XTX = np.dot(X.T,X)
    if np.linalg.det(XTX) == 0.0:
        # If the matrix is singular, this shows that some of the features
        # are not linearly independent
        print("Singular matrix error")
        return

    XTy = np.dot(X.T,y)

    # scale the learning rate by 1/n, where n=# training points
    scale = 1/y.shape[0]

    i = 1
    while LA.norm(W[1]-W[0]) > epsilon:
        if i != 1:
            W[0] = W[1]
        alpha = scale * eta_0 / (1 + beta*i)
        W[1] = W[0] - 2 * alpha * (np.dot(XTX, W[0]) - XTy)
        i+=1

        #print(LA.norm(W[1]-W[0]))
        # Store the error
        if i % 3 == 0:
            err_record.append(error_function(W[1], X, y))

        if i > 100000:
            break

    return W[1], err_record

def data_preprocessing(data, word_list):
    # Create training sets
    X_text, X_root, X_contro, X_chil, X_logchil, X_sqrtchil, X_pop = information_extraction(data)

    # The reason I partition features into severl lists is to make it convenient when
    # we choose to discard one or more features
    ones = onelist(len(X_pop))

    num_words = comment_words(X_text)
    num_words_recip = np.reciprocal(num_words)
    num_words_log = np.log([x+1 for x in num_words])

    words_freq = topfreq_highscore(data, word_list)

    # Merge feature lists into dataframe
    X_train = pd.DataFrame(
        {'is_root': X_root,
         'controversiality': X_contro,
         'children': X_chil,
         'log_children': X_logchil,
         'sqrt_children': X_sqrtchil,

         # New features (uncomment to see the effect)
         'words_freq': words_freq,
         'num_words_log': num_words_log,
         'num_words_recip': num_words_recip,

         'dummy var': ones
        })

# =============================================================================
#     # Add text features (uncomment to see the influence of the text features)
    word_matrix = text_features(data, word_list)
    i = 0
    for x in word_matrix:
        col_name = word_list[i]
        X_train[col_name] = x
        i = i+1
# =============================================================================

    y = pd.DataFrame({'popularity_score': X_pop})

    return X_train, y

# =============================================================================
# Trainig Part
# =============================================================================

def training_runner(data, init_eta, init_beta, init_epsilon, word_list):

    X_train, y_train = data_preprocessing(data, word_list)
    # Closed-form
    start = time.time()
    err_ls, W_ls = closed_form_runner(X_train, y_train)
    end = time.time()
    time_ls = end-start

    err_gd = []
    W_gd = 0

    # Gradient descent
    # (ONLY uncomment gradient descent when text features are not included!!!!!!!!!)
#    start = time.time()
#    W_gd, err_gd, err = gradient_descent_setup(X_train, y_train, init_eta, init_beta, init_epsilon)
#    end = time.time()
#    time_gd = end -start

# =============================================================================
# Task 3
    # Closed-form plot
    print("Time using closed form on training set:"+str.format('{0:.20f}', time_ls))
    print("Error using closed form on training set:", err_ls)

    y_train = np.array(y_train.values.tolist())
    t = np.arange(0., len(y_train),1)

    fig = plt.figure()
    plt.plot(t, y_train, 'k.', t, np.dot(X_train, W_ls), 'r.')
    fig.suptitle("Closed-form on training set (in red) vs True values", fontsize=14)
    plt.xlabel('data points', fontsize=14)
    plt.ylabel('popularity score', fontsize=14)
    plt.show()

    # Gradient descent plot
    # (ONLY uncomment gradient descent when text features are not included!!!!!!!!!)
#    print("Time using gradient descent on training set:"+str.format('{0:.20f}', time_gd))
#    print("Error using gradient descent on training set:", err)
#
#    fig = plt.figure()
#    plt.plot(t, y_train, 'k.', t, np.dot(X_train, W_gd), 'r.')
#    fig.suptitle("Gradient descent fitting values (red) vs Label values", fontsize=14)
#    plt.xlabel('data points', fontsize=14)
#    plt.ylabel('popularity score', fontsize=14)
#    plt.show()
# =============================================================================

    return W_ls, W_gd, err_gd

# =============================================================================
# Validation Part
# =============================================================================

def validation_runner(data, W_ls, word_list):
    X_val, y_val = data_preprocessing(data, word_list)

    val_err_ls = error_function(W_ls, X_val, y_val)
    print("Error using closed form on validation set:", val_err_ls)

    # Plot
    y_val = np.array(y_val.values.tolist())
    t = np.arange(0., len(y_val),1)
    fig = plt.figure()
    plt.plot(t, y_val, 'k.', t, np.dot(X_val, W_ls), 'r.')
    fig.suptitle("Closed-form on validation set (in red) vs True values", fontsize=14)
    plt.xlabel('data points', fontsize=14)
    plt.ylabel('popularity score', fontsize=14)
    plt.show()

# =============================================================================
# Test Part
# =============================================================================

def testing_runner(data, W_ls, word_list):
    X_test, y_test = data_preprocessing(data, word_list)

    err_ls = error_function(W_ls, X_test, y_test)
    print("Error using closed-form on test set:", err_ls)

    # Plot
    y_test = np.array(y_test.values.tolist())
    t = np.arange(0., len(y_test),1)
    fig = plt.figure()
    plt.plot(t, y_test, 'k.', t, np.dot(X_test, W_ls), 'r.')
    fig.suptitle("Closed-form on test set (in red) vs True values", fontsize=14)
    plt.xlabel('data points', fontsize=14)
    plt.ylabel('popularity score', fontsize=14)
    plt.show()

# =============================================================================
# Main function of the file
# =============================================================================
def main():
    # load data
    with open("proj1_data.json") as fp:
        data = json.load(fp)

    # deal with comment text
    data = text_modification(data)

    # Training, Validation, Test sets split
    training, validation, test = np.split(data, [int((5/6)*len(data)), int((11/12)*len(data))])

    # get top frequent words
    wordList = top_frequent_words(training, 60)

#    # Gradient descent hyparameters
    init_eta = 0.00003 # initial learning rate
    init_beta = 0.0000001 # controls the speed of the decay
    init_epsilon = 0.000001

    # training part
    W_ls, W_gd, err_record = training_runner(training, init_eta, init_beta, init_epsilon, wordList)

    # Gradient descent error plot
    # (ONLY uncomment when text features are not included)
#    i = 0
#    fig = plt.figure()
#    plt.plot(err_record)
#    name = "eta=" + str(init_eta) + " beta=" + str(init_beta) + " epsilon=" + str(init_epsilon)
#    fig.suptitle(name, fontsize=16)
#    plt.xlabel('iterations', fontsize=14)
#    plt.ylabel('Error', fontsize=14)
#    i = i+1

    # validation part
    validation_runner(validation, W_ls, wordList)
    # test part
    testing_runner(test, W_ls, wordList)

main()


import sys
import time
import torch
import argparse
import PIL.ImageOps
from tqdm import tqdm
import torch.nn as nn
from torch import optim
import torchvision.utils
import matplotlib.pyplot as plt
import torch.nn.functional as F
import matplotlib.pyplot as plt
import torchvision.datasets as dset
import torchvision.datasets.fakedata
import torchvision.transforms as transforms
from loss import ContrastiveLoss
from torch.utils.data import DataLoader
from siamesenetwork import SiameseNetwork
from siamesenetwork import SiameseNetworkDataset

def show_plot(iteration, loss):
    plt.plot(iteration, loss)
    plt.xlabel("iteration")
    plt.ylabel("loss")
    plt.savefig("loss_figure.png")
    plt.show()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"training device: {device}")

args = argparse.ArgumentParser()
args.add_argument("--model")
res = args.parse_args()
model_name = res.model

net = None
loaded_model = None

net = SiameseNetwork().to(device)
loaded_model = torch.load(model_name, map_location=torch.device("cuda"))
net.load_state_dict(loaded_model["state_dict"])

criterion = ContrastiveLoss().to(device)
optimizer = optim.Adam(net.parameters(), lr=0.0005)

optimizer.load_state_dict(loaded_model["optim_dict"])
epochs = loaded_model["epoch"]
loss_history = loaded_model["loss_history"]
counter = loaded_model["counter"]

if net is not None:
    print("model loaded... ")
else:
    print("model loading failed... try again")

print(f"resuming training from: {model_name}")

class Config:
    training_dir = "./images/train_q/"
    train_batch_size = 64
    train_number_epochs = 5000

folder_dataset = dset.ImageFolder(root=Config.training_dir)

siamese_dataset = SiameseNetworkDataset(
    imageFolderDataset=folder_dataset,
    transform=transforms.Compose(
        [
            transforms.Resize((100, 100)),
            transforms.RandomHorizontalFlip(),
            transforms.RandomVerticalFlip(),
            transforms.ColorJitter(),
            transforms.ToTensor(),
        ]
    ),
    should_invert=False,
)

train_dataloader = DataLoader(
    siamese_dataset,
    shuffle=True,
    # num_workers=8,
    batch_size=Config.train_batch_size,
)

iteration_number = counter[-1]

for epoch in tqdm(
    range(epochs, Config.train_number_epochs + epochs),
    ascii=True if sys.platform == "win32" else False,
):
    for i, data in enumerate(train_dataloader, 0):
        start = time.time()
        img0, img1, label = data
        img0, img1, label = img0.to(device), img1.to(device), label.to(device)
        optimizer.zero_grad()
        output1, output2 = net(img0, img1)
        loss_contrastive = criterion(output1, output2, label)
        loss_contrastive.backward()
        optimizer.step()
        if i % 10 == 0:
            print(
                "epoch: {} current loss: {} time taken: {}s\n".format(
                    epoch + 1, loss_contrastive.item(), str(time.time() - start)[:5]
                )
            )
            iteration_number += 10
            counter.append(iteration_number)
            loss_history.append(loss_contrastive.item())

    if epoch % 500 == 0 and epoch > 0:
        state_dict = {
            "epoch": epoch + 1,
            "state_dict": net.state_dict(),
            "optim_dict": optimizer.state_dict(),
            "loss_history": loss_history,
            "counter": counter,
        }
        torch.save(state_dict, f"model_state_dict_{epoch}.pt")
        print()
        print(f"model checkpoint saved to models/model_state_dict_{epoch    }")

show_plot(counter, loss_history)


"""Support for tracking a Nissan Leaf."""
import logging

from homeassistant.helpers.dispatcher import dispatcher_connect
from homeassistant.util import slugify

from . import DATA_LEAF, DATA_LOCATION, SIGNAL_UPDATE_LEAF

_LOGGER = logging.getLogger(__name__)

ICON_CAR = "mdi:car"

def setup_scanner(hass, config, see, discovery_info=None):
    """Set up the Nissan Leaf tracker."""
    if discovery_info is None:
        return False

    def see_vehicle():
        """Handle the reporting of the vehicle position."""
        for vin, datastore in hass.data[DATA_LEAF].items():
            host_name = datastore.leaf.nickname
            dev_id = 'nissan_leaf_{}'.format(slugify(host_name))
            if not datastore.data[DATA_LOCATION]:
                _LOGGER.debug("No position found for vehicle %s", vin)
                return
            _LOGGER.debug("Updating device_tracker for %s with position %s",
                          datastore.leaf.nickname,
                          datastore.data[DATA_LOCATION].__dict__)
            attrs = {
                'updated_on': datastore.last_location_response,
            }
            see(dev_id=dev_id,
                host_name=host_name,
                gps=(
                    datastore.data[DATA_LOCATION].latitude,
                    datastore.data[DATA_LOCATION].longitude
                ),
                attributes=attrs,
                icon=ICON_CAR)

    dispatcher_connect(hass, SIGNAL_UPDATE_LEAF, see_vehicle)

    return True


from __future__ import print_function

import os

try:
    from StringIO import StringIO
except ImportError:
    from io import StringIO
from tqdm import trange
from models import *
from tensorflow.contrib.slim.python.slim.nets import resnet_v1
from tensorflow.contrib.slim.python.slim.nets.mobilenet import mobilenet_v2

class Trainer(object):
    def __init__(self, config, get_loader):

        self.config = config
        self.get_loader = get_loader
        self.dataset = config.dataset

        self.beta1 = config.beta1
        self.beta2 = config.beta2
        self.batch_size = config.batch_size

        self.step = tf.Variable(0, name='step', trainable=False)

        self.j = config.j
        self.lr = tf.Variable(config.lr, name='lr')
        self.lr_mtplr = tf.Variable(config.lr_mtplr, name='lr_mtplr')

        self.model_dir = config.model_dir
        self.load_path = config.load_path
        self.summary_path = os.path.join(config.ckpt_dir)
        self.pre_dir = config.data_dir[0] + '/' + config.pre_dir

        self.use_gpu = config.use_gpu
        self.data_format = config.data_format

        self.start_step = 0
        self.log_step = config.log_step
        self.save_step = config.save_step
        self.k_t = tf.Variable(0., trainable=False)

        data = self.get_loader
        self.iterator = data.make_initializable_iterator()
        self.data = self.iterator.get_next()

        self.is_train = config.is_train
        self.build_model()

        self.saver = tf.compat.v1.train.Saver()
        self.summary_writer = tf.compat.v1.summary.FileWriter(self.summary_path)

        sv = tf.compat.v1.train.Supervisor(logdir=self.model_dir,
                                 is_chief=True,
                                 saver=self.saver,
                                 summary_op=None,
                                 summary_writer=self.summary_writer,
                                 save_model_secs=1200,
                                 global_step=self.step,
                                 ready_for_local_init_op=None)

        gpu_options = tf.compat.v1.GPUOptions(allow_growth=True)
        sess_config = tf.compat.v1.ConfigProto(allow_soft_placement=True, gpu_options=gpu_options)

        self.sess = sv.prepare_or_wait_for_session(config=sess_config)

        if not self.is_train:
            g = tf.get_default_graph()
            g._finalized = False

    def train(self):
        self.initial_iterator()
        self.init_fn(self.sess)
        self.init_fn_1(self.sess)
        # teacher first, student last
        # 2d train first, 3d last
        for step in trange(self.start_step, self.max_step):
            if step < self.two_d_pose:
                fetch_dict = {
                    '2d_t': self.two_d_teacher_optim,
                }
            elif step < self.two_d_pose + self.three_d_pose:
                if step is self.two_d_pose:
                    self.data_change()
                fetch_dict = {
                    '3d_t': self.three_d_teachter_optim,
                }
            elif step < self.two_d_pose + self.three_d_pose + self.two_d_pose:
                fetch_dict = {
                    '2d_s': self.two_d_teacher_optim,
                }
            else:
                fetch_dict = {
                    '3d_s': self.three_d_student_optim,
                }
            if step % self.log_step == 0:
                fetch_dict.update({
                    "summary": self.summary_op,
                    "teacher_heat_map_loss": self.t_hm_loss,
                    "teacher_local_map_loss": self.t_lm_loss,
                    "student_heat_map_loss": self.s_hm_loss,
                    "student_local_map_loss": self.s_lm_loss,
                })
            result = self.sess.run(fetch_dict)

            if step % self.log_step == 0:
                self.summary_writer.add_summary(result['summary'], step)
                self.summary_writer.flush()
                t_hm_loss = result['teacher_heat_map_loss']
                t_lm_loss = result['teacher_local_map_loss']
                s_hm_loss = result['student_heat_map_loss']
                s_lm_loss = result['student_local_map_loss']
                print(
                    "[{}/{}] Loss_FD: [{:.6f}] Loss_SD: [{:.6f}] || Loss_G: [{:.6f}]" \
                        .format(step, self.max_step, t_hm_loss, t_lm_loss, s_hm_loss, s_lm_loss))

    def build_model(self):
        with tf.compat.v1.variable_scope('preprocessing'):
            two_d_x, two_d_gt = self.preprocessing()
            x, gt = two_d_x, two_d_gt

        batch_input = utils.upscale_to(x, [224, 224], 'NCHW')
        batch_input, batch_mean_rgb = self.mean_rgb(batch_input)
        t, self.t_var = self.teacher(batch_input, self.j)
        s, self.s_var = self.student(batch_input, self.j)
        # print(t, s)
        with tf.compat.v1.variable_scope("loss_design"):
            with tf.compat.v1.variable_scope("optimizer"):
                two_d_optimizer = tf.compat.v1.train.RMSPropOptimizer(self.lr)
                three_d_optimizer = tf.compat.v1.train.AdamOptimizer(self.lr * self.lr_mtplr)

            with tf.compat.v1.variable_scope("loss_function"):
                t = utils.upscale_to_bil(t, [256, 256], 'NCHW')
                s = utils.upscale_to_bil(s, [256, 256], 'NCHW')
                t_h, t_l = tf.split(t, [self.j, -1], axis=1)  # heat for joint, x/y/z for joint
                s_h, s_l = tf.split(s, [self.j, -1], axis=1)  # heat for joint, x/y/z for joint
                gt_h, gt_l = two_d_gt, tf.concat([two_d_gt, two_d_gt, two_d_gt], 1)  # but data has only where is heat and x/y/z for joint?
                with tf.compat.v1.variable_scope("t_loss"):
                    self.t_hm_loss = tf.reduce_mean(tf.pow(tf.abs(t_h - gt_h), 2.))
                    self.t_lm_loss = tf.reduce_mean(tf.pow(tf.abs(t_l - gt_l), 2.))
                with tf.compat.v1.variable_scope("s_loss"):
                    alpha = 0.5
                    self.s_hm_loss = tf.reduce_mean(alpha * tf.pow(tf.abs(s_h - gt_h), 2.)
                                            + (1.-alpha) * tf.pow(tf.abs(s_h - t_h), 2.)
                    )
                    self.s_lm_loss = tf.reduce_mean(alpha * tf.pow(tf.concat([gt_h, gt_h , gt_h], 1) * tf.abs(s_l - gt_l), 2.)
                                             + (1. - alpha) * tf.pow(tf.concat([gt_h, gt_h , gt_h], 1) * tf.abs(s_l - t_l), 2.)
                                             )
            with tf.compat.v1.variable_scope("minimizer"):
                self.two_d_teacher_optim = two_d_optimizer.minimize(self.t_hm_loss + self.t_lm_loss, var_list=self.t_var)
                self.two_d_studnet_optim = two_d_optimizer.minimize(self.s_hm_loss + self.s_lm_loss, var_list=self.s_var)
                self.three_d_teachter_optim = three_d_optimizer.minimize(self.t_hm_loss + self.t_lm_loss, var_list=self.t_var)
                self.three_d_student_optim = three_d_optimizer.minimize(self.s_hm_loss + self.s_lm_loss, var_list=self.s_var)

        with tf.compat.v1.variable_scope("de_normal"):
            self.t_h = tf.clip_by_value((t_h + 1.) * 127.5, 0., 255.)
            self.t_l = tf.clip_by_value((t_l + 1.) * 127.5, 0., 255.)
            self.s_h = tf.clip_by_value((s_h + 1.) * 127.5, 0., 255.)
            self.s_l = tf.clip_by_value((s_l + 1.) * 127.5, 0., 255.)
            self.gt_h = tf.clip_by_value((gt_h + 1.) * 127.5, 0., 255.)
            self.gt_l = tf.clip_by_value((gt_l + 1.) * 127.5, 0., 255.)

        self.summary_op = tf.compat.v1.summary.merge([
            tf.compat.v1.summary.scalar("loss/t/hm_loss", self.t_hm_loss),
            tf.compat.v1.summary.scalar("loss/t/lm_loss", self.t_lm_loss),
            tf.compat.v1.summary.scalar("loss/s/hm_loss", self.s_hm_loss),
            tf.compat.v1.summary.scalar("loss/s/lm_loss", self.s_lm_loss),
            tf.compat.v1.summary.image("t/h", self.t_h),
            tf.compat.v1.summary.image("t/l", self.t_l),
            tf.compat.v1.summary.image("s/h", self.s_h),
            tf.compat.v1.summary.image("s/l", self.s_l),
            tf.compat.v1.summary.image("gt/h", self.gt_h),
            tf.compat.v1.summary.image("gt/l", self.gt_l),
        ])

    def build_test_model(self):
        print("test_model")

    def test(self):
        self.build_test_model()

    def initial_iterator(self):
        self.sess.run(self.iterator.initializer)

    def preprocessing(self):
        with tf.compat.v1.variable_scope("Pre"):
            two_d_x, two_d_joint = self.data  # [B, 256, 256, 3]  [B, 256, 256, 15]
            print(two_d_x, two_d_joint)
            two_d_x= self.norm(two_d_x)
            two_d_x = utils.nhwc_to_nchw(two_d_x)
            two_d_joint = utils.nhwc_to_nchw(two_d_joint)
        return two_d_x, two_d_joint

    def norm(self, x):
        return tf.clip_by_value((x/127.5)-1., -1., 1.)

    def load_data(self):
        self.initial_iterator()
        x = self.data
        return x

    def generate(self, image, fixed, root_path=None, path=None, idx=None, save=True):
        x = self.sess.run(self.Gz, {self.z: image})
        if path is None and save:
            path = os.path.join(root_path, '{}_G.mp4'.format(idx))
            utils.save_video(x, fixed, path)
            print("[*] Samples saved: {}".format(path))
        return x

    def mean_rgb(self, x):
        r, g, b = tf.split(x, 3, axis=1)
        mean_r = tf.reduce_mean(r, axis=[2, 3], keep_dims=True)
        mean_g = tf.reduce_mean(g, axis=[2, 3], keep_dims=True)
        mean_b = tf.reduce_mean(b, axis=[2, 3], keep_dims=True)
        r = r - mean_r
        g = g - mean_g
        b = b - mean_b
        return tf.concat([r, g, b], 1), tf.concat([mean_r, mean_g, mean_b], 0)

    def teacher(self, x, j):
        with slim.arg_scope(resnet_v1.resnet_arg_scope()):
            x = utils.nchw_to_nhwc(x)
            batch_out, batch_list = resnet_v1.resnet_v1_50(x, 1000, is_training=True)
            feature = batch_list['resnet_v1_50/block2/unit_4/bottleneck_v1/conv1']

        self.init_fn_1 = slim.assign_from_checkpoint_fn(
            self.pre_dir + '/resnet_v1_50.ckpt', slim.get_model_variables('resnet_v1_50'))
        '''
        del which has no gradient
        '''
        # print(batch_list)
        x = utils.nhwc_to_nchw(feature)
        x, var = vnect(x, j)
        return x, var

    def student(self, x, j):
        with tf.contrib.slim.arg_scope(mobilenet_v2.training_scope()):
            x = utils.nchw_to_nhwc(x)
            batch_out, batch_list = mobilenet_v2.mobilenet_v2_035(x, 1000, is_training=True)
            # print(batch_list)
            mobilenet_feature = batch_list['layer_12/output']

        self.init_fn = slim.assign_from_checkpoint_fn(
            self.pre_dir + '/mobilenet_v2_035.ckpt', slim.get_model_variables('MobilenetV2'))

        x = utils.nhwc_to_nchw(mobilenet_feature)
        x, var = movnect(x, j)
        return x, var


import os
import sys
import json
import numpy as np
import pandas as pd
from pathlib import Path
from sklearn.metrics import mean_squared_error
from scipy.stats import linregress

def print_json(data):
    print(json.dumps(data, indent=2))

def eprint(*args, **kwargs):
    print(*args, file=sys.stderr, **kwargs)

def GenDataFrameFromPath(path, pattern='*.png', fs=False):
    """
    generate a dataframe for all file in a dir with the specific pattern of file name.
    use: GenDataFrameFromPath(path, pattern='*.png')
    """
    fnpaths = list(path.glob(pattern))
    df = pd.DataFrame(dict(zip(['fnpath'], [fnpaths])))
    df['dir'] = df['fnpath'].apply(lambda x: x.parent)
    df['fn'] = df['fnpath'].apply(lambda x: x.name)
    if fs:
        df['size'] = df['fnpath'].apply(lambda x: os.path.getsize(x))
    return df

def ConciseVcf(fn):
    """
    concise the vcf file by remove the header, useless columns and simplfied genotype
    ConciseVcf(fn)
    """
    n = 0
    f = open(fn)
    for i in f:
        if i.startswith('##'):
            n += 1
        else:
            break
    df = pd.read_csv(fn, header=n, delim_whitespace=True)
    df = df.drop(['INFO', 'FORMAT', 'ID', 'REF', 'ALT', 'QUAL', 'FILTER'], axis=1)
    for idx in df.columns[2:]:
        df[idx] = df[idx].map(lambda x: x.split(':')[0])
    df = df.replace(['0/0', '0/1', '1/0', '1/1', './.'], [0, 1, 1, 2, 9])
    return df

class SimpleStats(object):
    """
    This class will do the simple statistics on two series objecjts.
    a) linear regressoin: slope, intercept, r^2, p_value
    b) mean, std of the difference and absolute differnece
    c) MSE (mean squared error) and RMSE (root mean squared error)
    d) agreement
    e) plot the regreesion figure and the difference distribution figure
    """
    def __init__(self, series1, series2):
        self.s1 = series1
        self.s2 = series2
        self.length = series1.shape[0]
        self.diff = series1 - series2
        self.absdiff = (series1 - series2).abs()

    def regression(self):
        slope, intercept, r_value, p_value, __ = linregress(self.s1, self.s2)
        return slope, intercept, r_value**2, p_value

    def mean_std_diff(self):
        mean, std = self.diff.mean(), self.diff.std()
        return mean, std

    def mean_std_absdiff(self):
        abs_mean, abs_std = self.absdiff.mean(), self.absdiff.std()
        return abs_mean, abs_std

    def mse(self):
        mse = mean_squared_error(self.s1, self.s2)
        return mse

    def rmse(self):
        rmse = mean_squared_error(self.s1, self.s2)**0.5
        return rmse

    def agreement(self, cutoff):
        return (self.absdiff<=float(cutoff)).sum()/self.length

def simulate_gwas_results(num_chr=10, num_snp=50000, max_chr_size=60000000,
                            num_gwas=25, sig_cutoff=6.5, num_sig_snps=5):
    '''
    num_chr: number of chromosomes
    num_snp: number of snps in each chr
    max_chr_size: the maximum chromosome size for each chr
    num_gwas: number of gwas analyses
    sig_cutoff: the -log transformed p value cutoff
    num_sig_snps: number of significant snps in each gwas analysis
    '''
    chrs = np.array([np.full(num_snp, i) for i in range(1, chrom+1)]).flatten()
    # simulate snp positions in each chromosomes
    snp_pos = np.array([np.sort(np.random.randint(0, max_chr_size, num_snp)) for i in range(1, 11)]).flatten()
    # simuate -log transformed pvalues in 25 GWAS analyses (assume all pvalues are not significant)
    df_pvalues = pd.DataFrame(np.random.uniform(0, sig_cutoff, (num_snp, num_gwas))) # simuate pvalues

    # simuate positions and pvalues for significant SNPs, assume only 5 significant snps in each GWAS
    sig_pos = pd.DataFrame(np.random.randint(0, num_snp, (num_sig_snps, num_gwas)))
    sig_pvalues = pd.DataFrame(np.random.uniform(sig_cutoff, sig_cutoff+5, (num_sig_snps, num_gwas)))
    # introduce the simuated snps
    for col in range(num_gwas):
        for row in range(num_sig_snps):
            pos = sig_pos.iloc[row, col]
            pv = sig_pvalues.iloc[row, col]
            #print(pos, col)
            df_pvalues.iloc[pos, col] = pv
    df_pvalues.columns = np.arange(num_gwas)
    df_snp = pd.DataFrame(chrs, columns=['chr'])
    df_snp['pos'] = snp_pos
    df_gwas = pd.concat([df_snp, df_pvalues], axis=1)
    return df_gwas

def bin_gwas_snps(df, bin_size=16):
    '''
    df: the dataframe of your GWAS results
        Make sure the first column is 'chr' and 2nd column is 'pos'
        The rest columns are genus names.
    bin_size: how many bins for each chromosme (even integer)
    '''
    df = df.set_index(['chr', 'pos'])
    chros, sts, eds, ps = [], [], [], []
    for chrom, tmp in df.groupby(level=0):
        tmp = tmp.reset_index()
        bins = pd.qcut(tmp['pos'], bin_size)
        grps = tmp.groupby(bins)
        for _, grp in grps:
            idx = grp['pos'].values
            st, ed = idx[0], idx[-1]
            maxs = grp.iloc[:,2:].max(axis=0)
            chros.append(chrom)
            sts.append(st)
            eds.append(ed)
            ps.append(maxs)
    df_bins = pd.DataFrame(ps)
    df_bins.index = chros
    df_plot = df_bins.transpose().fillna(0)
    df_bins_info = pd.DataFrame(dict(zip(['chr', 'st', 'ed'], [chros, sts, eds])))
    return df_plot, df_bins_info


import datetime
import time
import os
import sys
import logging
import argparse
import pkg_resources

from esrally import config, paths, racecontrol
from esrally.utils import io

__version__ = pkg_resources.require("esrally")[0].version

# we want to use some basic logging even before the output to log file is configured
def preconfigure_logging():
    logging.basicConfig(level=logging.INFO)

def configure_logging(cfg):
    log_dir = paths.Paths(cfg).log_root()
    io.ensure_dir(log_dir)
    cfg.add(config.Scope.application, "system", "log.dir", log_dir)
    log_file = "%s/rally_out.log" % log_dir

    print("\nWriting additional logs to %s\n" % log_file)

    # Remove all handlers associated with the root logger object so we can start over with an entirely fresh log configuration
    for handler in logging.root.handlers[:]:
        logging.root.removeHandler(handler)

    if cfg.opts("system", "quiet.mode"):
        log_level = logging.ERROR
    else:
        log_level = logging.INFO

    ch = logging.FileHandler(filename=log_file, mode="a")
    ch.setLevel(log_level)
    formatter = logging.Formatter("%(asctime)s,%(msecs)d %(name)s %(levelname)s %(message)s", datefmt="%Y-%m-%d %H:%M:%S")
    formatter.converter = time.gmtime
    ch.setFormatter(formatter)
    logging.root.addHandler(ch)

def parse_args():
    parser = argparse.ArgumentParser(prog="esrally", description="Benchmark Elasticsearch")
    parser.add_argument('--version', action='version', version="%(prog)s " + __version__)

    subparsers = parser.add_subparsers(
        title="subcommands",
        dest="subcommand",
        help="")

    race_parser = subparsers.add_parser("race", help="Run the benchmarking pipeline. This subcommand should typically be used.")
    # change in favor of "list telemetry", "list tracks", "list pipelines"
    list_parser = subparsers.add_parser("list", help="List configuration options")
    list_parser.add_argument(
        "configuration",
        metavar="configuration",
        help="The configuration for which Rally should show the available options. Possible values are: telemetry, tracks, pipelines, races",
        choices=["telemetry", "tracks", "pipelines", "races"])
    list_parser.add_argument(
        "--limit",
        help="Limit the number of search results for recent races (default: 10).",
        default=10,
    )

    config_parser = subparsers.add_parser("configure", help="Write the configuration file or reconfigure Rally")
    for p in [parser, config_parser]:
        p.add_argument(
                "--advanced-config",
                help="show additional configuration options when creating the config file (intended for CI runs) (default: false)",
                default=False,
                action="store_true")

    for p in [parser, race_parser]:
        p.add_argument(
            "--pipeline",
            help="Selects a specific pipeline to run. A pipeline defines the steps that are executed (default: from-sources-complete).",
            default="from-sources-complete")
        p.add_argument(
            "--quiet",
            help="Suppresses as much as output as possible. Activate it unless you want to see what's happening during the "
                 "benchmark (default: false)",
            default=False,
            action="store_true")
        p.add_argument(
            "--offline",
            help="Assume that Rally has no connection to the Internet (default: false)",
            default=False,
            action="store_true")
        p.add_argument(
            "--preserve-install",
            help="preserves the Elasticsearch benchmark candidate installation including all data. Caution: This will take lots of disk "
                 "space! (default: false)",
            default=False,
            action="store_true")
        p.add_argument(
            "--telemetry",
            help="Rally will enable all of the provided telemetry devices (i.e. profilers). Multiple telemetry devices have to be "
                 "provided as a comma-separated list.",
            default="")
        p.add_argument(
            "--revision",
            help="defines which sources to use when building the benchmark candidate. 'current' uses the source tree as is,"
                 " 'latest' fetches the latest version on master. It is also possible to specify a commit id or a timestamp."
                 " The timestamp must be specified as: \"@ts\" where \"ts\" must be a valid ISO 8601 timestamp, "
                 "e.g. \"@2013-07-27T10:37:00Z\" (default: current).",
            default="current")  # optimized for local usage, don't fetch sources
        p.add_argument(
            "--distribution-version",
            help="defines the version of the Elasticsearch distribution to download. Check https://www.elastic.co/downloads/elasticsearch "
                 "for released versions.",
            default="")
        p.add_argument(
            "--track",
            help="defines which track should be run. Only one is allowed at a time (default: geonames)",
            default="geonames")
        p.add_argument(
            "--track-setup",
            help="defines which track-setups should be run. Multiple track setups can be specified as a comma-separated "
                 "list (default: defaults).",
            default="defaults")  # optimized for local usage
        p.add_argument(
            "--target-hosts",
            help="defines a comma-separated list of host:port pairs which should be targeted iff using the pipeline 'benchmark-only' "
                 "(default: localhost:9200).",
            default="localhost:9200")
        p.add_argument(
            "--user-tag",
            help="defines a user-specific key-value pair that is separated by a ':' and added to each metric record as meta info. "
                 "Example: intention:baseline-ticket-12345",
            default="")

    ###############################################################################
    #
    # The options below are undocumented and can be removed or changed at any time.
    #
    ###############################################################################
    for p in [parser, race_parser]:
        # This option is intended to tell Rally to assume a different start date than 'now'. This is effectively just useful for things like
        # backtesting or a benchmark run across environments (think: comparison of EC2 and bare metal) but never for the typical user.
        p.add_argument(
            "--effective-start-date",
            help=argparse.SUPPRESS,
            type=lambda s: datetime.datetime.strptime(s, "%Y-%m-%d %H:%M:%S"),
            default=datetime.datetime.utcnow())
        # This is a highly experimental option and will likely be removed
        p.add_argument(
            "--data-paths",
            help=argparse.SUPPRESS,
            default=None)

    for p in [parser, config_parser, race_parser]:
        # This option is needed to support a separate configuration for the integration tests on the same machine
        p.add_argument(
            "--configuration-name",
            help=argparse.SUPPRESS,
            default=None)

    return parser.parse_args()

def print_banner():
    print("    ____        ____     ")
    print("   / __ \____ _/ / /_  __")
    print("  / /_/ / __ `/ / / / / /")
    print(" / _, _/ /_/ / / / /_/ / ")
    print("/_/ |_|\__,_/_/_/\__, /  ")
    print("                /____/   ")

def derive_subcommand(args, cfg):
    subcommand = args.subcommand
    # first, trust the user...
    if subcommand is not None:
        return subcommand
    # we apply some smarts in case the user did not specify a subcommand
    if cfg.config_present():
        return "race"
    else:
        return "configure"

def csv_to_list(csv):
    if csv is None:
        return None
    else:
        return [e.strip() for e in csv.split(",")]

def main():
    preconfigure_logging()
    args = parse_args()
    cfg = config.Config(config_name=args.configuration_name)
    subcommand = derive_subcommand(args, cfg)

    if subcommand == "configure":
        cfg.create_config(advanced_config=args.advanced_config)
        exit(0)
    else:
        if cfg.config_present():
            cfg.load_config()
            if not cfg.config_compatible():
                # logger.info("Detected incompatible configuration file. Trying to upgrade.")
                cfg.migrate_config()
                # Reload config after upgrading
                cfg.load_config()
        else:
            print("Error: No config present. Please run 'esrally configure' first.")
            exit(64)

    # Add global meta info derived by rally itself
    cfg.add(config.Scope.application, "meta", "time.start", args.effective_start_date)
    cfg.add(config.Scope.application, "system", "rally.root", os.path.dirname(os.path.realpath(__file__)))
    cfg.add(config.Scope.application, "system", "invocation.root.dir", paths.Paths(cfg).invocation_root())
    # Add command line config
    cfg.add(config.Scope.applicationOverride, "source", "revision", args.revision)
    cfg.add(config.Scope.applicationOverride, "source", "distribution.version", args.distribution_version)
    cfg.add(config.Scope.applicationOverride, "system", "pipeline", args.pipeline)
    cfg.add(config.Scope.applicationOverride, "system", "track", args.track)
    cfg.add(config.Scope.applicationOverride, "system", "quiet.mode", args.quiet)
    cfg.add(config.Scope.applicationOverride, "system", "offline.mode", args.offline)
    cfg.add(config.Scope.applicationOverride, "system", "user.tag", args.user_tag)
    cfg.add(config.Scope.applicationOverride, "telemetry", "devices", csv_to_list(args.telemetry))
    cfg.add(config.Scope.applicationOverride, "benchmarks", "tracksetups.selected", csv_to_list(args.track_setup))
    cfg.add(config.Scope.applicationOverride, "provisioning", "datapaths", csv_to_list(args.data_paths))
    cfg.add(config.Scope.applicationOverride, "provisioning", "install.preserve", args.preserve_install)
    cfg.add(config.Scope.applicationOverride, "launcher", "external.target.hosts", csv_to_list(args.target_hosts))
    if subcommand == "list":
        cfg.add(config.Scope.applicationOverride, "system", "list.config.option", args.configuration)
        cfg.add(config.Scope.applicationOverride, "system", "list.races.max_results", args.limit)

    configure_logging(cfg)
    print_banner()

    race_control = racecontrol.RaceControl(cfg)
    success = race_control.start(subcommand)
    if not success:
        sys.exit(1)

if __name__ == "__main__":
    main()


# -*- coding: utf-8 -*-
from __future__ import unicode_literals

from django.contrib.admin.views.decorators import staff_member_required
from django.http import HttpResponseRedirect
from django.shortcuts import get_object_or_404
from django.urls import reverse_lazy
from django.utils.decorators import method_decorator
from django.views.generic.base import TemplateView

from rest_framework import generics
from rest_framework.authentication import SessionAuthentication
from rest_framework.response import Response

from .forms import PrimaryForm, SecondaryForm, TertiaryForm
from .models import (
    SubstanceSurvey,
    Region,
    City,
    AcquiredFrom,
    Origin,
    UserCode,
    Alias,
    Substance,
    Apperance,
    Kind,
    Color,
    TestMethod,
    Drug,
)
from .serializers import (
    SubstanceSurveySerializer,
    RegionSerializer,
    CitySerializer,
    AcquiredFromSerializer,
    OriginSerializer,
    UserCodeSerializer,
    AliasSerializer,
    SubstanceSerializer,
    ApperanceSerializer,
    KindSerializer,
    ColorSerializer,
    TestMethodSerializer,
    DrugSerializer,
)

class CsrfExemptSessionAuthentication(SessionAuthentication):
    def enforce_csrf(self, request):
        return  # To not perform the csrf check previously happening

@method_decorator(staff_member_required, name='dispatch')
class SubstanceSurveyAddView(TemplateView):
    template_name = 'survey/substancesurvey_add.html'

    def get_context_data(self, **kwargs):
        context = super(SubstanceSurveyAddView, self).get_context_data(**kwargs)
        context['primary_form'] = PrimaryForm()
        context['secondary_form'] = SecondaryForm()
        context['tertiary_form'] = TertiaryForm()
        return context

class SubstanceSurveyAPIView(generics.ListAPIView):
    serializer_class = SubstanceSurveySerializer

    def get_queryset(self):
        if self.request.user.is_staff:
            queryset = SubstanceSurvey.objects.all()
        else:
            sample = self.request.query_params.get('sample', None)
            if sample is not None:
                queryset = SubstanceSurvey.objects.filter(sample_code=sample)
            else:
                queryset = SubstanceSurvey.objects.exclude(user_code__isnull=False)

        return queryset

    def post(self, request, format=None):
        data = request.data
        s = SubstanceSurvey.objects.create()

        SURVEY_ONLY_FIELDS = (
            'date_acquired',
            'price',
            'sample_code',
            'image',
            'observations',
            'contact'
        )
        SURVEY_FOREIGN_FIELDS = (
            ('region', Region, 'name'),
            ('city', City, 'name'),
            ('acquired_from', AcquiredFrom, 'name'),
            ('origin', Origin, 'name'),
            ('user_code', UserCode, 'code'),
            ('alias', Alias, 'name'),
            ('substance', Substance, 'name'),
            ('apperance', Apperance, 'name')
        )
        SURVEY_M2M_FIELDS = (
            ('kinds', Kind),
            ('color', Color),
            ('testmethods', TestMethod)
        )

        for field in SURVEY_ONLY_FIELDS:
            value = data.get(field, '')
            if value:
                setattr(s, field, value)

        for atr, model, arg_name in SURVEY_FOREIGN_FIELDS:
            get_data = data.get(atr)
            if get_data != '':
                obj, crt = model.objects.get_or_create(**{arg_name: get_data})
                if not crt:
                    obj.save()
                setattr(s, atr, obj)

        for attr, model in SURVEY_M2M_FIELDS:
            get_data = data.get(attr)
            if get_data:
                data_split = get_data.split(",")
                for el in data_split:
                    add_el, crt = model.objects.get_or_create(name=el)
                field = getattr(s, attr)
                field.add(add_el)

        get_data = data.get('detected')
        if get_data:
            d_l = get_data.split(",")
            for d in d_l:
                try:
                    dect = Drug.objects.get(name=d)
                    s.detected.add(dect)
                except Drug.DoesNotExist:
                    pass

        s.save()
        return Response("OK")

class DuplicateSurveyView(TemplateView):
    def get(self, request, uuid, **kwargs):
        ss = get_object_or_404(SubstanceSurvey, uuid=uuid)
        ss.pk = None
        ss.save()
        return HttpResponseRedirect(reverse_lazy('index'))

class SubstanceSurveyListView(TemplateView):
    template_name = 'survey/substancesurvey_list.html'

class RegionAPIView(generics.ListAPIView):
    queryset = Region.objects.all()
    serializer_class = RegionSerializer

class CityAPIView(generics.ListAPIView):
    queryset = City.objects.all()
    serializer_class = CitySerializer

class AcquiredFromAPIView(generics.ListAPIView):
    queryset = AcquiredFrom.objects.all()
    serializer_class = AcquiredFromSerializer

class OriginAPIView(generics.ListAPIView):
    queryset = Origin.objects.all()
    serializer_class = OriginSerializer

class UserCodeAPIView(generics.ListAPIView):
    queryset = UserCode.objects.all()
    serializer_class = UserCodeSerializer

class AliasAPIView(generics.ListAPIView):
    queryset = Alias.objects.all()
    serializer_class = AliasSerializer

class SubstanceAPIView(generics.ListAPIView):
    queryset = Substance.objects.all()
    serializer_class = SubstanceSerializer

class ApperanceAPIView(generics.ListAPIView):
    queryset = Apperance.objects.all()
    serializer_class = ApperanceSerializer

class KindAPIView(generics.ListAPIView):
    queryset = Kind.objects.all()
    serializer_class = KindSerializer

class ColorAPIView(generics.ListAPIView):
    queryset = Color.objects.all()
    serializer_class = ColorSerializer

class TestMethodAPIView(generics.ListAPIView):
    queryset = TestMethod.objects.all()
    serializer_class = TestMethodSerializer

class DetectedAPIView(generics.ListAPIView):
    queryset = Drug.objects.all()
    serializer_class = DrugSerializer