import matplotlib
matplotlib.use('Qt5Agg')
from matplotlib.figure import Figure
import copy
import numpy as np
#import sys
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import seaborn as sns
import plotly.graph_objects as go
import scipy.ndimage as ndimage
#from afm_analyze import JPKAnalyze
#from plotly_viewer import PlotlyViewer
from plot2widget import PlotWidget

class AFMPlot:

    def __init__(self, jpk_anal=None, output_path=None):

        self.plotwin = None
        #mapping plot types (from ANALYSIS_MODE_DICT) to functions
        PLOT_DICT = {'2d': self.plot_2d,
                     '3d': self.plot_3d,
                     'line': self.plot_line}
        self.CLICK_STATUS = False #True when clicked first and plot doesn't exist

        if jpk_anal != None:
            self.jpk_anal = jpk_anal
            self.file_path = jpk_anal.file_path
            #plot data
            for mode in jpk_anal.df.keys():
                plot_params =  jpk_anal.anal_dict[mode]['plot_parameters']
                for plot_type in plot_params['type']:
                    fig = PLOT_DICT[plot_type](jpk_anal.df[mode], plot_params,
                                         file_path=output_path,
                                         points=plot_params['points_flag'],
                                        return_fig=True)
                    jpk_anal.anal_dict['Misc']['figure_list'].append(fig)

            #plt.show(block=False)
##        plt.pause(0.05)
##        if self.plotwin != None:
####            self.plotwin.show()
##            self.plotwin.app.exec_()

    def plot_2d(self, df, plot_params, file_path=None, points=False,return_fig=False):
        x = plot_params['x']
        y = plot_params['y']
        z = plot_params['z']
        title = plot_params['title']
        self.points_data = np.empty((0, 3), float)
        #organize data into matrix for heatmap plot
        df_data = df.pivot_table(values=z, index=y, columns=x,
                                 aggfunc='first')

        #plot
        plt.rcParams.update(plt.rcParamsDefault)
        fig2d = plt.figure(f'{title}')
        ax2d = fig2d.add_subplot(111)
        im2d = ax2d.pcolormesh(df_data.columns, df_data.index,
                               df_data, cmap='afmhot')
        ax2d.ticklabel_format(style='sci', scilimits=(0,0))
        ax2d.set_xlabel(x)
        ax2d.set_ylabel(y)
        fig2d.colorbar(im2d, ax=ax2d, label=z, format='%.1e')
        fig2d.suptitle(title)

        fig2d.savefig(f'{file_path}/{title}.png', bbox_inches = 'tight',
                      transparent = True)

        canvas = fig2d.canvas

##        rgb_data = np.fromstring(canvas.tostring_rgb(),
##                                 dtype=np.uint8, sep='')
##        rgb_data = rgb_data.reshape(canvas.get_width_height()[::-1] + (3,))

        if points == False:
            df_reference = df.pivot_table(values='Segment folder', index=y,
                                          columns=x, aggfunc='first')
            self.cid = fig2d.canvas.mpl_connect('button_press_event',
                                           lambda event: self.on_click(event,
                                                                       df_reference))
        else:
            self.cid = fig2d.canvas.mpl_connect('button_press_event',
                                           lambda event: self.on_click(event,
                                                                       df_data, points))

        #BUG: program doesn't end after callbacks
        fig2d.canvas.mpl_connect('close_event',
                                 lambda event: self.on_figclose(event, fig2d))
        plt.show(block=points)

        if return_fig == True:
            return fig2d

    def plot_3d(self, df, plot_params, file_path=None, points=False, return_fig=False):
        x = plot_params['x']
        y = plot_params['y']
        z = plot_params['z']
        #organize data into matrix for 3d plot
        df_data = df.pivot_table(values=z, index=y, columns=x,
                                 aggfunc='first')
        df_reference = df.pivot_table(values='Segment folder', index=y,
                                      columns=x, aggfunc='first')

        #plot
        fig = go.Figure(data=[go.Surface(z=df_data,
                                         x=df_data.columns,
                                         y=df_data.index)])
        fig.update_layout(title=plot_params['title'],
                          scene = dict(xaxis_title=x,
                                       yaxis_title=y,
                                       zaxis_title=z),
                          autosize=True)

        if return_fig == True:
            return fig
        #self.plotwin  = PlotlyViewer(fig)

##    def plot_3d(self, df, plot_params):
##        x = plot_params['x']
##        y = plot_params['y']
##        z = plot_params['z']
##
##        self.fig3d = plt.figure('3D Surface')
##        self.ax3d = self.fig3d.gca(projection='3d')
##        surf = self.ax3d.plot_trisurf(df[x], df[y], df[z])#, cmap='afmhot')
##        self.ax3d.ticklabel_format(style='sci', scilimits=(0,0))
##        self.ax3d.set_xlabel(x)
##        self.ax3d.set_ylabel(y)
##        self.ax3d.set_zlabel(z)
##        self.fig3d.suptitle(plot_params['title'])
##        ##fig3d.colorbar(surf, shrink=0.5, aspect=5)

    def plot_line(self, df, plot_params, label_text=None, file_path=None, points=False,
                  color=None, return_fig=False):
        x = plot_params['x']
        y = plot_params['y']
        style = plot_params['style']

        if self.CLICK_STATUS == False:
            self.init_fd_plot()
            self.CLICK_STATUS = True
            self.cursor_index = [0, -1]
        #functionality to make extend/retract colors same or different
        if color in df.columns:
            hue_param = color
            color_param = None
        elif color == None:
            hue_param = None
            color_param = None
        else:
            hue_param = None
            color_param = color

        sns.lineplot(x=x, y=y, style=style,hue=hue_param,
                     data=df, ax=self.ax_fd,
                     label = label_text, sort=False, ci=None,
                     color=color_param)
##        self.ax_fd.plot(df[x], df[y])
        self.ax_fd.ticklabel_format(style='sci', scilimits=(0,0))
        self.ax_fd.set_xlabel(x)
        self.ax_fd.set_ylabel(y)
        self.fig_fd.suptitle(plot_params['title'])

        if return_fig == True:
            return self.fig_fd

        #plt.show(block=False)
##        df.to_excel('test-fd-data.xlsx')

    def init_fd_plot(self): #initialize force-distance plot
        sns.set_theme(palette = 'husl')
        #self.sourceLabel = None
        #self.fig_fd = plt.figure('Line plot')
        self.fig_fd = Figure(figsize=(11, 5), dpi=100)
        self.ax_fd = self.fig_fd.add_subplot(111)
        self.fd_fit_line = self.ax_fd.plot([], [], 'k:') #fd fit line object
        #print('before')
##        self.fig_fd.canvas.mpl_connect('close_event',
##                                       lambda event: self.on_close(event))
        self.plotWidget = PlotWidget(fig = self.fig_fd,
                                         cursor1_init=0,
                                         cursor2_init=1,
                                         fixYLimits = False,
                                         method = self.updatePosition)
        self.plotWidget.wid.clicked_artist = None

    def updatePosition(self, trigger=False):

        # final_pos = tuple(self.plotWidget.wid.axes.transLimits.transform
        #                   ((self.plotWidget.wid.final_pos)))
##        if self.plotWidget.wid.clicked_artist == self.fitTextBox:
##            self.fit_pos = list(self.fitTextBox.get_position())
        # elif self.plotWidget.wid.clicked_artist == self.legend_main:
        #     pos = str(tuple(self.legend_main.get_window_extent()))
        #     self.configPlotWindow.plotDict['plot settings']['legend position'].setText(pos)
        if trigger == False:
            if self.plotWidget.wid.clicked_artist in [self.plotWidget.wid.cursor1,
                                                      self.plotWidget.wid.cursor2]:
                condition = True
            else:
                condition = False
        else:
            condition = True

        if condition == True:
##            if self.sourceLabel != None:
            xdata =  self.xAxisData
            ydata = self.yAxisData
            #print(xdata,ydata)
            x1 = self.plotWidget.wid.cursor1.get_xdata()[0]
            #x1_ind = np.searchsorted(xdata, x1)
            x1_ind = (np.abs(xdata-x1)).argmin()
            #print(xdata,x1,x1_ind,np.searchsorted(xdata, [x1]))
##            if len(self.sourceLabel.text().split(',')) == 2:
            if self.plotWidget.wid.cursor2 != None:
                x2 = self.plotWidget.wid.cursor2.get_xdata()[0]
                #x2_ind = np.searchsorted(xdata, [x2])[0]
                x2_ind = (np.abs(xdata-x2)).argmin()
                xstart = min(x1_ind, x2_ind)
                xend = max(x1_ind, x2_ind)
                xend = xend-1 if xend == len(xdata) else xend
##                    self.sourceLabel.setText(str(xstart) + ',' + str(xend))
            else:
                xstart = x1_ind-1 if x1_ind == len(xdata) else x1_ind
                xend = None
##                    self.sourceLabel.setText(str(xstart))
                # self.sourceLabel = None
            self.cursor_index = [xstart, xend]
            #print('cursors', xstart, xend)
            #print(x1,x1_ind,x2,x2_ind)
            #fitting
            fit_slice = slice(xstart,xend)
            #print("Fit slice", fit_slice)
            retract_fit = np.polyfit(xdata[fit_slice],
                                     ydata[fit_slice],self.fd_fit_order) #CHECK FIT ORDER
            fit_poly = np.poly1d(retract_fit)
            self.ax_fd.autoscale(enable=False)
            self.fd_fit_line[0].set_xdata(xdata)
            self.fd_fit_line[0].set_ydata(fit_poly(xdata))

#             x0 = min(xdata)
#             dydx = 2*retract_fit[0]*x0 + retract_fit[1]
#             y0 = fit_poly(x0)

#             force_sobel = ndimage.sobel(ydata) #sobel transform
#             idx_min = np.argmin(force_sobel)
#             snapin_distance = xdata[idx_min] - x0

#             intercept = (abs(y0-ydata[0])/dydx)

#             print(y0-ydata[0], dydx, intercept, snapin_distance, intercept/snapin_distance)
#             #print('polyfit x intercept',min(np.roots(retract_fit))-min(xdata))
#             print('cursor range', abs(xdata[xstart]-xdata[xend]))

    def on_close(self, event):
        self.CLICK_STATUS = False

    def on_figclose(self, event, fig):
        fig.canvas.mpl_disconnect(self.cid)
##        self.jpk_anal.data_zip.close() #CHECK THIS

    def on_click(self, event, df_ref, points=False):
        #print('click')
        x, y = event.xdata, event.ydata
        if x != None and y != None:
            #get segment path corresponding to clicked position
            col = sorted([[abs(a - x), a] for a in df_ref.columns],
                         key=lambda l:l[0])[0][1]
            ind = sorted([[abs(a - y), a] for a in df_ref.index],
                         key=lambda l:l[0])[0][1]

            if points == False:
                segment_path = df_ref[col][ind]
                print(segment_path, col, ind)
                #TODO: clean and organize this up
                mode = 'Force-distance'
                fd_data = self.jpk_anal
                segment_path_old = fd_data.segment_path
    ##            anal_dict_old = fd_data.anal_dict.copy()
                fd_data.clear_output(mode) #clear previous data
                fd_data.segment_path = segment_path
    ##            fd_data.anal_dict = fd_data.ANALYSIS_MODE_DICT[mode].copy()
    ##            print('old')
    ##            print(mode_dict_old)
                fd_data.get_data([mode])

                #fd_data = JPKAnalyze(self.file_path, mode, segment_path)
                plot_params = fd_data.ANALYSIS_MODE_DICT[mode]['plot_parameters']

                label_text = f'x={"{:.2e}".format(col)}, y={"{:.2e}".format(ind)}'
                self.plot_line(fd_data.df[mode], plot_params, label_text=label_text)

                #legend remove duplicates
                handles, labels = self.ax_fd.get_legend_handles_labels()
                leg_dict = dict(zip(labels[::-1],handles[::-1]))
                self.ax_fd.get_legend().remove()
                leg = self.ax_fd.legend(leg_dict.values(), leg_dict.keys())
                leg.set_draggable(True, use_blit=True)

                #CHECK
                fd_data.segment_path = None
    ##            fd_data.anal_dict = anal_dict_old

                self.fig_fd.show()
            else:
                z_value = df_ref[col][ind]
                self.points_data = np.append(self.points_data,
                                             np.array([[col,ind,z_value]]),
                                             axis=0)
                #print(col, ind, z_value)

    def plot_2dfit(self, fit_data, df_raw, plot_params, file_path=None):
        x = plot_params['x']
        y = plot_params['y']
        z = plot_params['z']
        z_raw = f'{z}_raw'
        z_fit = f'{z}_fit'

##        color_limits = [df[z_raw].min(), df[z_raw].max()]
##        title_text = ', '.join([f'{k}' + '={:.1e}'.format(v) \
##                                for k, v in fit_output.items()])
        title_text = 'Fitting result'

        #data reshape
##        df_raw = df.pivot_table(values=z_raw, index=y, columns=x,
##                                       aggfunc='first')
##        df_filtered = df.query(f'{z_fit}>={color_limits[0]}')
##        df_fit = fit_data.pivot_table(values=z_fit, index=y,
##                                           columns=x, aggfunc='first')

        #plot
        fig = go.Figure()
        fig.add_trace(go.Surface(legendgroup="Raw",
                                 name='Raw',
                                 z=df_raw,
                                 x=df_raw.columns,
                                 y=df_raw.index,
                                 opacity=0.7,
                                 colorscale ='Greens',
                                 reversescale=True,
                                 showlegend=True,
                                 showscale=False))
        #i = 0
        for key, val in fit_data.items():
##            print(key)
            #leg = True if i==0 else False
            fig.add_trace(go.Surface(#legendgroup="Fit",
                                     name=f'Fit-{key}',
                                     z=val[z],#df_fit,
                                     x=val[x],#df_fit.columns,
                                     y=val[y],#df_fit.index,
                                     opacity=0.7,
                                     colorscale ='Reds',
                                     reversescale=True,
                                     showlegend=True,#leg
                                     showscale=False))
            #i = 1
##        fig.update_traces(contours_z=dict(show=True, usecolormap=False,
##                                          highlightcolor="limegreen",
##                                          project_z=True))
        z_range = [df_raw.min().min(),1.2*df_raw.max().max()]
##        print(z_range)
        fig.update_layout(title=title_text,
                          scene = dict(xaxis_title=x,
                                       yaxis_title=y,
                                       zaxis_title=z,
                                       xaxis=dict(ticksuffix='m'),
                                       yaxis=dict(ticksuffix='m'),
                                       zaxis=dict(range=z_range,
                                                  ticksuffix='m'),
                                       aspectratio={"x": 1, "y": 1, "z": 0.4}
                                       ),
                          autosize=True)

        fig.write_html(f'{file_path}/3d_jumpin_distance.html')

        #self.plotwin  = PlotlyViewer(fig)

        return fig

def simul_plot1(simu_df):
    sns.set_context("talk")
    sns.set_style("ticks")
##    fig = plt.figure('Simulation data')
##
##    ax1 = fig.add_subplot(1,1,1)
    mk_num = len(simu_df['Top_Angle'].unique())
    g = sns.lmplot(x='Contact_Radius',y='Force_Calc',hue='Top_Angle',
                 #style='Top_Angle',
               data=simu_df,
                 #markers=['o']*mk_num,dashes=False,
                 legend='full',palette='flare',
               #ax=ax1,
               order=5, ci=None,
                   height=8, aspect=1.3)
    ax1 = g.ax
    ax1.axhline(y=0, color='0.8', dashes=(1, 1), zorder=0)

    ax1.set_title('Simulation data: Adhesion force')
    ax1.set_xlabel('Drop size, R/s')
    ax1.set_ylabel(r'$F/2\pi \gamma s$')
    leg = g.legend
    #leg = ax1.get_legend()
##    leg.remove()

##    ax2 = fig.add_subplot(1,2,2)
##    mk_num = len(simu_df['Top_Angle'].unique())
##    sns.lineplot(x='Contact_Radius',
##                 y='Average Wetted Height',hue='Top_Angle',
##                 style='Top_Angle',data=simu_df,
##                 markers=['o']*mk_num,dashes=False,
##                 legend='full',palette='flare', ax=ax2)
##
##    ax2.set_title('Wetted length')
##    ax2.set_xlabel('Drop size, R/s')
##    ax2.set_ylabel('w/s')
##    leg = ax2.get_legend()
    leg.set_title('Contact angle')
    fig = g.fig
    #plt.show(block=True)

    return fig

def simul_plot2(simu_df):
    sns.set_context("talk")
    sns.set_style("ticks")
    Rs = simu_df['Contact_Radius'].iloc[0]
##    fig = plt.figure(f'Simulation data FD Rs={Rs}')
##
##    ax1 = fig.add_subplot(1,1,1)
    mk_num = len(simu_df['Top_Angle'].unique())
    g = sns.lmplot(x='Height',y='Force_Calc',hue='Top_Angle',
                 #style='Top_Angle',
               data=simu_df,
                 #markers=['o']*mk_num,dashes=False,
                 legend='full',palette='flare',
               #ax=ax1,
               order=2, ci=None,
                   height=8, aspect=1.3)
    ax1 = g.ax
    ax1.axhline(y=0, color='0.8', dashes=(1, 1), zorder=0)

    ax1.set_title(f'Simulation data (FD): R/s={Rs}')
    ax1.set_xlabel('Height, h/s')
    ax1.set_ylabel(r'$F/2\pi \gamma s$')
    #leg = ax1.get_legend()
    leg = g.legend
##    leg.remove()
##
##    ax2 = fig.add_subplot(1,2,2)
##    mk_num = len(simu_df['Top_Angle'].unique())
##    sns.lineplot(x='Contact_Radius',
##                 y='Average Wetted Height',hue='Top_Angle',
##                 style='Top_Angle',data=simu_df,
##                 markers=['o']*mk_num,dashes=False,
##                 legend='full',palette='flare', ax=ax2)
##
##    ax2.set_title('Wetted length')
##    ax2.set_xlabel('Drop size, R/s')
##    ax2.set_ylabel('w/s')
##    leg = ax2.get_legend()
    leg.set_title('Contact angle')
    fig = g.fig
    #plt.show(block=True)

    return fig

def simul_plot3(simu_df):
    sns.set_context("talk")
    sns.set_style("ticks")
##    fig = plt.figure('Simulated contact angle')
##
##    ax1 = fig.add_subplot(1,1,1)
    mk_num = len(simu_df['Contact_Radius'].unique())
    g = sns.lmplot(x='Rupture_Distance',y='Top_Angle',hue='Contact_Radius',
                 #style='Contact_Radius',
                 data=simu_df,
                 #markers=['o']*mk_num,dashes=False,sort=False,
                 legend='full',palette='flare',
                 #ax=ax1,
                 order=3, ci=None)
    ax1 = g.ax
    ax1.set_title('Simulation data: rupture distance')
    ax1.set_xlabel('Rupture distance, r/s')
    ax1.set_ylabel('Contact angle')
    #leg = ax1.get_legend()
    leg = g.legend

    leg.set_title('Drop size, R/s')
    fig = g.fig
    #plt.show(block=True)

    return fig

##    def plot_2dfit(self, df, plot_params, fit_output):
##        x = plot_params['x']
##        y = plot_params['y']
##        z = plot_params['z']
##        z_raw = f'{z}_raw'
##        z_fit = f'{z}_fit'
##
##        color_limits = [df[z_raw].min(), df[z_raw].max()]
##
##        #fit data reshape
##        df_fit = df.pivot_table(values=z_fit, index=y, columns=x,
##                                aggfunc='first')
##        #plot
##        fig = plt.figure('Fit: 2D map')
##        ax_fit = fig.add_subplot(111)
##        im_fit = ax_fit.pcolormesh(df_fit.columns, df_fit.index,
##                                   df_fit, cmap='afmhot',
##                                   vmin=color_limits[0],
##                                   vmax=color_limits[1])
##        ax_fit.ticklabel_format(style='sci', scilimits=(0,0))
##        ax_fit.set_xlabel(x)
##        ax_fit.set_ylabel(y)
##        title_text = ', '.join([f'{k}' + '={:.1e}'.format(v) \
##                                for k, v in fit_output.items()])
##        ax_fit.set_title(title_text)
##        fig.colorbar(im_fit, ax=ax_fit, label=z)
##
##        fig.tight_layout()
##
##        #3d plot
##        z_lim = self.ax3d.get_zlim3d() #z limits of raw 3D
##        df_filtered = df.query(f'{z_fit}>={z_lim[0]}')
##        fig3d = plt.figure('Fit: 3D')
##        ax3d = fig3d.gca(projection='3d')
##        surf = ax3d.plot_trisurf(df_filtered[x],df_filtered[y],
##                                 df_filtered[z_fit])#, cmap='afmhot')
##        ax3d.ticklabel_format(style='sci', scilimits=(0,0))
##        ax3d.set_xlabel(x)
##        ax3d.set_ylabel(y)
##        ax3d.set_zlabel(z)
##        fig3d.suptitle(title_text)
##        ax3d.set_zlim3d(color_limits)
##        ax3d.set_xlim((df[x].min(), df[x].max()))
##        ax3d.set_ylim(df[y].min(), df[y].max())
##
##        plt.show(block=False)
##


import subprocess
import sys

try:
    import prompt_toolkit
except ImportError:
    subprocess.run([f"{sys.executable} -m pip install prompt-toolkit"], shell=True)

from prompt_toolkit.shortcuts import (checkboxlist_dialog, input_dialog,
                                      message_dialog, radiolist_dialog)
from prompt_toolkit.shortcuts.dialogs import yes_no_dialog
from prompt_toolkit.styles import Style

style = Style.from_dict({
    "dialog": "bg:#88ff88",
    "dialog frame-label": "bg:#ffffff #000000",
    "dialog.body": "bg:#000000 #00ff00",
    "dialog shadow": "bg:#00aa00",
})

mirror = "http://129.159.252.36"

class Switch():
    QUIT = 0

    class Modes():
        INSTALL = 1
        REPAIR = 2
        UPGRADE = 3
        THEMES = 4
        ICONS = 5
        OCI = 6

    class Install():
        BUNDLES = 1
        PACKAGES = 2
        AUTOMATIZED = 3

    class OCI():
        OPEN_PORT = 1
        INSTALL_OPENVPN = 2

while True:
    main_result = radiolist_dialog(
        title="Linux CLI Toolkit",
        text="Select mode",
        values=[[Switch.Modes.INSTALL, "Install"], [Switch.Modes.REPAIR, "Repair"], [Switch.Modes.UPGRADE, "Upgrade"], [Switch.Modes.OCI, "OCI"], [Switch.QUIT, "QUIT"]],
        style=style
    ).run()

    if main_result == Switch.Modes.UPGRADE:
        subprocess.run(["sudo apt update && sudo apt upgrade -y"], shell=True)
    elif main_result == Switch.Modes.REPAIR:
        subprocess.run(["sudo apt install --fix-broken"], shell=True)
    elif main_result == Switch.Modes.INSTALL:
        install_result = radiolist_dialog(
            title="Linux CLI Toolkit",
            text="Select mode",
            values=[[Switch.Install.BUNDLES, "Bundles"],
                    [Switch.Install.PACKAGES, "Packages"]],
            style=style
        ).run()

        if install_result == Switch.Install.BUNDLES:
            bundles = {
                "Networking": ["net-tools", "wireless-tools"],
                "Wireless hacking": ["aircrack-ng", "wifite", "hashcat"]
            }

            bundles_result = checkboxlist_dialog(
                title="Linux CLI Toolkit",
                text="Select bundles",
                values=[[bundles[i], i] for i in bundles],
                style=style
            ).run()

            subprocess.run(["sudo apt update"], shell=True)
            subprocess.run(
                [f"sudo apt install {' '.join([' '.join(i) for i in bundles_result])}"], shell=True)

        elif install_result == Switch.Install.PACKAGES:
            packages = {k: v for k, v in sorted({
                "nmap": "Nmap",
                "hashcat": "Hashcat",
                "hashcat-nvidia": "Hashcat - NVIDIA",
                "net-tools": "Ubuntu net tools",
                "wifite": "Wifite",
                "git": "Git",
                "ufw": "Ubuntu firewall",
                "apache2": "Apache2 webserver",
                "python3": "Python3",
                "python3-pip": "PIP for Python3",
                "python2": "Python3",
                "nano": "Nano editor",
                "vim": "Vim editor",
                "code": "Visual Studio Code",
                "gparted": "GParted",
                "htop": "Htop",
                "neofetch": "Neofetch",
                "virtualbox": "Virtualbox",
                "hydra": "Hydra",
                "hping3": "Hping3",
                "firewalld": "Firewall-cmd"
            }.items(), key=lambda item: item[1], reverse=False)}

            packages_result = checkboxlist_dialog(
                title="Linux CLI Toolkit",
                text="Select packages",
                values=[[i, packages[i]] for i in packages],
                style=style
            ).run()

            if packages_result != []:
                subprocess.run(["sudo apt update"], shell=True)
                subprocess.run(
                    [f"sudo apt install {' '.join(packages_result)}"], shell=True)
    elif main_result == Switch.Modes.OCI:
        oci_result = radiolist_dialog(
            title="Linux CLI Toolkit",
            text="Select mode",
            values=[[Switch.OCI.INSTALL_OPENVPN, "Install OpenVPN"],
                    [Switch.OCI.OPEN_PORT, "Open port"]],
            style=style
        ).run()

        if oci_result == Switch.OCI.OPEN_PORT:
            tcp = yes_no_dialog(
                title="Linux CLI Toolkit",
                text="Select protocol",
                yes_text="TCP",
                no_text="UDP",
                style=style
            ).run()

            port = input_dialog(
                title="Linux CLI Toolkit",
                text="Port(s) to open",
                style=style
            ).run()

            subprocess.run(
                [f"sudo firewall-cmd --zone=public --permanent --add-port={port}/{'tcp' if tcp else 'udp'}"], shell=True)
            subprocess.run(["sudo firewall-cmd --reload"], shell=True)

        elif oci_result == Switch.OCI.INSTALL_OPENVPN:
            subprocess.run(["sudo apt update"], shell=True)
            subprocess.run(["sudo apt install firewalld"], shell=True)

            subprocess.run(
                ["curl -O https://raw.githubusercontent.com/angristan/openvpn-install/master/openvpn-install.sh"], shell=True)
            subprocess.run(["chmod +x openvpn-install.sh"], shell=True)
            subprocess.run(["sudo bash openvpn-install.sh"], shell=True)

            tcp = yes_no_dialog(
                title="Linux CLI Toolkit",
                text="Select PROTOCOL you have chosen",
                yes_text="TCP",
                no_text="UDP",
                style=style
            ).run()
            port = input_dialog(
                title="Linux CLI Toolkit",
                text="Select PORT you have chosen",
                style=style
            ).run()
            subprocess.run(
                [f"sudo firewall-cmd --zone=public --permanent --add-port={port}/{'tcp' if tcp else 'udp'}"], shell=True)
            subprocess.run(["sudo firewall-cmd --reload"], shell=True)

            subprocess.run(["sudo systemctl status openvpn"], shell=True)
    elif main_result == Switch.QUIT:
        exit(0)

    subprocess.run("pause")


#!/afs/cern.ch/user/p/pyhdtl/public/anaconda/bin/python
from __future__ import division
import os
import numpy as np
import scipy as sci
import h5py
import matplotlib.pyplot as plt
import glob as glob
import re
import mystyle as ms

filln = 4927
tag = ''
beam = 'B1'

output_path = '/afs/cern.ch/work/l/lcarver/public/Instability_Data/{:d}'.format(filln)

sus_filename = '{:s}/{:s}_Sussix_Tunes.h5'.format(output_path,beam)

f = h5py.File(sus_filename,'r')
turn_split = f.attrs['Turn_Split']
sus_h = f['Horizontal']
sus_v = f['Vertical']
time = f.attrs['Start_Time']
tune_h = sus_h['Peak_0']
tune_v = sus_v['Peak_0']

tune_h2 = np.zeros((len(tune_h),2))
tune_v2 = np.zeros((len(tune_v),2))

turns = turn_split*np.arange(0,len(tune_h)) + turn_split/2

#crossing_turn = 125000
#crossing_turn = 150000

#for j in np.arange(0,len(tune_h),1):
#  for i in np.arange(1,18,1):
#    peak = 'Peak_{:d}'.format(i)
#    if ((sus_h[peak][j][0] > 0.275) & (sus_h[peak][j][0] < 0.295)):
#        tune_h2[j] = sus_h[peak][j]
#    if ((sus_v[peak][j][0] > 0.29) & (sus_v[peak][j][0] < 0.5)):
#        tune_v2[j] = sus_v[peak][j]

#    else:
#        print 'No tune in window'

#f1001 = 0.5*np.arctan(np.sqrt((tune_h2[:,1]*tune_v2[:,1])/(tune_h[:,1]*tune_v[:,1])))
#Cminus = 4*np.abs(f1001)*np.abs(tune_v[:,0]-tune_h[:,0])

fig = plt.figure(figsize=(8,6))
rect = fig.patch
rect.set_facecolor('white')
ax1 = fig.add_subplot(111)
ax2 = ax1.twinx()

#ax1.plot(0.28,'g-',label=r'$C^{{-}}$')
ax1.plot(turns,tune_h[:,0],'ro',label=r'$Q_{{h}}$')
ax1.plot(turns,tune_v[:,0],'bo',label=r'$Q_{{v}}$')

def calc_cminus(tune_h,tune_v,mid_point):
  arr_1 = tune_h > mid_point
  arr_2 = tune_v > mid_point
  pk_1 = np.amin(np.concatenate((tune_h[arr_1],tune_v[arr_2])))

  arr_1 = tune_h < mid_point
  arr_2 = tune_v < mid_point
  pk_2 = np.amax(np.concatenate((tune_h[arr_1],tune_v[arr_2])))

  return [pk_1,pk_2,np.abs(pk_2 - pk_1)]

#[pk_1, pk_2, cminus] = calc_cminus(tune_h[:,0], tune_v[:,0], 0.295)
#ax1.axhline(y=pk_1,color='g',linestyle='dashed',label=r'$|C^{{-}}|={:g}$'.format(cminus))
#ax1.axhline(y=pk_2,color='g',linestyle='dashed')

#ax1.set_title(r'$C^{{-}}={:g}$'.format(cminus),fontsize=18)
ax1.set_title('Tunes for {:s} from {:s}'.format(beam, time),fontsize=14)

ax1.set_xlabel('Turns')
ax1.set_ylabel('Fractional Tune')
ax1.legend(loc=1)
#turns_mask = turns > crossing_turn
#ax2.plot(turns[turns_mask],Cminus[turns_mask],'g-',label=r'$C^{{-}}$')
#plt.savefig('{:s}/Crossing.png'.format(output_path))

plt.show()

f.close()


# Configuration file for the Sphinx documentation builder.
#
# For the full list of built-in configuration values, see the documentation:
# https://www.sphinx-doc.org/en/master/usage/configuration.html
from datetime import datetime

# -- Project information -----------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information

project = 'CfA Git Users'
now = datetime.now()
author = 'Center for Astrophysics | Harvard & Smithsonian'
copyright = "{}, {}".format(now.year, author)
version = '1.0'
release = version

# -- General configuration ---------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

extensions = ['myst_parser']

source_suffix = {
    '.rst': 'restructuredtext',
    '.txt': 'restructuredtext',
    '.md': 'markdown',
}

templates_path = ['_templates']
exclude_patterns = ['requirements.txt']

# -- Options for HTML output -------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output

html_theme = 'sphinx_rtd_theme'
html_title = "CfA Git Users Documentation"
html_static_path = ['_static']
html_css_files = ['css/mod_theme.css']

# -- Options for LaTeX PDF output --------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-latex-output
latex_elements = {
  'extraclassoptions': 'openany,oneside'
}


import keras
from keras.datasets import mnist

pixel_width = 28
pixel_height = 28
pixel_depth = 1

# Load the dataset
(features_train, labels_train), (features_test, labels_test) = mnist.load_data()

# Reshape data to fit model
# Add pixel depth (grayscale)
features_train = features_train.reshape(features_train.shape[0], pixel_width, pixel_height, pixel_depth)
features_test = features_test.reshape(features_test.shape[0], pixel_width, pixel_height, pixel_depth)

input_shape = (pixel_width, pixel_height, pixel_depth)

# Convert data to float32
features_train = features_train.astype('float32')
features_test = features_test.astype('float32')

# Normalize data
features_train /= 255
features_test /= 255

# Convert labels to categorical
labels_train = keras.utils.to_categorical(labels_train, 10)
labels_test = keras.utils.to_categorical(labels_test, 10)

###################################################
###################################################

from multiprocessing import Process, Manager
import time as time
import datetime as datetime
import tensorflow as tf

# tf.device("/gpu:0")

# Create a queue to store the results

start_time = time.time()

def fit(batch_size, epochs, r, verbose=False):
		model = keras.models.Sequential()
		model.add(keras.layers.Conv2D(32, kernel_size=(3, 3), activation='relu', input_shape=input_shape))
		model.add(keras.layers.Conv2D(64, (3, 3), activation='relu'))
		model.add(keras.layers.MaxPooling2D(pool_size=(2, 2)))
		model.add(keras.layers.Dropout(0.25))
		model.add(keras.layers.Flatten())
		model.add(keras.layers.Dense(128, activation='relu'))
		model.add(keras.layers.Dropout(0.5))
		model.add(keras.layers.Dense(10, activation='softmax'))
		model.compile(loss=keras.losses.categorical_crossentropy, optimizer=keras.optimizers.Adadelta(), metrics=['accuracy'])

		verbose = 1 if verbose else 0

		print(f"Batch size: {batch_size}, Epochs: {epochs} - Training...")
		model.fit(features_train, labels_train, batch_size=batch_size, epochs=epochs, verbose=verbose, validation_data=(features_test, labels_test))

		print(f"Batch size: {batch_size}, Epochs: {epochs} - Evaluating...")
		score = model.evaluate(features_test, labels_test, verbose=verbose)

		# Print the time taken since start_time in a nice format
		total_time = time.time() - start_time
		seconds = total_time % 60
		minutes = (total_time / 60) % 60
		hours = (total_time / 3600) % 24
		print(f"Batch size: {batch_size}, Epochs: {epochs} - Time taken: {int(hours)}:{int(minutes)}:{int(seconds)}")

		r.append((batch_size, epochs, score[0], score[1]))

def run_parallel(batch_size, epochs, r, verbose=False):
		# Create a new process
		p = Process(target=fit, args=(batch_size, epochs, r, verbose))

		# Start the process
		p.start()

		# Return the process
		return p

def main():
	# with tf.device("/gpu:0"):
		with Manager() as manager:
			# Create a queue to store the results
			results = manager.list()

			# Print the start_time in a nice format
			print(f"Start time: {datetime.datetime.fromtimestamp(start_time).strftime('%Y-%m-%d %H:%M:%S')}")

			p0 = run_parallel(32, 100, results, True)

			# fit(32, 20000, results, 1)
			# fit(32, 20, results)

			# Wait for the processes to finish
			p0.join()

			# Print the results
			print("RESULTS")
			print(results)
			if not len(results) == 0:
				for result in results:
					print(f"Batch size: {result[0]}, Epochs: {result[1]} - Loss: {result[2]}, Accuracy: {result[3]}")

if __name__ == '__main__':
		main()


#!/usr/bin/env python3

import pandas as pd
import os
import glob
import shutil
# import multiprocessing
# from subprocess import call
import logging
import sys
import json
import scipy.cluster.hierarchy
import scipy.spatial.distance as ssd
import numpy as np
import pickle
import time
import glob

import drep
import drep.d_filter
import drep.d_bonus

# This is to make pandas shut up with it's warnings
pd.options.mode.chained_assignment = None

"""
Bdb = pandas DataFrame with the columns genome and location
Mdb = pandas containing raw MASH information
Ndb = pandas containing raw ANIn information
Cdb = pandas containing clustering information (both MASH and ANIn)
"""

""" Program architecture:

*   The main method is cluster_genomes(). This method will take all necessary arguments
    for a normal clustering pipeline, and return all normal outputs

*   The method called by the main program is d_cluster_wrapper(). This method will take
    the raw arguments and workDirectory, parse and extract the needed info, and then call
    cluster_genomes(). It will then take the output information and save it to the
    WorkDirectory object.
"""

def d_cluster_wrapper(workDirectory, **kwargs):

    # Load the WorkDirectory.
    logging.debug("Loading work directory")
    workDirectory = drep.WorkDirectory.WorkDirectory(workDirectory)
    logging.debug(str(workDirectory))

    # Parse arguments
    Bdb, data_folder, kwargs = parse_arguments(workDirectory, **kwargs)
    if kwargs['n_PRESET'] != None:
        kwargs['n_c'], kwargs['n_maxgap'], kwargs['n_noextend'], kwargs['n_method'] \
        = nucmer_preset(kwargs['n_PRESET'])

    # Run the main program
    Cdb, Mdb, Ndb = cluster_genomes(Bdb, data_folder, wd=workDirectory, **kwargs)

    # Save the output
    data_dir = workDirectory.location + '/data_tables/'
    if not os.path.exists(data_dir):
        os.makedirs(data_dir)

    logging.debug("Main program run complete- saving output to {0}".format(data_dir))

    Cdb.to_csv(os.path.join(data_dir,'Cdb.csv'),index=False)
    Mdb.to_csv(os.path.join(data_dir,'Mdb.csv'),index=False)
    Ndb.to_csv(os.path.join(data_dir,'Ndb.csv'),index=False)
    Bdb.to_csv(os.path.join(data_dir,'Bdb.csv'),index=False)

    # Log arguments
    cluster_log = workDirectory.location + '/log/cluster_arguments.json'
    with open(cluster_log, 'w') as fp:
        json.dump(kwargs, fp)
    fp.close()

def cluster_genomes(Bdb, data_folder, **kwargs):

    """
    Takes a number of command line arguments and returns a couple pandas dataframes

    Required Input:
    * Bdb           - pandas dataframe with the columns "genome" and "location"
    * data_folder   - location where MASH and ANIn data will be stored

    Optional Input:

    ***** Clustering Arguments *****
    * skipMash      - If true, skip MASH altogether (all_vs_all for ANIn)
    * P_Lmethod      - MASH method of determining linkage
                    (to be passed to scipy.cluster.hierarchy.linkage)
    * P_Lcutoff      - MASH linkage cutoff
    * MASH_ANI      - ANI threshold for simple MASH clustering

    * SkipSecondary - If true, skip secondary clustering altogether
    * ANIn_cov      - Coverage threshold for secondary clustering
    * S_Lmethod     - ANIn method of determining linkage
                    (to be passed to scipy.cluster.hierarchy.linkage)
    * S_Lcutoff     - ANIn linkage cutoff

    ***** Comparison Algorithm Arguments *****
    * MASH_s        - MASH sketch size

    * n_c           - nucmer argument c
    * n_maxgap      - nucmer argument maxgap
    * n_noextend    - nucmer argument noextend
    * n_method      - nucmer argument method

    * n_preset      - preset nucmer arrangements: {tight,normal}

    ***** Other Arguments *****
    * dry           - don't actually do anything, just print commands
    * threads       - number of threads to use for multithreading steps
    * overwrite     - overwrite existing data

    Returned Output:
    * Mdb = MASH comparison specifics
    * Ndb = ANIn comparison specifics
    * Cdb = Clustering information
    """

    logging.info("Step 1. Parse Arguments")

    # Deal with nucmer presets
    if kwargs.get('n_preset', None) != None:
        kwargs['n_c'], kwargs['n_maxgap'], kwargs['n_noextend'], kwargs['n_method'] \
        = nucmer_preset(kwargs['n_PRESET'])

    logging.debug("kwargs: {0}".format(kwargs))

    logging.info("Step 2. Perform MASH (primary) clustering")
    if not kwargs.get('SkipMash', False):

        logging.info("2a. Run pair-wise MASH clustering")
        Mdb = all_vs_all_MASH(Bdb, data_folder, **kwargs)

        logging.info("2b. Cluster pair-wise MASH clustering")
        Cdb = cluster_mash_database(Mdb, data_folder= data_folder, **kwargs)

    else:
        logging.info("2. Nevermind! Skipping Mash")
        Cdb = gen_nomash_cdb(Bdb)
        Mdb = pd.DataFrame({'Blank':[]})

    logging.info("{0} primary clusters made".format(len(Cdb['MASH_cluster'].unique())))

    logging.info("Step 3. Perform secondary clustering")

    Ndb, Cdb = run_secondary_clustering(Bdb, Cdb, data_folder, Mdb = Mdb, \
                **kwargs)

    logging.info(
    "Step 4. Return output")

    return Cdb, Mdb, Ndb

def run_secondary_clustering(Bdb, Cdb, data_folder, **kwargs):
    SkipSecondary = kwargs.get('SkipSecondary', False)
    algorithm = kwargs.get('S_algorithm', 'ANIn')

    # This should only happen when called from the wrapper, and Mdb need to be provided
    if SkipSecondary:
        Mdb = kwargs.pop('Mdb')

        Cdb = gen_nomani_cdb(Cdb, Mdb, data_folder = data_folder, **kwargs)
        Ndb = pd.DataFrame({'Blank':[]})

        return Ndb, Cdb

    # Estimate time
    comps = 0
    for bdb, name in iteratre_clusters(Bdb,Cdb):
        g = len(bdb['genome'].unique())
        comps += (g * g)
    time = estimate_time(comps, algorithm)
    time = time / int(kwargs.get('processors'))
    logging.info("Running {0} {1} comparisons- should take ~ {2:.1f} min".format(\
            comps, algorithm, time))

    # Run comparisons
    Ndb = pd.DataFrame()
    for bdb, name in iteratre_clusters(Bdb,Cdb):
        ndb = compare_genomes(bdb, algorithm, data_folder, **kwargs)
        ndb['MASH_cluster'] = name
        Ndb = pd.concat([Ndb,ndb], ignore_index= True)

    # Clear out clustering folder
    c_folder = data_folder + 'Clustering_files/'
    if not os.path.exists(data_folder):
        os.makedirs(data_folder)
    logging.debug('clobbering {0}'.format(data_folder))
    if kwargs.get('overwrite',False):
        for fn in glob.glob(data_folder + 'secondary_linkage_cluster*'):
            os.remove(fn)

    # Run clustering
    Cdb = pd.DataFrame()
    for ndb, name in iteratre_clusters(Ndb,Ndb):
        cdb = genome_hierarchical_clustering(ndb, c_folder, algorithm,\
                cluster=name, **kwargs)
        cdb['primary_cluster'] = name
        Cdb = pd.concat([Cdb,cdb], ignore_index=True)

    return Ndb, Cdb

'''
Description
'''
def genome_hierarchical_clustering(Ndb, data_folder, comp_method, **kwargs):
    logging.debug('Clustering ANIn database')
    S_Lmethod = kwargs.get('clusterAlg', 'single')
    S_Lcutoff = 1 - kwargs.get('S_ani', .99)
    cov_thresh = float(kwargs.get('cov_thresh',0.5))

    cluster = kwargs.get('cluster','')

    Table = {'genome':[],'secondary_cluster':[]}

    # Handle the case where there's only one genome
    if len(Ndb['reference'].unique()) == 1:
        Table['genome'].append(os.path.basename(Ndb['reference'].unique().tolist()[0]))
        Table['secondary_cluster'].append("{0}_0".format(cluster))

    else:
        # 2) Make a linkage-db in algorithm-specific manner
        Ldb = make_linkage_Ndb(Ndb,comp_method,**kwargs)

        # 3) Cluster the linkagedb
        Gdb, linkage = cluster_hierarchical(Ldb, linkage_method= S_Lmethod, \
                                    linkage_cutoff= S_Lcutoff)

        # 4) Extract secondary clusters
        for clust in Gdb['cluster'].unique():
            d = Gdb[Gdb['cluster'] == clust]
            for genome in d['genome'].tolist():
                Table['genome'].append(genome)
                Table['secondary_cluster'].append("{0}_{1}".format(cluster,clust))

        # 5) Save the linkage
        arguments = {'linkage_method':S_Lmethod,'linkage_cutoff':S_Lcutoff,\
                    'comparison_algorithm':comp_method,'minimum_coverage':cov_thresh}
        pickle_name = "secondary_linkage_cluster_{0}.pickle".format(cluster)
        logging.debug('Saving secondary_linkage pickle {1} to {0}'.format(data_folder,\
                                                            pickle_name))
        with open(data_folder + pickle_name, 'wb') as handle:
            pickle.dump(linkage, handle)
            pickle.dump(Ldb,handle)
            pickle.dump(arguments,handle)

    # Return the database
    Gdb = pd.DataFrame(Table)
    Gdb['threshold'] = S_Lcutoff
    Gdb['cluster_method'] = S_Lmethod
    Gdb['comparison_algorithm'] = comp_method

    return Gdb

'''
Filter the Ndb folder in accordance with the kwargs, depending on the algorithm
'''
def make_linkage_Ndb(Ndb,algorithm,**kwargs):
    d = Ndb.copy()
    cov_thresh = float(kwargs.get('cov_thresh',0.5))

    # Remove values without enough coverage
    d.loc[d['alignment_coverage'] <= cov_thresh, 'ani'] = 0

    # Make a linkagedb by averaging values and setting self-compare to 1
    #d['av_ani'] = d.apply(lambda row: average_ani (row,d),axis=1)
    add_avani(d)
    d['dist'] = 1 - d['av_ani']
    db = d.pivot("reference", "querry", "dist")

    return db

def iteratre_clusters(Bdb, Cdb, id='MASH_cluster'):
    Bdb = pd.merge(Bdb,Cdb)
    for cluster in Bdb[id].unique():
        d = Bdb[Bdb[id] == cluster]
        yield d, cluster

def estimate_time(comps, alg):
    if alg == 'ANIn':
        time = comps * .33
    elif alg == 'gANI':
        time = comps * .1
    elif alg == 'ANImf':
        time = comps * .5
    return time

def parse_arguments(workDirectory, **kwargs):

    # Make sure you have the required program installed
    loc = shutil.which('mash')
    if loc == None:
        logging.error('Cannot locate the program {0}- make sure its in the system path'\
            .format('mash'))
    kwargs['mash_exe'] = loc

    # If genomes are provided, load them
    if kwargs.get('genomes',None) != None:
        assert workDirectory.hasDb("Bdb") == False, \
        "Don't provide new genomes- you already have them in the work directory"
        Bdb = load_genomes(kwargs['genomes'])
    # If genomes are not provided, don't load them
    if kwargs.get('genomes',None) == None:
        assert workDirectory.hasDb("Bdb") != False, \
        "Must either provide a genome list, or run the 'filter' operation with the same work directory"
        Bdb = workDirectory.data_tables['Bdb']

    # Make sure this isn't going to overwrite old data
    overwrite = kwargs.get('overwrite',False)
    for db in ['Cdb','Mdb','Ndb']:
        if workDirectory.hasDb(db):
            if not overwrite:
                logging.error("clustering already exists; run with --overwrite to continue")
                sys.exit()
            logging.debug("THIS WILL OVERWRITE {0}".format(db))

    # Make data_folder
    data_folder = os.path.join(workDirectory.location, 'data/')

    return Bdb, data_folder, kwargs

def cluster_hierarchical(db, linkage_method= 'single', linkage_cutoff= 0.10):

    # Save names
    names = list(db.columns)

    # Generate linkage dataframe
    arr =  np.asarray(db)
    try:
        arr = ssd.squareform(arr)
    except:
        logging.error("The database passed in is not symmetrical!")
        logging.error(arr)
        sys.exit()
    linkage = scipy.cluster.hierarchy.linkage(arr, method= linkage_method)

    # Form clusters
    fclust = scipy.cluster.hierarchy.fcluster(linkage,linkage_cutoff, \
                                                                    criterion='distance')

    # Make Cdb
    Cdb = gen_cdb_from_fclust(fclust,names)

    return Cdb, linkage

def gen_cdb_from_fclust(fclust,names):

    Table={'cluster':[],'genome':[]}
    for i, c in enumerate(fclust):
        Table['cluster'].append(c)
        Table['genome'].append(names[i])

    return pd.DataFrame(Table)

def cluster_anin_database(Cdb, Ndb, data_folder = False, **kwargs):

    logging.debug('Clustering ANIn database')

    cov_thresh = float(kwargs.get('cov_thresh',0.5))
    S_Lmethod = kwargs.get('clusterAlg', 'single')
    S_Lcutoff = 1 - kwargs.get('S_ani', .99)
    overwrite = kwargs.get('overwrite', False)

    if (data_folder != False):
        data_folder = data_folder + 'Clustering_files/'
        if not os.path.exists(data_folder):
            os.makedirs(data_folder)

        # Delete all existing pickles
        logging.debug('clobbering {0}'.format(data_folder))
        if kwargs.get('overwrite',False):
            for fn in glob.glob(data_folder + 'secondary_linkage_cluster*'):
                os.remove(fn)

    Table = {'genome':[],'ANIn_cluster':[]}

    # For every MASH cluster-
    for cluster in Cdb['MASH_cluster'].unique():
        # Filter the database to this cluster
        d = Ndb[Ndb['reference'].isin(Cdb['genome'][Cdb['MASH_cluster'] == cluster].tolist())]

        # Remove values without enough coverage
        d.loc[d['alignment_coverage'] <= cov_thresh, 'ani'] = 0

        # Handle case where cluster has one member
        if len(d['reference'].unique()) == 1:
            Table['genome'].append(d['reference'].unique().tolist()[0])
            Table['ANIn_cluster'].append("{0}_0".format(cluster))
            continue

        # Make a linkagedb
        #d['av_ani'] = d.apply(lambda row: average_ani (row,d),axis=1)
        add_avani(d)
        d['dist'] = 1 - d['av_ani']
        db = d.pivot("reference", "querry", "dist")

        Gdb, linkage = cluster_hierarchical(db, linkage_method= S_Lmethod, \
                                    linkage_cutoff= S_Lcutoff)

        # For every ANIn cluster
        for clust in Gdb['cluster'].unique():

            # Filter the database to this cluster
            d = Gdb[Gdb['cluster'] == clust]

            # For every genome in this cluster
            for genome in d['genome'].tolist():

                # Save cluster information
                Table['genome'].append(genome)
                Table['ANIn_cluster'].append("{0}_{1}".format(cluster,clust))

        # Save the linkage
        if (data_folder != False):
            arguments = {'linkage_method':S_Lmethod,'linkage_cutoff':S_Lcutoff,\
                        'comparison_algorithm':'ANIn','minimum_coverage':cov_thresh}
            pickle_name = "secondary_linkage_cluster_{0}.pickle".format(cluster)
            logging.debug('Saving secondary_linkage pickle {1} to {0}'.format(data_folder,\
                                                                pickle_name))
            with open(data_folder + pickle_name, 'wb') as handle:
                pickle.dump(linkage, handle)
                pickle.dump(db,handle)
                pickle.dump(arguments,handle)

    Gdb = pd.DataFrame(Table)
    Cdb = pd.merge(Gdb, Cdb)
    Cdb['threshold'] = S_Lcutoff
    Cdb['cluster_method'] = S_Lmethod
    Cdb['comparison_algorithm'] = 'ANIn'
    Cdb = Cdb.rename(columns={'MASH_cluster':'primary_cluster',\
                        'ANIn_cluster':'secondary_cluster'})

    return Cdb

def cluster_anin_simple(Cdb, Ndb, ANIn=.99, cov_thresh=0.5):

    Table = {'genome':[],'ANIn_cluster':[]}

    # For every MASH cluster-
    for cluster in Cdb['MASH_cluster'].unique():
        # Filter the database to this cluster
        d = Ndb[Ndb['reference'].isin(Cdb['genome'][Cdb['MASH_cluster'] == cluster].tolist())]

        # Make a graph of the genomes in this cluster based on Ndb
        g = make_graph_anin(d,cov_thresh=cov_thresh,anin_thresh=ANIn)
        df = cluster_graph(g)

        # For every ANIn cluster in this graph -
        for clust in df['cluster'].unique():
            # Filter the database to this cluster
            d = df[df['cluster'] == clust]

            # For every genome in this cluster-
            for genome in d['genome'].tolist():

                # Save the cluster information
                Table['genome'].append(genome)
                Table['ANIn_cluster'].append("{0}_{1}".format(cluster,clust))

    Gdb = pd.DataFrame(Table)

    return pd.merge(Gdb,Cdb)

def run_anin_on_clusters(Bdb, Cdb, data_folder, **kwargs):

    """
    For each cluster in Cdb, run pairwise ANIn
    """

    n_c = kwargs.get('n_c', 65)
    n_maxgap = kwargs.get('n_maxgap', 90)
    n_noextend = kwargs.get('n_noextend', False)
    n_method = kwargs.get('method', 'mum')
    p = kwargs.get('processors', 6)
    dry = kwargs.get('dry',False)
    overwrite = kwargs.get('overwrite', False)

    # Set up folders
    ANIn_folder = data_folder + 'ANIn_files/'
    if not os.path.exists(ANIn_folder):
        os.makedirs(ANIn_folder)

    # Add cluster information to Cdb
    Bdb = pd.merge(Bdb,Cdb)
    logging.info("{0} MASH clusters were made".format(len(Bdb['MASH_cluster']\
                .unique().tolist())))

    # Step 1. Make the directories and generate the list of commands to be run
    cmds = []
    files = []
    for cluster in Bdb['MASH_cluster'].unique():
        d = Bdb[Bdb['MASH_cluster'] == cluster]
        genomes = d['location'].tolist()
        for g1 in genomes:
            for g2 in genomes:
                file_name = "{0}{1}_vs_{2}".format(ANIn_folder, \
                            get_genome_name_from_fasta(g1),\
                            get_genome_name_from_fasta(g2))
                files.append(file_name + '.delta')

                # If the file doesn't already exist, add it to what needs to be run
                if not os.path.isfile(file_name + '.delta'):
                    cmds.append(gen_nucmer_cmd(file_name,g1,g2,c=n_c,noextend=n_noextend,\
                                maxgap=n_maxgap,method=n_method))

    # Step 2. Run the nucmer commands

    if not dry:
        if len(cmds) > 0:
            thread_nucmer_cmds_status(cmds,p)

    # Step 3. Parse the nucmer output

    org_lengths = {y:dm.fasta_length(x) for x,y in zip(Bdb['location'].tolist(),Bdb['genome'].tolist())}
    Ndb = process_deltadir(files, org_lengths)
    Ndb['MASH_cluster'] = None
    for cluster in Bdb['MASH_cluster'].unique():
        d = Bdb[Bdb['MASH_cluster'] == cluster]
        Ndb['MASH_cluster'][Ndb['reference'].isin(d['genome'].tolist())] = cluster

    return Ndb

def gen_nucmer_commands(genomes,outf,c=65,maxgap=90,noextend=False,method='mum'):
    cmds = []
    for g1 in genomes:
        for g2 in genomes:
            out = "{0}{1}_vs_{2}".format(outf,get_genome_name_from_fasta(g1),get_genome_name_from_fasta(g2))
            cmds.append(gen_nucmer_cmd(out,g1,g2,c=c,noextend=noextend,maxgap=maxgap,method=method))

    return cmds

def run_nucmer_genomeList(genomes,outf,b2s,c=65,maxgap=90,noextend=False,method='mum',dry=False):
    """
    genomes is a list of locations of genomes in .fasta file. This will do pair-wise
    comparisons of those genomes using the nucmer settings given
    """

    # Run commands on biotite
    cmds = []
    for g1 in genomes:
        for g2 in genomes:
            out = "{0}{1}_vs_{2}".format(outf,get_genome_name_from_fasta(g1),get_genome_name_from_fasta(g2))
            cmds.append(gen_nucmer_cmd(out,g1,g2,c=c,noextend=noextend,maxgap=maxgap,method=method))
    if not dry:
        thread_nucmer_cmds_status(cmds, t=p)

    # Parse resulting folder
    data = process_deltadir(outf, b2s)
    data['c'] = c
    data['maxgap'] = maxgap
    data['noextend'] = noextend
    data['method'] = method

    return data

def all_vs_all_MASH(Bdb, data_folder, **kwargs):
    """
    Run MASH pairwise within all samples in Bdb
    """

    MASH_s = kwargs.get('MASH_sketch',1000)
    dry = kwargs.get('dry',False)
    overwrite = kwargs.get('overwrite', False)
    mash_exe = kwargs.get('mash_exe', 'mash')
    p = kwargs.get('processors',6)

    # set up logdir
    if 'wd' in kwargs:
        logdir = kwargs.get('wd').get_dir('cmd_logs')
    else:
        logdir = False

    # Set up folders
    MASH_folder = data_folder + 'MASH_files/'
    if not os.path.exists(MASH_folder):
        os.makedirs(MASH_folder)

    sketch_folder = MASH_folder + 'sketches/'
    if not os.path.exists(sketch_folder):
        os.makedirs(sketch_folder)

    # Make the MASH sketches
    cmds = []
    for fasta in Bdb['location'].unique():
        genome = Bdb['genome'][Bdb['location'] == fasta].tolist()[0]
        file = sketch_folder + genome
        if not os.path.isfile(file + '.msh'):
            cmd = [mash_exe, 'sketch', fasta, '-s', str(MASH_s), '-o',
                file]
            cmds.append(cmd)

    if not dry:
        if len(cmds) > 0:
            dm.thread_cmds(cmds, logdir=logdir, t=int(p))

    # Combine MASH sketches
    cmd = [mash_exe, 'paste', MASH_folder + 'ALL.msh'] + glob.glob(sketch_folder+ '*')
    # cmd = ' '.join(cmd)
    dm.run_cmd(cmd, dry, shell=False, logdir=logdir)

    # Calculate distances
    all_file = MASH_folder + 'ALL.msh'
    cmd = [mash_exe, 'dist', all_file, all_file, '>', MASH_folder
            + 'MASH_table.tsv']
    cmd = ' '.join(cmd)
    dm.run_cmd(cmd, dry, shell=True, logdir=logdir)

    # Make Mdb based on all genomes in the MASH folder
    Mdb = pd.DataFrame()
    file = MASH_folder + 'MASH_table.tsv'

    table = pd.read_csv(file,sep='\t',header = None)
    table.columns = ['genome1','genome2','dist','p','kmers']
    table['genome1'] = table['genome1'].apply(get_genome_name_from_fasta)
    table['genome2'] = table['genome2'].apply(get_genome_name_from_fasta)
    Mdb = pd.concat([Mdb,table],ignore_index=True)
    Mdb['similarity'] = 1 - Mdb['dist'].astype(float)

    # Filter out those genomes that are in the MASH folder but shouldn't be in Mdb
    genomes = Bdb['genome'].unique()
    Mdb = Mdb[Mdb['genome1'].isin(genomes)]
    Mdb = Mdb[Mdb['genome2'].isin(genomes)]

    return Mdb

def cluster_mash_database(db, data_folder= False, **kwargs):

    logging.debug('Clustering MASH database')

    P_Lmethod = kwargs.get('clusterAlg','single')
    P_Lcutoff = 1 - kwargs.get('P_ani',.9)
    dry = kwargs.get('dry',False)
    overwrite = kwargs.get('overwrite', False)

    if data_folder != False:
        data_folder = data_folder + 'Clustering_files/'
        if not os.path.exists(data_folder):
            os.makedirs(data_folder)

    db['dist'] = 1 - db['similarity']
    linkage_db = db.pivot("genome1","genome2","dist")
    Cdb, linkage = cluster_hierarchical(linkage_db, linkage_method= P_Lmethod, \
                                linkage_cutoff= P_Lcutoff)
    Cdb = Cdb.rename(columns={'cluster':'MASH_cluster'})

    if (data_folder != False):
        arguments = {'linkage_method':P_Lmethod,'linkage_cutoff':P_Lcutoff,\
                        'comparison_algorithm':'MASH'}
        logging.debug('Saving primary_linkage pickle to {0}'.format(data_folder))
        with open(data_folder + 'primary_linkage.pickle', 'wb') as handle:
            pickle.dump(linkage, handle)
            pickle.dump(linkage_db, handle)
            pickle.dump(arguments, handle)

    return Cdb

def get_genome_name_from_fasta(fasta):
    return str(os.path.basename(fasta))

# Parse NUCmer delta file to get total alignment length and total sim_errors
def parse_delta(filename):
    """Returns (alignment length, similarity errors) tuple from passed .delta.
    - filename - path to the input .delta file
    Extracts the aligned length and number of similarity errors for each
    aligned uniquely-matched region, and returns the cumulative total for
    each as a tuple.
    """
    aln_length, sim_errors = 0, 0
    for line in [l.strip().split() for l in open(filename, 'rU').readlines()]:
        if line[0] == 'NUCMER' or line[0].startswith('>'):  # Skip headers
            continue
        # We only process lines with seven columns:
        if len(line) == 7:
            aln_length += abs(int(line[1]) - int(line[0]))
            sim_errors += int(line[4])
    return aln_length, sim_errors

def gen_gANI_cmd(file, g1, g2, dir, exe):
    # Handle self comparison
    # Did a pretty exhaustive test- same comaprisons always give 100% ANI
    if g1 == g2:
        # make a copy of g1
        #g1T = g1 + '.GANI_TEMP'
        #shutil.copyfile(g1,g1T)

        #cmd = [exe,'-genome1fna',g1T,'-genome2fna',g2,'-outfile',file,'-outdir',file + 'TEMP']
        return []
    else:
        cmd = [exe,'-genome1fna',g1,'-genome2fna',g2,'-outfile',file,'-outdir',file + 'TEMP']
    #cmd = [exe,'-genome1fna',g1,'-genome2fna',g2,'-outfile',file,'-outdir',dir]
    return cmd

def process_deltadir(deltafiles, org_lengths, logger=None):
    """Returns a tuple of ANIm results for .deltas in passed directory.
    - delta_dir - path to the directory containing .delta files
    - org_lengths - dictionary of total sequence lengths, keyed by sequence
    Returns the following pandas dataframes in a tuple; query sequences are
    rows, subject sequences are columns:
    - alignment_lengths - symmetrical: total length of alignment
    - percentage_identity - symmetrical: percentage identity of alignment
    - alignment_coverage - non-symmetrical: coverage of query and subject
    - similarity_errors - symmetrical: count of similarity errors
    May throw a ZeroDivisionError if one or more NUCmer runs failed, or a
    very distant sequence was included in the analysis.
    """
    # Process directory to identify input files
    #deltafiles = glob.glob(delta_dir + '*.delta')

    Table = {'querry':[],'reference':[],'alignment_length':[],'similarity_errors':[],
            'ref_coverage':[],'querry_coverage':[],'ani':[], 'reference_length':[],
            'querry_length':[],'alignment_coverage':[]}

    # Process .delta files assuming that the filename format holds:
    # org1_vs_org2.delta
    zero_error = False  # flag to register a divide-by-zero error
    for deltafile in deltafiles:
        qname, sname = os.path.splitext(os.path.split(deltafile)[-1])[0].split('_vs_')
        tot_length, tot_sim_error = parse_delta(deltafile)
        if tot_length == 0 and logger is not None:
            logging.info("Total alignment length reported in " +
                               "%s is zero!" % deltafile)
        query_cover = float(tot_length) / org_lengths[qname]
        sbjct_cover = float(tot_length) / org_lengths[sname]
        # Calculate percentage ID of aligned length. This may fail if
        # total length is zero.
        # The ZeroDivisionError that would arise should be handled
        # Common causes are that a NUCmer run failed, or that a very
        # distant sequence was included in the analysis.
        try:
            perc_id = 1 - float(tot_sim_error) / tot_length
        except ZeroDivisionError:
            #print("Alignment between {0} and {1} has 0 alignment!".format(qname,sname))
            perc_id = 0  # set arbitrary value of zero identity
            zero_error = True

        Table['querry'].append(qname)
        Table['querry_length'].append(org_lengths[qname])
        Table['reference'].append(sname)
        Table['reference_length'].append(org_lengths[sname])
        Table['alignment_length'].append(tot_length)
        Table['similarity_errors'].append(tot_sim_error)
        Table['ani'].append(perc_id)
        Table['ref_coverage'].append(sbjct_cover)
        Table['querry_coverage'].append(query_cover)
        Table['alignment_coverage'].append((tot_length * 2)/(org_lengths[qname]\
                                                             + org_lengths[sname]))

    df = pd.DataFrame(Table)
    return df

def process_deltafiles(deltafiles, org_lengths, logger=None, **kwargs):

    Table = {'querry':[],'reference':[],'alignment_length':[],'similarity_errors':[],
            'ref_coverage':[],'querry_coverage':[],'ani':[], 'reference_length':[],
            'querry_length':[],'alignment_coverage':[]}

    # Process .delta files assuming that the filename format holds:
    # org1_vs_org2.delta
    coverage_method = kwargs.get('coverage_method')
    #logging.debug('coverage_method is {0}'.format(coverage_method))

    for deltafile in deltafiles:
        qname, sname = os.path.splitext(os.path.split(deltafile)[-1])[0].split('_vs_')
        # Handle naming of "filtered"
        if sname.endswith('.delta'):
            sname = sname[:-6]

        tot_length, tot_sim_error = parse_delta(deltafile)
        if tot_length == 0 and logger is not None:
            logging.info("Total alignment length reported in " +
                               "%s is zero!" % deltafile)
        query_cover = float(tot_length) / org_lengths[qname]
        sbjct_cover = float(tot_length) / org_lengths[sname]

        try:
            perc_id = 1 - float(tot_sim_error) / tot_length
        except ZeroDivisionError:
            #print("Alignment between {0} and {1} has 0 alignment!".format(qname,sname))
            perc_id = 0  # set arbitrary value of zero identity
            zero_error = True

        Table['querry'].append(qname)
        Table['querry_length'].append(org_lengths[qname])
        Table['reference'].append(sname)
        Table['reference_length'].append(org_lengths[sname])
        Table['alignment_length'].append(tot_length)
        Table['similarity_errors'].append(tot_sim_error)
        Table['ani'].append(perc_id)
        Table['ref_coverage'].append(sbjct_cover)
        Table['querry_coverage'].append(query_cover)

        if coverage_method == 'total':
            Table['alignment_coverage'].append((tot_length * 2)/(org_lengths[qname]\
                                                             + org_lengths[sname]))
        elif coverage_method == 'larger':
            Table['alignment_coverage'].append(max((tot_length/org_lengths[qname]),\
                (tot_length/org_lengths[sname])))

    df = pd.DataFrame(Table)
    return df

### MAKE IT SO THAT YOU REMOVE THE _TEMP MARKER FROM THE GENOMES, AND DELTE THE
### FILE WHILE YOU'RE AT IT
def process_gani_files(files):
    Table = {'querry':[],'reference':[],'ani':[],'alignment_coverage':[]}
    for file in files:
        results = parse_gani_file(file)

        # calculate coverage using the "larger" method
        cov = max(results['rq_coverage'], results['qr_coverage'])

        # Add forward results
        Table['reference'].append(results['reference'])
        Table['querry'].append(results['querry'])
        Table['ani'].append(float(results['rq_ani'])/100)
        Table['alignment_coverage'].append(cov)
        #Table['alignment_coverage'].append(results['rq_coverage'])

        # Add reverse results [they're the same]
        Table['reference'].append(results['querry'])
        Table['querry'].append(results['reference'])
        Table['ani'].append(float(results['qr_ani'])/100)
        Table['alignment_coverage'].append(cov)
        #Table['alignment_coverage'].append(results['qr_coverage'])

    # Add self comparisons
    for g in set(Table['reference']):
        Table['reference'].append(g)
        Table['querry'].append(g)
        Table['ani'].append(1)
        Table['alignment_coverage'].append(1)

    Gdb = pd.DataFrame(Table)
    return Gdb

'''
This method takes in bdb (a table with the columns location and genome), runs
pair-wise comparisons between all genomes in the sample, and returns a table
with at least the columns 'reference', 'querry', 'ani','coverage', depending
on what algorithm is called
'''
def compare_genomes(bdb, algorithm, data_folder, **kwargs):
    # To handle other versions of this method which passed in a WorkDirectory
    # instead of data_folder string
    if isinstance(data_folder,drep.WorkDirectory.WorkDirectory):
        data_folder = data_folder.location + '/data/'

    if algorithm == 'ANIn':
        genome_list = bdb['location'].tolist()
        working_data_folder = data_folder + 'ANIn_files/'
        df = run_pairwise_ANIn(genome_list, working_data_folder, **kwargs)
        return df

    elif algorithm == 'ANImf':
        genome_list = bdb['location'].tolist()
        working_data_folder = data_folder + 'ANImf_files/'
        df = run_pairwise_ANImf(genome_list, working_data_folder, **kwargs)
        return df

    elif algorithm == 'gANI':
        working_data_folder = data_folder + 'gANI_files/'
        prod_folder = data_folder + 'prodigal/'
        df = run_pairwise_gANI(bdb, working_data_folder, \
                prod_folder = prod_folder, **kwargs)
        return df

    elif algorithm == 'mauve':
        working_data_folder = data_folder + 'mauve_files/'
        df = run_pairwise_mauve(bdb, working_data_folder, **kwargs)
        return df

    else:
        logging.error("{0} not supported".format(algorithm))
        sys.exit()

def run_pairwise_ANIn(genome_list, ANIn_folder, **kwargs):
    p = kwargs.get('processors',6)
    genomes = genome_list

    # Make folder
    if not os.path.exists(ANIn_folder):
        os.makedirs(ANIn_folder)

    # Gen commands
    cmds = []
    files = []
    for g1 in genomes:
        for g2 in genomes:
            file_name = "{0}{1}_vs_{2}".format(ANIn_folder, \
                        get_genome_name_from_fasta(g1),\
                        get_genome_name_from_fasta(g2))
            files.append(file_name)

            # If the file doesn't already exist, add it to what needs to be run
            if not os.path.isfile(file_name + '.delta'):
                cmds.append(gen_nucmer_cmd(file_name,g1,g2))

    # Run commands
    if len(cmds) > 0:
        for c in cmds:
            logging.debug(' '.join(c))

        if 'wd' in kwargs:
            logdir = kwargs.get('wd').get_dir('cmd_logs')
        else:
            logdir = False
        dm.thread_cmds(cmds, logdir=logdir, t=int(p))
        #thread_nucmer_cmds_status(cmds,p,verbose=False)

    # Parse output
    org_lengths = {}
    for genome in genomes:
        #logging.debug("getting genome length of {0}".format(genome))
        org_lengths[get_genome_name_from_fasta(genome)] = dm.fasta_length(genome)

    deltafiles = ["{0}.delta".format(file) for file in files]
    df = process_deltafiles(deltafiles, org_lengths, **kwargs)

    return df

def run_pairwise_ANImf(genome_list, ANIn_folder, **kwargs):
    p = kwargs.get('processors',6)
    genomes = genome_list

    # Make folder
    if not os.path.exists(ANIn_folder):
        os.makedirs(ANIn_folder)

    # Gen commands
    cmds = []
    files = []
    for g1 in genomes:
        for g2 in genomes:
            file_name = "{0}{1}_vs_{2}".format(ANIn_folder, \
                        get_genome_name_from_fasta(g1),\
                        get_genome_name_from_fasta(g2))
            files.append(file_name)

            # If the file doesn't already exist, add it to what needs to be run
            if not os.path.isfile(file_name + '.delta.filtered'):
                cmds.append(gen_animf_cmd(file_name,g1,g2))

    # Run commands
    if len(cmds) > 0:
        for c in cmds:
            logging.debug(' '.join(c))

        if 'wd' in kwargs:
            logdir = kwargs.get('wd').get_dir('cmd_logs')
        else:
            logdir = False
        dm.thread_cmds(cmds, shell=True, logdir=logdir, t=int(p))

    # Parse output
    org_lengths = {}
    for genome in genomes:
        #logging.debug("getting genome length of {0}".format(genome))
        org_lengths[get_genome_name_from_fasta(genome)] = dm.fasta_length(genome)

    deltafiles = ["{0}.delta.filtered".format(file) for file in files]
    df = process_deltafiles(deltafiles, org_lengths, **kwargs)

    return df

def run_pairwise_mauve(bdb, data_folder, **kwargs):
    p = kwargs.get('processors',6)

    # Make folder
    if not os.path.exists(data_folder):
        os.makedirs(data_folder)

    # Gen commands
    cmds = []
    files = []
    for g1 in genomes:
        for g2 in genomes:
            '''
            file_name = "{0}{1}_vs_{2}".format(ANIn_folder, \
                        get_genome_name_from_fasta(g1),\
                        get_genome_name_from_fasta(g2))
            files.append(file_name)

            # If the file doesn't already exist, add it to what needs to be run
            if not os.path.isfile(file_name + '.delta'):
                cmds.append(gen_nucmer_cmd(file_name,g1,g2))
            '''

    # Run commands
    if len(cmds) > 0:
        thread_nucmer_cmds_status(cmds,p,verbose=False)

    # Parse output
    '''
    org_lengths = {get_genome_name_from_fasta(y):dm.fasta_length(x) \
                    for x,y in zip(genomes,genomes)}
    deltafiles = ["{0}.delta".format(file) for file in files]
    df = process_deltafiles(deltafiles, org_lengths)
    '''

    return df

def run_pairwise_gANI(bdb, gANI_folder, verbose = False, **kwargs):
    # gANI_exe = kwargs.get('gANI_exe','ANIcalculator')
    # loc = shutil.which('ANIcalculator')
    loc, works = dBonus.find_program('ANIcalculator')
    if loc == None:
        logging.error('Cannot locate the program {0}- make sure its in the system path'\
            .format('ANIcalculator (for gANI)'))
    if works == False:
        logging.error('Program {0} is not working'\
            .format('ANIcalculator (for gANI)'))
    gANI_exe = loc

    p = kwargs.get('processors',6)
    prod_folder = kwargs.get('prod_folder')
    genomes = bdb['location'].tolist()

    # Make folder
    if not os.path.exists(gANI_folder):
        os.makedirs(gANI_folder)

    # Remove crap folders- they shouldn't exist and if they do it messes things up
    crap_folders = glob.glob(gANI_folder + '*.gANITEMP')
    for crap_folder in crap_folders:
        if os.path.exists(crap_folder):
            logging.debug("CRAP FOLDER EXISTS FOR gANI- removing {0}".format(crap_folder))
            shutil.rmtree(crap_folder)

    if not os.path.exists(prod_folder):
        os.makedirs(prod_folder)

    # Run prodigal
    if verbose:
        logging.info("Running prodigal...")
    dFilter.run_prodigal(bdb, prod_folder, verbose=verbose, **kwargs)

    # Gen gANI commands
    if verbose:
        logging.info("Running gANI...")
    cmds = []
    files = []
    for i, g1 in enumerate(genomes):
        for j, g2 in enumerate(genomes):
            if i > j:
                name1= get_genome_name_from_fasta(g1)
                name2= get_genome_name_from_fasta(g2)
                file_name = "{0}{1}_vs_{2}.gANI".format(gANI_folder, \
                            name1, name2)
                files.append(file_name)

                # If the file doesn't already exist, add it to what needs to be run
                if not os.path.isfile(file_name):
                    fna1 = "{0}{1}.fna".format(prod_folder,name1)
                    fna2 = "{0}{1}.fna".format(prod_folder,name2)
                    cmds.append(gen_gANI_cmd(file_name,fna1,fna2,gANI_folder,gANI_exe))

    # Run commands
    if len(cmds) > 0:
        logging.debug('Running gANI commands: {0}'.format('\n'.join([' '.join(x) for x in cmds])))
        if 'wd' in kwargs:
            logdir = kwargs.get('wd').get_dir('cmd_logs')
        else:
            logdir = False
        dm.thread_cmds(cmds, logdir=logdir, t=int(p))

    else:
        if verbose:
            logging.info("gANI already run- will not re-run")

    # Parse output
    df = process_gani_files(files)

    # Handle self-comparisons
    #df['reference'] = [i.replace('.fna.GANI_','') for i in df['reference']]
    #df['querry'] = [i.replace('.fna.GANI_','') for i in df['querry']]
    #df = df.drop_duplicates()

    # Add self-comparisons if there is only one genome
    if len(genomes) == 1:
        Table = {'querry':[],'reference':[],'ani':[],'alignment_coverage':[]}
        for g in genomes:
            Table['reference'].append(get_genome_name_from_fasta(g))
            Table['querry'].append(get_genome_name_from_fasta(g))
            Table['ani'].append(1)
            Table['alignment_coverage'].append(1)
        d = pd.DataFrame(Table)
        df = pd.concat([df,d],ignore_index=True)

    return df

def parse_gani_file(file):
    try:
        x = pd.read_table(file)
    except:
        logging.error('gANI file {0} does not exist. The most likely reason is that '.format(file)\
            + 'one of the genomes has a .fasta header than gANI doesnt like.'\
            + ' Known issues include having a header over 160 characaters long, or any '\
            + 'kind of special character besides "_" (including ".", ":", and "-").'\
            + ' To fix this error either use ANIn (which is not as picky with fasta headers)'\
            + ', or fix the .fasta headers to conform to those rules.')
        sys.exit()
    x = x.rename(columns={'GENOME1':'reference','GENOME2':'querry','AF(1->2)':'rq_coverage',\
                        'AF(2->1)':'qr_coverage','ANI(1->2)':'rq_ani','ANI(2->1)':'qr_ani'})
    dict = x.to_dict(orient='list')
    dict = {k:v[0] for k,v in dict.items()}
    dict['reference'] = dict['reference'][:-4]
    dict['querry'] = dict['querry'][:-4]
    return dict

def gen_nucmer_cmd(prefix,ref,querry,c='65',noextend=False,maxgap='90',method='mum'):
    cmd = ['nucmer','--' + method,'-p',prefix,'-c',str(c),'-g',str(maxgap)]
    if noextend: cmd.append('--noextend')
    cmd += [ref,querry]
    return cmd

def gen_animf_cmd(prefix,ref,querry, **kwargs):
    '''
    return animf command. It will be a single string, with the format:
    "nucmer cmd ; filter cmd"
    '''

    nucmer_cmd = gen_nucmer_cmd(prefix,ref,querry, **kwargs)

    delta = prefix + '.delta'
    out = delta + '.filtered'

    filter_cmd = gen_filter_cmd(prefix + '.delta', out)
    return ' '.join(nucmer_cmd) + '; ' + ' '.join(filter_cmd)

def gen_filter_cmd(delta, out):
    cmd = ["delta-filter", '-r', '-q', delta, '>', out]
    return cmd

def thread_nucmer_cmds(cmds,t=10):
    pool = multiprocessing.Pool(processes=t)
    pool.map(run_nucmer_cmd,cmds)
    pool.close()
    pool.join()
    return

def thread_nucmer_cmds_status(cmds,t=10,verbose=True):
    total = len(cmds)
    if verbose:
        minutes = ((float(total) * (float(0.33))) /float(t))
        logging.info("Running {0} mummer comparisons: should take ~ {1:.1f} min".format(total,minutes))
    pool = multiprocessing.Pool(processes=int(t))
    rs = pool.map_async(run_nucmer_cmd,cmds)
    pool.close()

    pool.join()
    return

def thread_mash_cmds_status(cmds,t=10):
    total = len(cmds)
    pool = multiprocessing.Pool(processes=t)
    rs = pool.map_async(run_nucmer_cmd,cmds)
    pool.close()

    while(True):
        done = total - (rs._number_left * rs._chunksize)
        percR = (done/total) * 100
        sys.stdout.write('\r')
        sys.stdout.write("[{0:20}] {1:3.2f}%".format('='*int(percR/5), percR))
        sys.stdout.flush()
        if (rs.ready()):
            sys.stdout.write('\n')
            break
        time.sleep(0.5)

    pool.join()
    return

def gen_nomash_cdb(Bdb):
    Cdb = Bdb.copy()
    Cdb['MASH_cluster'] = 0
    return Cdb

# def run_nucmer_cmd(cmd,dry=False,shell=False):
#     devnull = open(os.devnull, 'w')
#     if shell:
#         if not dry: call(cmd,shell=True, stderr=devnull,stdout=devnull)
#         else: print(cmd)
#     else:
#         if not dry: call(cmd, stderr=devnull,stdout=devnull)
#         else: print(' '.join(cmd))
#     return
#
# def run_mash_cmd(cmds):
#     cmd = ' '.join(cmd)
#     dm.run_cmd(cmd,dry=False,shell=True)

def load_genomes(genome_list):
    Table = {'genome':[],'location':[]}

    for genome in genome_list:
        assert os.path.isfile(genome), "{0} is not a file".format(genome)
        Table['genome'].append(os.path.basename(genome))
        Table['location'].append(os.path.abspath(genome))

    Bdb = pd.DataFrame(Table)
    return Bdb

def average_ani(row,db):
    g1 = row['querry']
    g2 = row['reference']
    if g1 == g2:
        return 1
    else:
        ani1 = float(row['ani'])
        ani2 = float(db['ani'][(db['reference'] == g1) & (db['querry'] == g2)].tolist()[0])
        avg = np.mean([ani1,ani2])
        return avg

def add_avani(db):
    '''
    add a column titled 'av_ani' to the passed in dataframe

    dataframe must have rows reference, querey, and ani.
    '''

    logging.debug('making dictionary for average_ani')
    combo2value = {}
    for i, row in db.iterrows():
        combo2value["{0}-vs-{1}".format(row['querry'], row['reference'])] \
            = row['ani']

    logging.debug('list comprehension for average_ani')
    db['av_ani'] = [np.mean([combo2value["{0}-vs-{1}".format(q,r)], \
                        combo2value["{0}-vs-{1}".format(r,q)]]) if r != q else 1\
                        for q, r in zip(db['querry'].tolist(), \
                        db['reference'].tolist())]

    logging.debug('averageing done')

def nucmer_preset(preset):
   #nucmer argument c, n_maxgap, n_noextend, n_method

    assert preset in ['tight','normal']

    if preset == 'tight':
        return 65, 1, True, 'mum'

    elif preset == 'normal':
        return 65, 90, False, 'mum'

def gen_nomani_cdb(Cdb, Mdb, **kwargs):
    c_method = kwargs.get('clusterAlg','single')
    threshold = 1 - kwargs.get('P_ani',.9)
    data_folder = kwargs.get('data_folder') + 'Clustering_files/'

    # Make Cdb look like Cdb
    cdb = Cdb.copy()
    cdb['secondary_cluster'] = ["{0}_0".format(i) for i in cdb['MASH_cluster']]
    cdb.rename(columns={'MASH_cluster': 'primary_cluster'}, inplace=True)
    cdb['threshold'] = threshold
    cdb['cluster_method'] = c_method
    cdb['comparison_algorithm'] = 'MASH'

    # Delete any only secondary clusters
    if kwargs.get('overwrite',False):
        for fn in glob.glob(data_folder + 'secondary_linkage_cluster*'):
            os.remove(fn)

    return cdb

def test_clustering():
    # Get test genomes
    test_genomes = glob.glob(str(os.getcwd()) + '/../test/genomes/*')
    names = [os.path.basename(g) for g in test_genomes]
    assert names == ['Enterococcus_faecalis_T2.fna',
	                 'Escherichia_coli_Sakai.fna',
	                 'Enterococcus_casseliflavus_EC20.fasta',
	                 'Enterococcus_faecalis_TX0104.fa',
	                 'Enterococcus_faecalis_YI6-1.fna']
    Bdb = load_genomes(test_genomes)

    # Set test directory
    test_directory = str(os.getcwd()) + '/../test/test_backend/'
    dm.make_dir(test_directory,overwrite=True)

    # Perform functional test
    Cdb, Mdb, Ndb = cluster_genomes(Bdb,test_directory)

    # Confirm it's right by showing there is one group of 3 and two groups of one
    group_lengths = sorted([len(Cdb['genome'][Cdb['ANIn_cluster'] == x].tolist()) for x in Cdb['ANIn_cluster'].unique()])
    assert group_lengths == [1, 1, 3]

    print("Functional test success!")

if __name__ == '__main__':

    test_clustering()


from django.conf.urls.defaults import *
from os import path as os_path
from django.conf import settings
import economico.views
# Uncomment the next two lines to enable the admin:
from django.contrib import admin
admin.autodiscover()

urlpatterns = patterns('',
    (r'^admin/(.*)', admin.site.root),
    (r'^', include('economico.urls')),
    (r'^seguridad-alimentaria/', include('seguridad_alimentaria.urls')),
    (r'^demografico/', include('demografico.urls')),
    (r'^ajax/municipio/(?P<departamento>\d+)/$', 'lugar.views.get_municipios'),
    (r'^ajax/comunidad/(?P<municipio>\d+)/$', 'lugar.views.get_comunidad'),
    (r'^ajax/entrevista/(?P<comunidad>\d+)/$', 'sequias.views.get_entrevista'),
    (r'sequia/', include('sequias.urls')),
)

if settings.DEBUG:
    urlpatterns += patterns('',
            (r'^css/(.*)$', 'django.views.static.serve',
                {'document_root': os_path.join(settings.MEDIA_ROOT + '/css')}),
            (r'^js/(.*)$', 'django.views.static.serve',
                {'document_root': os_path.join(settings.MEDIA_ROOT + '/js')}),
            (r'^imagen/(.*)$', 'django.views.static.serve',
                {'document_root': os_path.join(settings.MEDIA_ROOT + '/imagen')}),
)


import asyncio
import random

from prosumer import Prosumer

class PrintNumbersProducer(Prosumer):
    """print numbers asynchronously"""

    def __init__(self, data: int, context):
        super().__init__(data, context)

    async def produce(self):
        self.logger.debug('%s producer starting...', self.name)
        for i in range(self.data):
            self.add_task(i)
        await asyncio.gather(*self.tasks)
        self.logger.debug('[%s] finished', self.name)
        self.context.prosumers.remove(self)

    async def fill_queue(self):
        for i in range(1, self.data):
            await self.queue.put(i)

    async def work(self, number):
        sleep_time = random.randint(1, 3)
        await asyncio.sleep(sleep_time)
        self.logger.info(number)
        self.append()

    @property
    def name(self):
        return 'print_number'


import sys
from collections import deque
# 09m 37s
input=sys.stdin.readline
for _ in range(int(input())):
    N,M=map(int,input().split())
    D=list(map(int,input().split()))
    q=deque()
    # 덱 사용
    for i in range(N): q.append((i,D[i]))
    # 덱에 (인덱스, 중요도) 추가

    res=[]
    while q:
        idx,doc=q[0]
        # 제일 앞에 있는 문서를 기준으로 뒤에 중요도가 높은 문서가 있다면 인쇄하지 않고 q에 다시 넣음
        # 아니면, res(정답 배열)에 넣음
        cnt=0
        for j in range(1,len(q)):
            # 제일 앞에 있는 문서를 기준으로 하기 때문에 j는 1부터 시작
            if q[j][1]>doc: cnt+=1
            # 제일 앞 문서의 중요도보다 큰 문서가 있으면 개수 증가
        if cnt>0: q.append(q.popleft())
        # 중요도가 큰 문서가 하나 이상이면 인쇄하지 않고 q에 제일 뒤에 배치
        else: res.append(q.popleft())
        # 아니면 res(정답)에 추가

    for i in range(N):
        if res[i][0]==M: print(i+1); break


from archive.CognitiveEngine.modules import CognitiveEngine
from flask import Blueprint
from flask import Response, request

from archive.CognitiveEngine.functions.module_functions import *
from archive.CognitiveEngine.functions.route_functions import *

from CPSBuilder.modules.archive.execution_module import ExecutionModule

import logging
logger = logging.getLogger(__name__)

# config mongodb
client = MongoClient(config.mongo_ip, config.mongo_port)

# initialize modules
cognitive_engine = CognitiveEngine(client)
execution_module = ExecutionModule(client)

# initialize blueprint
process_flow_control = Blueprint('process_flow_control', __name__)

"""
Route is used to test cognitive engine algorithm.
Not for production.
"""
@process_flow_control.route("/api/job-monitoring/test-js", methods=['POST'])
def test_js_monitoring():
    job_id = request.args.get('job_id')
    user_id = request.args.get('user_id')
    action_seq = int(request.args.get('action_seq'))
    step_seq = int(request.args.get('step_seq'))
    execution_module.post_job_to_board(job_id, user_id, action_seq, step_seq)
    return ""

@process_flow_control.route('/api/cognitive-engine/reasoner', methods=['GET', 'POST'])
# def get_process_status(job_id, user_id, action_status, resource_type, action_sequence):
def get_process_status():
    """
    This function is not being used in the CPS. Developed for testing of the reasoner.
    """
    if request.method == 'GET':
        return Response('pong', 200)
    data = request.get_json()
    job_id = data.get('job_id')
    user_id = data.get('user_id')
    action_status = data.get('action_status')
    action_seq = data.get('action_seq')
    step_seq = data.get('step_seq')
    total_action_count = data.get('total_action_count')
    total_step_count = data.get('total_step_count')

    step_status, next_action_seq, next_step_seq = cognitive_engine.get_inference_class(
        job_id, user_id, action_status, action_seq, step_seq, total_action_count, total_step_count)
    res = {
        "step_status": step_status,
        "next_action_seq": next_action_seq,
        "next_step_seq": next_step_seq
    }
    headers = {'Content-Type': 'application/json'}
    return Response(json.dumps(res), 200, headers=headers)

@process_flow_control.route('/api/cognitive-engine/reasoner/get-next-step', methods=['GET', 'POST'])
# def get_process_status(job_id, user_id, action_status, resource_type, action_sequence):
def get_next_step():
    """
    This function is not being used in the CPS. Developed for testing of the reasoner.
    """
    if request.method == 'GET':
        return Response('pong', 200)
    data = request.get_json()
    job_id = data.get('job_id')
    user_id = data.get('user_id')
    action_status = data.get('action_status')
    action_seq = int(data.get('action_seq'))
    step_seq = int(data.get('step_seq'))
    total_action_count = data.get('total_action_count')
    total_step_count = data.get('total_step_count')
    print(f"data received!")

    step_status, next_action_seq, next_step_seq = cognitive_engine.get_next_step(
        job_id, user_id, action_status, action_seq, step_seq, total_action_count, total_step_count)
    res = {
        "step_status": step_status,
        "next_action_seq": next_action_seq,
        "next_step_seq": next_step_seq
    }
    headers = {'Content-Type': 'application/json'}
    return Response(json.dumps(res), 200, headers=headers)


from multiprocessing import Pool
from statistics import mode
from symbols import ExemplarySymbols
import scipy.signal as sig
import numpy as np

# class for decoding DTMF signal

# signal is divided into frames (n samples)
# for each frame its PSD is calculated and compared with pure symbols (symbols.py)
# then the resulting sequence is parsed

class DTMFreader:
    def __init__(self, consts, signal, f_samp):
        # passed parameters
        # consts - dictionary containing config values (const.py)
        # signal - audio signal to be anlyzed
        # f_samp - sampling fequency

        self.signal = signal

        self.frame_size = consts["frame_size"]
        self.frame_step = consts["frame_step"]
        self.min_symbol_frames = consts["min_symbol_frames"]
        self.max_symbol_break = consts["max_symbol_break"]
        self.min_error_ratio = consts["min_error_ratio"]
        self.read_entire_signal = consts["read_entire_signal"]
        self.max_frames = consts["max_frames"]

        self.symbols = ExemplarySymbols(self.frame_size, f_samp)

    # run method performs all the tasks and returns reqovered sequence

    def run(self):
        res = self.read_signal()

        reading = list()

        for i in res:
            reading.append(i["symbol"])

        self.sequence = self.parse_symbols(reading)

        return ''.join(self.sequence)

    # parse_symbols method coverts result of reading frames into a sequence

    def parse_symbols(self, read_results):

        symbols_parsed = list()

        i = 0
        n_of_symbols = len(read_results)

        while i < n_of_symbols - self.min_symbol_frames:

            # analyzing fragment of reading result

            analyzed_fragment = read_results[i : i + self.min_symbol_frames]
            most_frequent_symbol = mode(analyzed_fragment)

            # whent the most frequent symbol is None the analyzed fragment is too noisy to recover

            if most_frequent_symbol == None:
                i += 1
                continue

            symbol_count = analyzed_fragment.count(most_frequent_symbol)

            # if there are enough occurences of the most frequent symbol in the fragment
            # the symbol is considered to be present in the original signal

            if symbol_count >= self.min_symbol_frames - self.max_symbol_break:
                # symbol was read
                symbols_parsed.append(most_frequent_symbol)
                i += self.min_symbol_frames

                # if the symbol is present in the next frames it is considered to be the same symbol,
                # and the following reads are skipped

                try:
                    while read_results[i : i + self.max_symbol_break].count(most_frequent_symbol) >= self.max_symbol_break // 2:
                        i += self.max_symbol_break // 2
                except IndexError:
                    pass

            i += 1

        return symbols_parsed

    # read_signal method reads the frames in parallel

    def read_signal(self):

        # determining the number of frames to be read

        signal_length_in_frames = int((len(self.signal) - self.frame_size) / self.frame_step)

        if self.read_entire_signal:
            n_of_frames = signal_length_in_frames
        else:
            n_of_frames = self.max_frames if self.max_frames <= signal_length_in_frames else signal_length_in_frames

        # defining function for running in parallel
        global process_frame
        def process_frame(i):
            frame_index = i * self.frame_step
            return self.read_frame(frame_index)

        pool = Pool(processes = 16)

        frames_result = pool.map(process_frame, range(n_of_frames))
        return frames_result

    # read_frame method reads individual frame and compares it to clear DTMF symbols

    def read_frame(self, frame_index: int):

        if not isinstance(frame_index, int):
            raise TypeError("frame_index must be int")

        # dafining the frame and computing its PSD
        frame = self.signal[frame_index : frame_index + self.frame_size]
        (freqspace, frame_PSD) = sig.periodogram(frame, self.symbols.f_samp, scaling='density')
        frame_max = max(frame_PSD)

        errors = list()

        # comparing the frame to all DTMF symbols

        for i in range(12):
            horz_index = i % 3
            vert_index = i // 3

            # to compute MAE correctly the frame is scaled so its max values is the same as the DTMF symbol
            scalar = self.symbols.symbol_maxima[vert_index][horz_index] / frame_max

            # mean absolute error
            mae = np.mean(np.abs(self.symbols.PSDs[vert_index][horz_index] - frame_PSD * scalar))
            errors.append(mae)

        # characteristics of the mae
        mean_error = np.mean(errors)
        min_error = min(errors)
        mean_to_min_err = mean_error / min_error

        # if mean_to_min_err is larger than the value specified in config,
        # the symbol is considered to be read reliably
        if mean_to_min_err >= self.min_error_ratio:
            index_of_min = errors.index(min_error)
            symbol_read = self.symbols.symbols[index_of_min // 3][index_of_min % 3]
        else:
            symbol_read = None

        return {
            "symbol" : symbol_read,
            "error_ratio" : mean_to_min_err
        }


import numpy as np
import scipy.io
from polyrat import *
from pgf import PGF

# Load data
dat = scipy.io.loadmat('penzl3.dat')
X = dat['X']
y = dat['y'].flatten()

n1, n2, n3 = [],[],[]
d1, d2, d3 = [],[],[]
err = []
for d in [5, 6, 7, 8, 9, 10]:
	num_degree = [d-1, 1,3]
	denom_degree = [d, 2, 4]

	ssk = SKRationalApproximation(num_degree, denom_degree, refine = False, maxiter = 10)
	ssk.fit(X, y)

	n1.append(num_degree[0])
	n2.append(num_degree[1])
	n3.append(num_degree[2])

	d1.append(denom_degree[0])
	d2.append(denom_degree[1])
	d3.append(denom_degree[2])
	err.append(np.linalg.norm(ssk(X) - y))

	pgf = PGF()
	pgf.add('n1', n1)
	pgf.add('n2', n2)
	pgf.add('n3', n3)

	pgf.add('d1', d1)
	pgf.add('d2', d2)
	pgf.add('d3', d3)
	pgf.add('err', err)
	pgf.write('data/fig_penzl3_ssk.dat')


import torch
from enum import IntEnum

class ElemType(IntEnum):
    Line2 = 0
    Quad4 = 1

class Mesh:
    def __init__(self, device):
        self.device = device

    def elems(self):
        return self.nodes[self.connectivity]

    def num_nodes(self):
        return len(self.nodes)

    def num_elems(self):
        return len(self.connectivity)

class LineMesh(Mesh):
    def __init__(self, params, device):
        super().__init__(device)
        self.nodes = torch.column_stack(
            (
                torch.linspace(
                    params["xmin"], params["xmax"], params["N"] + 1, device=self.device
                ),
                torch.zeros((params["N"] + 1, 1), device=self.device),
                torch.zeros((params["N"] + 1, 1), device=self.device),
            )
        )

        self.connectivity = torch.column_stack(
            (
                torch.arange(0, params["N"], device=self.device),
                torch.arange(1, params["N"] + 1, device=self.device),
            )
        )

        self.elem_type = torch.full(
            (params["N"], 1), ElemType.Line2, device=self.device
        )

class RectangleMesh(Mesh):
    def __init__(self, params, device):
        super().__init__(device)
        grid_x, grid_y = torch.meshgrid(
            torch.linspace(
                params["xmin"], params["xmax"], params["Nx"] + 1, device=self.device
            ),
            torch.linspace(
                params["ymin"], params["ymax"], params["Ny"] + 1, device=self.device
            ),
            indexing="ij",
        )
        self.nodes = torch.column_stack(
            (grid_x.flatten(), grid_y.flatten(), torch.zeros_like(grid_x.flatten()))
        )

        self.connectivity = torch.empty(
            (params["Nx"] * params["Ny"], 4), dtype=torch.int64, device=self.device
        )
        for i in range(params["Nx"]):
            for j in range(params["Ny"]):
                elem_id = i * params["Ny"] + j
                self.connectivity[elem_id, 2] = elem_id + i
                self.connectivity[elem_id, 1] = elem_id + i + 1
                self.connectivity[elem_id, 3] = elem_id + i + 1 + params["Ny"]
                self.connectivity[elem_id, 0] = elem_id + i + 2 + params["Ny"]

        self.elem_type = torch.full(
            (params["Nx"] * params["Ny"], 1), ElemType.Quad4, device=self.device
        )