import matplotlib as mpl
import collections
from matplotlib.ticker import ScalarFormatter, NullFormatter, PercentFormatter
from matplotlib import gridspec
from matplotlib.ticker import FormatStrFormatter
import matplotlib.cm as cm

params = { 'figure.figsize': (8, 10),
           'legend.fontsize': 8,
           'lines.color':'black',
           'lines.linewidth':1,
            'xtick.labelsize':10,
            'xtick.major.pad'  : 3,
            'xtick.major.size' : 2,
            'ytick.major.pad'  : 3,
            'ytick.major.size' : 2,
            'ytick.labelsize':10,
            'axes.labelsize':8,
            'font.size':8,
            'axes.labelpad':2,
            'pdf.fonttype' : 42,
            'figure.dpi': 600,
            }
mpl.rcParams.update(params)

def ylabl( text, axtrans ):
   pylab.text(-0.15, 0.5,text, fontsize = 10, horizontalalignment='center',
   verticalalignment='center',rotation='vertical',
   transform = axtrans)
def ylablsup( text, axtrans ):
   pylab.text(-0.10, 0.5,text, fontsize = 8, horizontalalignment='center',
   verticalalignment='center',rotation='vertical',
   transform = axtrans)
def ylabr( text, axtrans ):
   pylab.text(1.15, 0.5,text,fontsize = 10, horizontalalignment='center',
   verticalalignment='center',rotation='vertical',
   transform = axtrans)

import sys, os
import numpy as np
import pylab
import matplotlib.pyplot as plt
import pickle, pickle
import scipy.stats
from scipy.optimize import leastsq, curve_fit
import scipy
import statsmodels.api as sm # recommended import according to the docs
from sklearn.metrics import auc
import pandas as pd
from statsmodels.formula.api import ols
from copy import deepcopy
import matplotlib.pyplot as plt
import metrics_mod

save_type = 'pdf'
folder_name = './' #FOLDER WHERE NUMERICAL EXPERIMENTS ARE SAVED
save_folder = './' #FOLDER WHERE FIGURES ARE SAVED

try:
    arg = sys.argv
    ci_test = str(arg[1]) #par_corr
    variant = str(arg[2]) #To choose from "sample_size_highdegree" "highdim_highdegree"
                          # "tau_max_highdegree" "autocorr_highdegree"
except:
    arg = ''
    ci_test = 'par_corr'
    variant = 'sample_size_highdegree'

print(variant)
name = {'par_corr':r'ParCorr', 'gp_dc':r'GPDC'}

def f1score(precision_,recall_):
    return 2 * (precision_ * recall_) / (precision_ + recall_)

def get_metrics_from_file(para_setup):

    name_string = '%s-'*len(para_setup)  # % para_setup
    name_string = name_string[:-1]

    try:
        print("load from metrics file  %s_metrics.dat " % (folder_name + name_string % tuple(para_setup)))
        results = pickle.load(open(folder_name + name_string % tuple(para_setup) + '_metrics.dat', 'rb'), encoding='latin1')
    except Exception as e:
        print('failed from metrics file '  , tuple(para_setup))
        print(e)
        return None

    return results

def print_time(seconds, precision=1):
    if precision == 0:
        if seconds > 60*60.:
            return "%.0fh" % (seconds/3600.)
        elif seconds > 60.:
            return "%.0fmin" % (seconds/60.)
        else:
            return "%.0fs" % (seconds)
    else:
        if seconds > 60*60.:
            return "%.1fh" % (seconds/3600.)
        elif seconds > 60.:
            return "%.1fmin" % (seconds/60.)
        else:
            return "%.1fs" % (seconds)

def print_time_std(time, precision=1):

    mean = time.mean()
    std = time.std()
    if precision == 0:
        if mean > 60*60.:
            return r"%.0f$\pm$%.0fh" % (mean/3600., std/3600.)
        elif mean > 60.:
            return r"%.0f$\pm$%.0fmin" % (mean/60., std/60.)
        else:
            return r"%.0f$\pm$%.0fs" % (mean, std)
    else:
        if mean > 60*60.:
            return r"%.1f$\pm$%.1fh" % (mean/3600., std/3600.)
        elif mean > 60.:
            return r"%.1f$\pm$%.1fmin" % (mean/60., std/60.)
        else:
            return r"%.1f$\pm$%.1fs" % (mean, std)

def draw_it(paras, which):

    figsize = (4, 3)
    capsize = .5
    marker1 = 'o'
    marker2 = 's'
    marker3 = '+'
    alpha_marker = 1.

    params = {
           'legend.fontsize': 5,
           'legend.handletextpad': .05,
           'lines.color':'black',
           'lines.linewidth':.5,
           'lines.markersize':2,
            'xtick.labelsize':6.5,
            'xtick.major.pad'  : 1,
            'xtick.major.size' : 2,
            'ytick.major.pad'  : 1,
            'ytick.major.size' : 2,
            'ytick.labelsize':6.5,
            'axes.labelsize':8,
            'font.size':8,
            'axes.labelpad':2,
            'axes.spines.right' : False,
            'axes.spines.top' : False,
            }
    mpl.rcParams.update(params)
    fig = plt.figure(figsize=figsize,dpi=600)
    gs = fig.add_gridspec(1, 2)
    ax1a = fig.add_subplot(gs[0, 0])
    ax2a = fig.add_subplot(gs[0, 1])

    anova = {}
    auc_dict = {}
    fpr_dict = {}
    pr_df = pd.DataFrame(columns=['method', 'n_bs', 'param', 'pc_alpha','recall','precision','contemp_recall',
    'contemp_precision'])

    for method in methods:
        anova[method] = {}
        anova[method]['adj_anylink_tpr'] = []
        anova[method]['adj_anylink_fpr'] = []
        anova[method]['adj_lagged_tpr'] = []
        anova[method]['adj_lagged_fpr'] = []
        anova[method]['adj_contemp_tpr'] = []
        anova[method]['adj_contemp_fpr'] = []
        anova[method]['edgemarks_contemp_recall'] = []
        anova[method]['edgemarks_contemp_precision'] = []
        anova[method]['computation_time'] = []

        if "bootstrap" in method:
            n_bs_method = n_bs
        else:
            n_bs_method= [0]

        for idx,n_bs_here in enumerate(n_bs_method):
            for para in paras:
                for pc_alpha in pc_alpha_here:
                    pc_alpha = np.format_float_positional(float(pc_alpha),trim='-')
                    if which == 'auto':
                        auto_here = para
                        N_here = N
                        tau_max_here = tau_max
                        frac_unobserved_here = frac_unobserved
                        T_here = T

                    elif which == 'N':
                        N_here = para
                        auto_here = auto
                        tau_max_here = tau_max
                        frac_unobserved_here = frac_unobserved
                        T_here = T

                    elif which == 'tau_max':
                        N_here = N
                        auto_here = auto
                        tau_max_here = para
                        frac_unobserved_here = frac_unobserved
                        T_here = T

                    elif which == 'sample_size':
                        N_here = N
                        auto_here = auto
                        tau_max_here = tau_max
                        frac_unobserved_here = frac_unobserved
                        T_here = para

                    n_links_here = links_from_N(N_here)

                    para_setup = (model, N_here, n_links_here, min_coeff, coeff, auto_here, contemp_fraction, frac_unobserved_here,
                                    max_true_lag, T_here, ci_test, method, pc_alpha, tau_max_here, n_bs_here)
                    metrics_dict = get_metrics_from_file(para_setup)
                    if metrics_dict is not None:
                        current_res = pd.DataFrame([[
                                    method, n_bs_here, para, pc_alpha, metrics_dict['adj_anylink_recall'][0], metrics_dict['adj_anylink_precision'][0],
                                    metrics_dict['edgemarks_contemp_recall'][0],metrics_dict['edgemarks_contemp_precision'][0]
                                    ]],
                                    columns=['method', 'n_bs', 'param', 'pc_alpha','recall','precision','contemp_recall','contemp_precision'])
                        pr_df = pd.concat([pr_df,current_res])

    #Find common recall coverage for all parameters settings (methods, varying_param)
    pc_alpha_min, pc_alpha_max = min(pc_alpha_here),max(pc_alpha_here)
    pc_alpha_min = np.format_float_positional(float(pc_alpha_min),trim='-')
    pc_alpha_max = np.format_float_positional(float(pc_alpha_max),trim='-')
    lower_common_recall = pr_df[pr_df.pc_alpha == pc_alpha_min].max().recall
    upper_common_recall = pr_df[pr_df.pc_alpha == pc_alpha_max].min().recall
    print(pr_df[pr_df.pc_alpha == pc_alpha_min])
    pr_df[pr_df.pc_alpha == pc_alpha_max]
    pc_alpha_contemp_min, pc_alpha_contemp_max = min(pc_alpha_here),0.2
    pc_alpha_contemp_min = np.format_float_positional(float(pc_alpha_contemp_min),trim='-')
    pc_alpha_contemp_max = np.format_float_positional(float(pc_alpha_contemp_max),trim='-')
    lower_common_contemp_recall = pr_df[pr_df.pc_alpha == pc_alpha_contemp_min].max().contemp_recall
    upper_common_contemp_recall = pr_df[pr_df.pc_alpha == pc_alpha_contemp_max].min().contemp_recall
    print(pr_df[pr_df.pc_alpha == pc_alpha_contemp_min])
    print(pr_df[pr_df.pc_alpha == pc_alpha_contemp_max])
    for method in methods:
        if "bootstrap" in method:
            n_bs_method = n_bs
        else:
            n_bs_method= [0]
        for idx,n_bs_here in enumerate(n_bs_method):
            for para in paras:
                print(para)
                fx = []
                x=[]
                fx_c = []
                x_c=[]
                for pc_alpha in pc_alpha_here:
                    pc_alpha = np.format_float_positional(float(pc_alpha),trim='-')
                    data_query = pr_df.query('method==@method and n_bs==@n_bs_here and param==@para and pc_alpha == @pc_alpha')
                    x.append(data_query.recall.values[0])
                    fx.append(data_query.precision.values[0])
                    if float(pc_alpha) <= 0.1:
                        x_c.append(data_query.contemp_recall.values[0])
                        fx_c.append(data_query.contemp_precision.values[0])
                lower_interp_precision = np.interp(lower_common_recall,x,fx)
                upper_interp_precision = np.interp(upper_common_recall,x,fx)
                lower_interp_precision_c = np.interp(lower_common_contemp_recall,x_c,fx_c)
                upper_interp_precision_c = np.interp(upper_common_contemp_recall,x_c,fx_c)
                final_recall = [lower_common_recall]
                final_precision = [lower_interp_precision]
                final_recall_c = [lower_common_contemp_recall]
                final_precision_c = [lower_interp_precision_c]
                for ix in range(len(x)):
                    if upper_common_recall>x[ix]>lower_common_recall:
                        final_recall.append(x[ix])
                        final_precision.append(fx[ix])
                for ix_c in range(len(x_c)):
                    if upper_common_contemp_recall>x_c[ix_c]>lower_common_contemp_recall:
                        final_recall_c.append(x_c[ix_c])
                        final_precision_c.append(fx_c[ix_c])
                final_recall.append(upper_common_recall)
                final_precision.append(upper_interp_precision)
                final_recall_c.append(upper_common_contemp_recall)
                final_precision_c.append(upper_interp_precision_c)

                auc_value = auc(final_recall,final_precision)
                auc_c_value = auc(final_recall_c,final_precision_c)
                if "bootstrap" in method:
                    para_plot = paras.index(para) + (methods.index(method)+idx)/float(len(n_bs)+len(methods)-1)*.6
                else:
                    para_plot = paras.index(para) + (methods.index(method)+idx)/float(len(n_bs)+len(methods)-1)*.6

                ax1a.errorbar(para_plot, auc_value, capsize=capsize,  alpha=alpha_marker,
                             color=color_picker(method,idx), marker=marker2, linestyle='dashed')
                ax2a.errorbar(para_plot, auc_c_value, capsize=capsize,  alpha=alpha_marker,
                             color=color_picker(method,idx), marker=marker2, linestyle='dashed')

    axes = {'ax1a':ax1a, 'ax2a':ax2a}
    for axname in axes:

        ax = axes[axname]

        if which == 'N':
            ax.set_xlim(-0.5, len(paras))
            if ci_test == 'par_corr':
                ax.xaxis.set_ticks([paras.index(p) for p in paras] )
                ax.xaxis.set_ticklabels([str(p) for p in paras] )
            else:
                ax.xaxis.set_ticks([paras.index(p) for p in paras] )
                ax.xaxis.set_ticklabels([str(p) for p in paras] )

        elif which == 'auto':
            ax.xaxis.set_ticks([paras.index(p) for p in paras] )
            ax.xaxis.set_ticklabels([str(p) for p in paras] )
        elif which == 'pc_alpha':
            ax.set_xlim(0, len(paras))
            ax.xaxis.set_ticks([paras.index(p) for p in paras])
            ax.set_xticklabels([str(p) for p in paras], rotation=75, horizontalalignment="center")

        elif which == 'tau_max':
            ax.set_xlim(-0.5, len(paras))
            ax.xaxis.set_ticks([paras.index(p) for p in paras] )
            ax.xaxis.set_ticklabels([str(p) for p in paras] )

        elif which == 'unobserved':
            ax.set_xlim(0, len(paras))
            ax.xaxis.set_ticks([paras.index(p) for p in paras] )
            ax.xaxis.set_ticklabels([str(p) for p in paras] )

        elif which == 'sample_size':
            print(paras)
            ax.set_xlim(0, len(paras))
            ax.xaxis.set_ticks([paras.index(p) for p in paras] )
            ax.xaxis.set_ticklabels([str(p) for p in paras] )

        # ax.set_xlabel(xlabel, fontsize=8)
        for line in ax.get_lines():
            line.set_clip_on(False)
        # Disable spines.
        # if not 'ax3a' in axname:
        ax.spines['right'].set_color('none')
        ax.spines['top'].set_color('none')
        ax.spines['left'].set_position(('outward', 3))
        ax.spines['bottom'].set_position(('outward', 3))

        ax.grid(axis='y', linewidth=0.3)

        pad = 3

        if axname == 'ax1a':
            ax.set_title('Adj. precision-recall AUC', fontsize=10, pad=pad)
            ax.yaxis.set_major_formatter(FormatStrFormatter('%.2f'))

        elif axname == 'ax2a':
            ax.yaxis.set_major_formatter(FormatStrFormatter('%.2f'))
            ax.set_title('Contemp. \n precision-recall  AUC', fontsize=10, pad=pad)

    axlegend = fig.add_axes([0.0, .93, 1., .05])
    axlegend.axis('off')
    for method in methods:
        if "bootstrap" in method:
            for idx,n_bs_here in enumerate(n_bs):
                method_label ="Bagged("+str(n_bs_here)+")-PCMCI+"
                color_ = color_picker(method,idx)

                axlegend.errorbar([], [], linestyle='',
                capsize=capsize, label=method_label,
                color= color_, marker='s')
        else:
            method_label ="PCMCI+"
            color_= color_picker(method)
            axlegend.errorbar([], [], linestyle='',
            capsize=capsize, label=method_label,
            color=color_, marker='s')

    if not 'paper' in variant:
        ncol = 4
        fontsize = 10
    else:
        ncol = 4
        fontsize = 10
    axlegend.legend(ncol=ncol,
             loc='lower left',
            markerscale=3,
            columnspacing=.75,
            labelspacing=.01,
            fontsize=fontsize, framealpha=.5
            )
    if not 'paper' in variant:
        ncol = 5
        fontsize = 6
        ncol2 =3
        ncolroc=1
    else:
        ncol = 5
        fontsize = 6
        ncol2 =3
        ncolroc=1

    axlegend.legend(ncol=ncol,
             loc='lower left',
            markerscale=3,
            columnspacing=.75,
            labelspacing=.01,
            fontsize=fontsize, framealpha=.5
            )

    #produce letter mark top left
    if 'paper' in variant:
        if 'autocorr' in variant:
            fig.text(0., 0.93, "C", fontsize=12, fontweight='bold',
                ha='left', va='top')
        elif 'highdim' in variant:
            fig.text(0., 0.93, "A", fontsize=12, fontweight='bold',
                ha='left', va='top')
        elif 'sample_size' in variant:
            fig.text(0., 0.93, "B", fontsize=12, fontweight='bold',
                ha='left', va='top')
        elif 'tau_max' in variant:
            fig.text(0., 0.93, "D", fontsize=12, fontweight='bold',
                ha='left', va='top')

    if 'random_lineargaussian' in model:
        model_name = r"$\mathcal{N}$"
    elif 'random_linearmixed' in model:
        model_name = r"$\mathcal{N{-}W}$"
    elif 'random_nonlinearmixed' in model:
        model_name = r"$\mathcal{N{-}W}^2$"
    elif 'random_nonlineargaussian' in model:
        model_name = r"$\mathcal{N}^2$"
    else:
        model_name = model
    if 'fixeddensity' in model:
        model_name += r"$_{d{=}0.3}$"
    elif 'highdegree' in model:
        model_name += r"$_{d{=}1.5}$"

    if which == 'N':
        fig.text(0.5, 0., r"Number of variables $N$", fontsize=10,
            horizontalalignment='center', va='bottom')
        fig.text(1., 1., r"%s: $T=%d, a=%s$, " %(model_name, T, auto)
                            + r"%s, $\tau_{\max}=%d$" %(name[ci_test],tau_max),
         fontsize=7.5, ha='right', va='top')
    elif which == 'auto':
        fig.text(0.5, 0., r"Autocorrelation $a$", fontsize=10,
            horizontalalignment='center', va='bottom')
        fig.text(1., 1., r"%s: $N=%d, T=%d$" %(model_name, N, T)
                        + r", %s, $\tau_{\max}=%d$" %(name[ci_test], tau_max),
         fontsize=8, ha='right', va='top')

    elif which == 'tau_max':
        fig.text(0.5, 0., r"Time lag $\tau_{\max}$", fontsize=10,
            horizontalalignment='center', va='bottom')
        fig.text(1., 1., r"%s: $N=%d, T=%d, a=%s$" %(model_name, N, T, auto)
                         + r", %s" %(name[ci_test]),
         fontsize=8, ha='right', va='top')
    elif which == 'unobserved':
        fig.text(0.5, 0., r"Frac. unobserved", fontsize=9,
            horizontalalignment='center', va='bottom')
        fig.text(1., 1., r"%s: $N=%d, T=%d, a=%s, \tau_{\max}=%d$" %(model_name, N, T, auto)
                            +"\n" + r"%s, $\alpha=%s, \tau_{\max}=%d$" %(name[ci_test], pc_alpha, tau_max),
         fontsize=8, ha='right', va='top')
    elif which == 'sample_size':
        fig.text(0.5, 0., r"Sample size $T$", fontsize=10,
            horizontalalignment='center', va='bottom')
        fig.text(1., 1., r"%s: $N=%d, a=%s$" %(model_name, N, auto)
                        + r", %s, $\tau_{\max}=%d$" %(name[ci_test], tau_max),
         fontsize=8, ha='right', va='top')
    elif which == 'pc_alpha':
        fig.text(0.5, 0., r"$\alpha$", fontsize=9,
            horizontalalignment='center', va='bottom')

        fig.text(1., 1., r"%s: $N=%d, T=%d, a=%s$" %(model_name, N, T, auto)
                        + r", %s, $\tau_{\max}=%d$" %(name[ci_test], tau_max),
         fontsize=6, ha='right', va='top')
    fig.subplots_adjust(left=0.07, right=0.93, hspace=.3, bottom=0.12, top=0.85, wspace=.3)
    print("Saving plot to %s" %(save_folder + '%s.%s' %(save_suffix, save_type)))
    fig.savefig(save_folder + '%s.%s' %(save_suffix, save_type))
    plot_files.append(save_folder + '%s.%s' %(save_suffix, save_type))

def adjust_lightness(color, amount=0.5):
    #Function to create a color shading
    # amount <1 darkens the color, amount >1 lightens the color
    import matplotlib.colors as mc
    import colorsys
    try:
        c = mc.cnames[color]
    except:
        c = color
    c = colorsys.rgb_to_hls(*mc.to_rgb(c))
    return colorsys.hls_to_rgb(c[0], max(0, min(1, amount * c[1])), c[2])

def color_picker(method,idx=0):

    if "bootstrap" in method:
        color_scaling = [adjust_lightness('green', amount=i) for i in [3.4,2.4,1,0]]
        return color_scaling[2]
    else: return 'orange'

def method_label(method):
    # return method
    if not 'paper' in variant:
        return method

    if 'standard_pcmci' in method:
        if 'allpx0' in method:
            return r'PCMCI$^+_0$'
        elif 'laggedpx0' in method:
            return r'PCMCI$^+_{0-}$'
        elif 'resetlagged' in method:
           return r'PCMCI$^+_{\rm reset}$'
        else:
            return r'PCMCI$^+$'

    elif 'residualPC' in method or 'GCresPC' in method:
        return r"GCresPC"
    elif 'lingam' in method:
         return r"LiNGAM"
    elif 'pcalg' in method:
         return r"PC"
    else:
        return method

def links_from_N(num_nodes):

    # if 'highdim' in variant and :
    #     return num_nodes
    if num_nodes == 2:
        return 1

    if 'fixeddensity' in model:
        return max(num_nodes, int(0.2*num_nodes*(num_nodes-1.)/2.))   # CONSTANT LINK DENSITY 0.2 FOR NON_TIME SERIES !!!
    elif 'highdegree' in model:
        return int(1.5*num_nodes)
    else:
        return num_nodes

if __name__ == '__main__':

    save_type = 'pdf'
    paper = True
    fpr_precision = 'precision'

    if 'versions' in variant:
        methods = [
            'standard_pcmci+',
            'bootstrap_pcmci+'
            ]
    else:
        methods = []

        methods += [
            'standard_pcmci+',
            'bootstrap_pcmci+',
            ]

    if ci_test == 'par_corr':
        if 'mixed' in variant:
            model  = 'random_linearmixed'  # random_lineargaussian random_linearmixed random_nonlinearmixed
        else:
            model  = 'random_lineargaussian'  # random_lineargaussian random_linearmixed random_nonlinearmixed
    else:
        if 'mixed' in variant:
            model  = 'random_nonlinearmixed'  # random_lineargaussian random_linearmixed random_nonlinearmixed
        else:
            model  = 'random_nonlineargaussian'  # random_lineargaussian random_linearmixed random_nonlinearmixed

    if 'fixeddensity' in variant:
        model += '_fixeddensity'
    elif 'highdegree' in variant:
        model += '_highdegree'

    if  'autocorr' in variant:
        if ci_test == 'par_corr':
            T_here = [500]
            N_here = [5]
            num_rows = 4
        else:
            T_here = [500]
            N_here = [5]
            num_rows = 3

        tau_max = 5
        # vary_auto = [0., 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.95, 0.98, 0.99, 0.995, 0.999]
        vary_auto = [ 0.4, 0.5, 0.6, 0.8, 0.9, 0.95, 0.98]#, 0.99, 0.999]#[0., 0.2,

        pc_alpha_here = [0.00001,0.00005,0.0001,0.0005,0.001,0.005,0.01,0.02,0.05,0.1,0.2,0.3,
                                    0.4, 0.5, 0.6, 0.8, 0.9, 0.999] #

        min_coeff = 0.1
        coeff = 0.5
        frac_unobserved = 0.
        contemp_fraction = 0.3
        max_true_lag = 5
        # T = 500
        n_bs= [100]

        pdfjam_suffix = variant + "_" + str(ci_test)
        for T in T_here:
         for N in N_here:
            n_links = links_from_N(N)
            para_setup_name = (variant, N, n_links, min_coeff, coeff, contemp_fraction, frac_unobserved,
                            max_true_lag, T, ci_test, "auc", tau_max,n_bs)

            save_suffix = '%s-'*len(para_setup_name) % para_setup_name
            save_suffix = save_suffix[:-1]

            print(save_suffix)
            draw_it(paras=vary_auto, which='auto')

            if test:
                sys.exit(0)

    elif 'highdim' in variant:
        if ci_test == 'par_corr':
            T_here = [500]
            vary_N =   [3,5,10,]
            auto_here = [0.95]
            num_rows = 4
        else:
            T_here = [500] #, 1000]
            vary_N =  [2,3,5,8,10,15]
            auto_here = [0.95]
            num_rows = 4

        contemp_fraction = 0.3
        frac_unobserved = 0.
        max_true_lag = 5
        tau_max = 5

        min_coeff = 0.1
        coeff = 0.5
        pc_alpha_here = [0.00001,0.00005,0.0001,0.0005,0.001,0.005,0.01,0.02,0.05,0.1,0.2,0.3,
                                    0.4, 0.5, 0.6, 0.8, ]
        n_bs = [100]

        pdfjam_suffix = variant + "_" + str(ci_test)

        for T in T_here:
            for auto in auto_here:
                para_setup_name = (variant, min_coeff, coeff, auto, contemp_fraction, frac_unobserved, max_true_lag, T,
                                        ci_test, 'auc', tau_max, n_bs)
                save_suffix = '%s-'*len(para_setup_name) % para_setup_name
                save_suffix = save_suffix[:-1]
                draw_it(paras=vary_N, which='N')

                if test:
                    sys.exit(0)

    elif  'sample_size' in variant:

        if ci_test == 'par_corr':
            vary_T = [100,200,500,1000]
            N_here = [5]
            auto_here = [0.95]
            num_rows = 4
        else:
            vary_T = [200, 500]
            N_here = [5, 10]
            auto_here = [0., 0.5, 0.9, 0.95]
            num_rows = 4

        min_coeff = 0.1
        coeff = 0.5
        pc_alpha_here = [0.00001,0.00005,0.0001,0.0005,0.001,0.005,0.01,0.02,0.05,0.1,0.2,0.3,
                                    0.4, 0.5, 0.6, 0.8, 0.9, 0.999]
        contemp_fraction = 0.3
        frac_unobserved = 0.

        max_true_lag = 5
        tau_max = 5
        pdfjam_suffix = variant + "_" + str(ci_test)
        n_bs = [100]
        for N in N_here:
          n_links = links_from_N(N)
          for auto in auto_here:
                para_setup_name = (variant, N, n_links, min_coeff, coeff, contemp_fraction, frac_unobserved,
                                    max_true_lag, auto, ci_test, "auc", tau_max, n_bs)
                save_suffix = '%s-'*len(para_setup_name) % para_setup_name
                save_suffix = save_suffix[:-1]

                print(save_suffix)
                draw_it(paras=vary_T, which='sample_size')

                if test:
                    sys.exit(0)
    if 'tau_max' in variant:

        if ci_test == 'par_corr':
            T_here = [500]
            N_here = [5]
            auto_here = [0.95]
            # vary_N =  [5, 10, 20, 30, 40]
            num_rows = 4
            vary_tau_max = [5, 10, 15, 20, 25, 30, 35, 40]

        else:
            T_here = [200, 500]
            N_here = [5]
            auto_here = [0., 0.5, 0.9, 0.95]
            num_rows = 4
            vary_tau_max = [5, 10, 15, 20]

        min_coeff = 0.1
        coeff = 0.5
        pc_alpha_here = [0.00001,0.00005,0.0001,0.0005,0.001,0.005,0.01,0.02,0.05,0.1,0.2,0.3,
                                    0.4, 0.5, 0.6]

        contemp_fraction = 0.3
        frac_unobserved = 0.
        max_true_lag = 5
        n_bs = [100]
        pdfjam_suffix = variant + "_" + str(ci_test)

        for T in T_here:

         for N in N_here:
          n_links = links_from_N(N)
          for auto in auto_here:
                para_setup_name = (variant, N, n_links, min_coeff, coeff, contemp_fraction, frac_unobserved,
                                    max_true_lag, auto, T, ci_test, "auc", n_bs)
                save_suffix = '%s-'*len(para_setup_name) % para_setup_name
                save_suffix = save_suffix[:-1]

                print(save_suffix)
                draw_it(paras=vary_tau_max, which='tau_max')


import open3d as o3d
o3d.utility.set_verbosity_level(o3d.utility.VerbosityLevel.Error)
import numpy as np
from scipy.spatial.transform import Rotation as R

import time
import copy
import math

import os
import sys
from datetime import datetime as dt

# An installation agnostic method to find and link to root of the package which is mlfactory
#==========================================================
import re
try: #testing the functions locally without pip install
  import __init__
  cimportpath = os.path.abspath(__init__.__file__)
  if 'extensions' in cimportpath:
    print("Non local testing ")
    import mlfactory
    cimportpath = os.path.abspath(mlfactory.__file__)+'/applications/deep_modular_scene_mapper/__init__.py'

except: #testing while mlfactory is installed using pip
  print("Non local testing")
  import mlfactory
  cimportpath = os.path.abspath(mlfactory.__file__)+'/applications/deep_modular_scene_mapper/__init__.py'

idxlist = [m.start() for m in re.finditer(r"/", cimportpath)]
invoking_submodule = cimportpath[idxlist[-2]+1:idxlist[-1]]
print("In deep_modular_scene_mapper/main.py got invoking submodule using re",invoking_submodule)
main_package_loc = cimportpath[:cimportpath.rfind('mlfactory')+len('mlfactory')]
print("In deep_modular_scene_mapper/main.py got main package location ",main_package_loc)

os.environ['applications'] = main_package_loc+'/applications'
os.environ['top'] = main_package_loc
sys.path.append(os.path.join(os.environ['applications']))
sys.path.append(os.path.join(os.environ['top']))
#==========================================================

#from superglue_inference.match_pair import matcher

from applications.superglue_inference import match_pair
from applications.deep_modular_scene_mapper.tools.rigid_transform_3D import rigid_transform_3D
from applications.deep_modular_scene_mapper.tools import point_plane_icp
from applications.deep_modular_scene_mapper.tools.depth_estimator import monodepth

from applications.deep_modular_scene_mapper.tools import project_rgbd
from applications.deep_modular_scene_mapper.tools import sift_matching
from applications.deep_modular_scene_mapper.tools import pose_graph_optimizer

def create_pcd_from_points(points,colors):
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points)
    pcd.colors = o3d.utility.Vector3dVector(colors)
    return pcd

def join_pcds(pcdlist):
  all_points = []
  all_colors = []

  for sub_pcd in pcdlist:
      p1 = np.asarray(sub_pcd.points)
      c1 = np.asarray(sub_pcd.colors)

      if all_points==[]:
          all_points = p1
          all_colors = c1
      else:
          all_points = np.concatenate((all_points,p1), axis=0)
          all_colors = np.concatenate((all_colors,c1), axis=0)

  pcd = create_pcd_from_points(all_points,all_colors)
  return pcd

class mapper(object):
  def __init__(self, data_location = '/datasets/sample_videos/extracted/', frame_numbers = [0,10,20,30], init_pose = np.array([[1.0,0.0,0.0,0.0], [0.0,1.0,0.0,0.0], [0.0,0.0,1.0,0.0], [0.0,0.0,0.0,1.0] ] ) ):
    self.data_location = data_location
    self.feature_matcher = match_pair.matcher()
    self.md = monodepth() #the network that estimates mono depth
    self.voxel_size = 0.01
    self.init_pose = init_pose

    #camera cannot move extremely fast there needs to be continuity so maximum change in angles
    #and translations must be bounded according to some assumptions on how the robot moves
    self.geo_constraints = [500.0,500.0] # sum of squares of roll pitch yaw, sum of squares of translation x y and z

    #try to remove overlapping pcd points while registering
    self.combine_depth_feature_matching = False #uses sift like feature matching across depth images as well
    self.viz_load_pairs = False

    #gap between indices should be smaller when motion is faster
    #self.indices = [340, 350, 360, 370, 380, 390, 400, 410, 420, 430 ]
    self.indices = frame_numbers

    self.iphone_back_camera_params = {
                                "fx": 520.3,
                                "fy": 520.3,
                                "centerX": 320.0,
                                "centerY": 240.0,
                                "scalingFactor": 1
                            }

    self.transform_sequence = []
    self.correction_transforms = []

    self.pcds = []
    self.mapped_pcds = [] #stores the final perfectly aligned map
    self.pose_sequence = [] #stores the final perfectly aligned poses
    self.rgbds = []
    self.depth_images = []

    #initialize poses
    self.poses = [ np.array([[1.0,0.0,0.0,0.0], [0.0,1.0,0.0,0.0], [0.0,0.0,1.0,0.0], [0.0,0.0,0.0,1.0] ] ) ]

    print("estimating first frame depth ... ")
    self.rgb1, self.depth1, self.max_d_init = self.md.predict_depth(self.data_location+str(self.indices[0])+".jpg")

    print("rgb and depth image shapes ",self.rgb1.shape, self.depth1.shape)

  def estimate_pointcloud_t1t2(self, idx):
    print("time ",dt.now())

    self.index_pair = [self.indices[idx], self.indices[idx+1]]
    #===================================================================
    #load source and target rgbd images
    print("loading rgb images, predicting depths and making initial pcds for the pair ",self.index_pair)
    #rgbs,depths = project_rgbd.load_verify_extracted(data_location, index_pair)

    #===================================================================
    #actual resized rgb image, estimated depth , and also the maximum value of depth according to a global scale
    if idx!=0:
      self.rgb1, self.depth1, self.max_d1 = copy.copy(self.rgb2), copy.copy(self.depth2), copy.copy(self.max_d2)
    else:
      self.max_d1 = copy.copy(self.max_d_init)

    #estimate for the adjacent time step
    self.rgb2, self.depth2, self.max_d2 = self.md.predict_depth(self.data_location+str(self.index_pair[1])+".jpg")

    #rescale depths back to absolute depths
    self.absd1 = (self.max_d1/self.max_d_init)*self.depth1
    self.absd2 = (self.max_d2/self.max_d_init)*self.depth2

    self.depth_images.append(o3d.geometry.Image(self.absd1))
    self.rgbds.append([self.rgb1, self.absd1])

    self.rgbs = [self.rgb1,self.rgb2]
    self.depths = [self.depth1,self.depth2]

    #project_rgbd.show_pcd([pcd1,pcd2])
    #===================================================================

    #end result of this step- get the pointclouds that now need to be aligned
    print("converting rgbd to pointclouds")
    self.pcd1, self.pixel_points1 = project_rgbd.project_rgbd_to_pointcloud(self.rgbs[0][:,:,::-1],(self.max_d1/self.max_d_init)*self.depths[0], self.iphone_back_camera_params)
    self.pcd2, self.pixel_points2 = project_rgbd.project_rgbd_to_pointcloud(self.rgbs[1][:,:,::-1],(self.max_d2/self.max_d_init)*self.depths[1], self.iphone_back_camera_params)

    #initialize poses of the pointclouds
    #self.pose1 = np.array([[1.0,0.0,0.0,0.0], [0.0,1.0,0.0,0.0], [0.0,0.0,1.0,0.0], [0.0,0.0,0.0,1.0] ] )
    #self.pose2 = np.array([[1.0,0.0,0.0,0.0], [0.0,1.0,0.0,0.0], [0.0,0.0,1.0,0.0], [0.0,0.0,0.0,1.0] ] )

    self.pose1 = self.init_pose
    self.pose2 = self.init_pose
    #diagnostic
    if self.viz_load_pairs:
      project_rgbd.show_pcd([self.pcd1,self.pcd2])

  def find_rgb_feature_matchest1t2(self):
    #===================================================================
    #find pairs of matching pixel points between consequtive images (for ex using sift matching), randomly chose half of those matches
    #find the xyz point of these pairs by cross checking with depth map
    #use P3P algorithm to find rigid transformation
    #if wrong points are chosen as matching pixel points, the transformation is going to be abrupt and will probably violate geo constraints
    #if geo constraints are violated try again back from step 1, keep doing loop until geo constraints match
    inflate = 0
    for _ in range(10): #try atmost 10 times

      #===================================================================
      #extract corresponding points using sift based feature matching in rgb domain
      print("extracting image features using superglue ")
      #s,d = sift_matching.compare_extra(data_location+str(index_pair[0])+".png", data_location+str(index_pair[1])+".png")

      i1,i2 = self.feature_matcher.prepare_images(self.rgbs[0], self.rgbs[1])
      s,d = self.feature_matcher.match(i1,i2, viz = False, ransac_refine = True)
      print("got number of matches rgb ",s.shape[0], d.shape[0])

      if self.combine_depth_feature_matching: #also add sift like feature match between mono depth estimated images (does not generally work well)
        d1,d2 = self.feature_matcher.prepare_images(self.depths[0],self.depths[1])
        s_,d_ = self.feature_matcher.match(d1,d2, viz = False, ransac_refine = True)
        print("got number of matches depth ",s_.shape[0], d_.shape[0])

        #trying to merge matching keypoints calculated on depth frames as well, not much benefit
        s = np.concatenate((s,s_))
        d = np.concatenate((d,d_))

      #===================================================================
      #get exact 3d corresponding points by mapping rgbs back to their pcds
      print("Doing rough alignment using least squares matching based on superglue matches")
      source_cor = []
      dest_cor = []

      n_frac_matches = 2
      if s.shape[0]<6:
        n_frac_matches = 1 #alg cannot work with too few matches

      rand_chose = np.random.choice(s.shape[0], s.shape[0]//n_frac_matches)

      #for i in range(s.shape[0]):
      for i in rand_chose:
        u,v = s[i]
        p,q = d[i]

        if list(self.pixel_points1[int(u),int(v)])[2]!=0.0 and list(self.pixel_points2[int(p),int(q)])!=0.0:
          source_cor.append( list(self.pixel_points1[int(u),int(v)]) )
          dest_cor.append( list(self.pixel_points2[int(p),int(q)]) )

      #===================================================================

      #This part uses the feature matches from rgb and finds the xyz in the corresponding pointclouds
      #using pairs of xyz matches now we can use P3P algorithm to find rough rotation and translation
      #===================================================================
      #estimate the rigid transform using SVD
      ret_R, ret_t = rigid_transform_3D(np.array(dest_cor).T, np.array(source_cor).T)
      print("estimated rotation and translation ",ret_R, ret_t)
      #===================================================================
      #store the final best rot and trans as the end result of this function
      self.ret_R, self.ret_t = ret_R, ret_t

      r = R.from_matrix(ret_R)
      pyr = r.as_euler('zyx', degrees=True)
      print("estimated pitch, yaw, roll change ",pyr)

      magn_r = 0.0
      magn_t = 0.0
      for e in pyr:
          magn_r+=e**2
      for e in ret_t:
          magn_t+= e**2
      print("got magnitudes of rot and trans transformations ",magn_r, magn_t)

      if magn_r<self.geo_constraints[0]+inflate and magn_t<self.geo_constraints[1]+inflate:
        print("transformation matched geometric motion constraint assumption ! exiting attempts")
        break
      else:
        print("proposed transformation did not match geo constraints !! ")
        inflate +=1

  def align_pointcloudst1t2(self):
    #===================================================================
    #register destination pointcloud to source by rgb matched festures transformation and show the result
    print("Rough transform last pointcloud of the most recent pair")
    #roll all the pointclouds so far through all the prior transformations
    if self.transform_sequence!=[]:
      count = 1
      prev_pose = self.poses[0]
      for t in range(len(self.transform_sequence)):
        self.pcd1 = self.pcd1.transform(self.transform_sequence[t])
        self.pcd2 = self.pcd2.transform(self.transform_sequence[t])

        self.pose1 = np.dot(self.pose1, self.transform_sequence[t])
        self.pose2 = np.dot(self.pose2, self.transform_sequence[t])

        if len(self.poses)>2 and t!=len(self.transform_sequence)-1:
          self.poses[-1] = np.dot(self.poses[-1], self.transform_sequence[t])

    #apply current transformation
    T = np.eye(4)
    T[:3, :3] = self.ret_R
    T[0, 3] = self.ret_t[0][0]
    T[1, 3] = self.ret_t[1][0]
    T[2, 3] = self.ret_t[2][0]

    self.pcd2 = self.pcd2.transform(T)
    self.pose2 = np.dot(self.pose2, T)

    print("Now showing registered pcds ")
    #project_rgbd.show_pcd([pcd1,pcd2])
    #poses.append(poses[-1])
    self.transform_sequence.append(T)
    self.poses.append(T)
    #poses[-1] = np.dot(poses[-1], T)

    if self.pcds==[]:
      self.pcds.append(self.pcd1)
      self.pcds.append(self.pcd2)

    else:
      self.pcds[-1] = self.pcd1
      self.pcds.append(self.pcd2)

    #o3d.io.write_point_cloud(str(indices[idx])+"_T.pcd", pcds[idx])

    #further refine the alignment by doing additional ICP step for the two pointclouds
    rpcd1, rpcd2, correction_transform = point_plane_icp.register_pair(self.pcds[-2],self.pcds[-1], viz = False)

    self.pose1 = np.dot(self.pose1, correction_transform)
    if self.mapped_pcds==[]:
      self.mapped_pcds.append(rpcd1)
      self.mapped_pcds.append(rpcd2)

      self.pose_sequence.append(self.pose1)
      self.pose_sequence.append(self.pose2)

    else:

      self.mapped_pcds[-1] = rpcd1
      self.mapped_pcds.append(rpcd2)

      self.pose_sequence[-1] = self.pose1
      self.pose_sequence.append(self.pose2)

    self.correction_transforms.append(correction_transform)

    #poses[-1] = np.matmul(poses[-1], correction_transform)

    #backpropagate the last correction transform to all previous mapped pcds other than the most recent pair
    if self.correction_transforms!=[]:
        for i in range(len(self.mapped_pcds)-2):
          self.mapped_pcds[i].transform(self.correction_transforms[-1])
          self.poses[i] = np.dot(self.poses[i], self.correction_transforms[-1])
          self.pose_sequence[i] = np.dot(self.pose_sequence[i],self.correction_transforms[-1])

  def run_sequence(self):
    for idx in range(len(self.indices)-1):
      self.estimate_pointcloud_t1t2(idx)
      self.find_rgb_feature_matchest1t2()
      self.align_pointcloudst1t2()

    print("all the feature matching transform sequence ",len(self.transform_sequence), self.transform_sequence)
    print("all the icp correction transforms ",len(self.correction_transforms), self.correction_transforms)

  def global_refinement(self):
    self.refined_pcds, self.refined_poses = pose_graph_optimizer.run(self.mapped_pcds[:-1])
    for i in range(len(self.poses)-1):
      self.pose_sequence[i] = np.dot(self.pose_sequence[i], self.refined_poses[i])

  def merge_map_poses(self, display_result = 'colab'):
    #===================================================================
    #(check the accuracy of poses by reconstructing the scene from untransformed pointclouds stored earlier in rgbds list)
    #(using only the pose list)
    print("Now trying to get back scene from rgbd and poses ...")
    opcds=[]
    #poses.insert(0, np.array([[1.0,0.0,0.0,0.0], [0.0,1.0,0.0,0.0], [0.0,0.0,1.0,0.0], [0.0,0.0,0.0,1.0] ] ))
    for i in range(len(self.poses)-1):
      pcd, _ = project_rgbd.project_rgbd_to_pointcloud(self.rgbds[i][0][:,:,::-1],self.rgbds[i][1], self.iphone_back_camera_params)
      #pcd = pcd.transform(self.poses[i])
      pcd = pcd.transform(self.pose_sequence[i])
      #pcd = pcd.voxel_down_sample(voxel_size=self.voxel_size)
      opcds.append(pcd)
    #===================================================================
    self.scene_map = join_pcds(opcds)
    self.scene_map = self.scene_map.voxel_down_sample(voxel_size=self.voxel_size)

    #===================================================================
    #(visualize change in poses as coord meshes)
    self.coord_meshes = []
    #poses = poses[1:]
    for i in range(len(self.poses)-1):
      #self.pose_sequence[i] = np.dot(self.pose_sequence[i], self.refined_poses[i])

      coord_mesh = o3d.geometry.TriangleMesh.create_coordinate_frame()
      coord_mesh.scale(0.1, center=coord_mesh.get_center())
      coord_mesh.transform(self.pose_sequence[i])
      #coord_mesh.transform(refined_poses[i])
      self.coord_meshes.append(coord_mesh)

    #self.scene_map.extend(self.coord_meshes)
    self.scene_with_coords = [self.scene_map]
    self.scene_with_coords.extend(self.coord_meshes)
    if display_result=='colab':
      project_rgbd.show_pcd_colab(self.scene_with_coords)
    elif display_result=='dont_display':
      print("finished mapping, check results by inspecting mapper.scene_with_coords")
    else:
      project_rgbd.show_pcd(self.scene_with_coords)

    #project_rgbd.show_pcd([self.scene_map, self.coord_meshes])
    print("Total number of points in the scene ",np.array(self.scene_map.points).shape[0])
    #===================================================================

if __name__ == '__main__':
  '''
  f = [240, 250, 260, 270, 280, 290, 300, 310, 320, 330 ,
       340, 350, 360, 370, 380, 390, 400, 410, 420, 430 ,
       440, 450, 460, 470, 480, 490, 500, 510, 520, 530 ]
  '''
  f = [340, 341, 342, 343, 344, 345, 346, 347, 348, 349, 350, 351, 352, 353, 354, 355, 356, 357, 358, 359, 360]

  #f = [340, 350, 360, 370, 380, 390, 400, 410, 420, 430] #map starts to lose details and gets highly overlapped on increasing number of frames
  #initial unit pose
  i_pose = np.array([[1.0,0.0,0.0,0.0], [0.0,1.0,0.0,0.0], [0.0,0.0,1.0,0.0], [0.0,0.0,0.0,1.0] ] )

  m = mapper(frame_numbers = f, init_pose = i_pose)
  m.voxel_size = 0.01
  m.run_sequence()
  m.global_refinement()
  m.merge_map_poses()

  ### to do
  #==============
  '''
  0. decide a global maximum depth value to normalize the depth map and not the max of the first depth map
  1. instead of final voxel downsample maintain an iterative map by downsampling iteratively and merging new points iteratively
  2. superglue matching, replace P3P alignment with fundamental matrix based alignment (dont use depth estimation points) - additional option to use the code for only VO and ignore mapping
  3. pair wise microbundle adjustment from the github code (possibly correct the mono depth estimation points as well) - https://github.com/Parskatt/micro-bundle-adjustment
  4. introduce tsdf fusion to get mesh - take code from Desktop/slam/rgbd_slam/kinfu_gui.py line 44
  '''

  '''
  i_pose = np.array([[1.0,0.0,0.0,0.0], [0.0,1.0,0.0,0.0], [0.0,0.0,1.0,0.0], [0.0,0.0,0.0,1.0] ] )

  local_maps = []
  sparse_poses = []
  all_coord_meshes = []

  frame_tuples = [[240, 250, 260, 270, 280, 290, 300, 310, 320, 330 ,
                   340, 350, 360, 370, 380, 390, 400, 410, 420, 430 ,
                   440, 450, 460, 470, 480, 490, 500, 510, 520, 530 ]]

  for f in frame_tuples:
    m = mapper(frame_numbers = f, init_pose = i_pose)
    m.run_sequence()
    #m.global_refinement()
    m.merge_map_poses()

    local_maps.append(m.scene_map)
    all_coord_meshes.extend(m.coord_meshes)
    sparse_poses.append(m.pose_sequence[-1])
    i_pose = np.dot(i_pose, m.pose_sequence[-1])

  print("showing final global map ")
  total_map = join_pcds(local_maps)
  total_map = total_map.voxel_down_sample(voxel_size=0.01)

  total_map_coords = [total_map]
  total_map_coords.extend(all_coord_meshes)
  project_rgbd.show_pcd(total_map_coords)

  print("Total number of points in the global scene ",np.array(total_map.points).shape[0])
  '''


import wave

import numpy as np

class file_handle:
    def readfile(self, filename):
        f = wave.open(filename, "rb")
        # 只读，返回wave_read
        params = f.getparams()
        # 返回一个tuple:
        # (声道数nchannels,量化位数sampwidth(byte),采样频率framerate,采样点数nframes,...)
        print(params)
        str_data = f.readframes(params[3])
        # readframes返回的是二进制数据（一大堆bytes，且不包含文件头)
        f.close()
        wave_data_origin = np.fromstring(str_data, dtype=np.short)
        # 根据声道数和量化单位，将读取的二进制数据转换为一个可以计算的数组
        # 由于我们的声音格式是以两个字节表示一个取样值，因此采用short数据类型转换。
        return params, wave_data_origin

    def writefile(self, filename, data):
        fw = open(filename, 'w')
        for i in range(len(data)):
            m = str(data[i]) + "\n"
            fw.writelines(m)
        fw.close()

    def write_vad(self, filename, wave_vad_final, framerate):
        fw = wave.open(filename, 'wb')
        fw.setnchannels(1)  # 双声道
        fw.setsampwidth(2)  # 16位
        fw.setframerate(framerate)  # 取样频率
        fw.writeframes(wave_vad_final.tostring())  # 转换为二进制数据写入文件
        fw.close()


'''
Created on 16-Jan-2020

@author: Srinivas Gunti
'''
from main.model import State

FLAG_VECTOR = [4, 2, 16, 2, 16, 20, 14, 20, 14, 10, 10, 10]

WIN_TEXT = 'You Win!!'
LOSS_TEXT = 'Game Over'

class MinesweeperView(object):
    CS = 20

    def __init__(self, app):
        self.app = app
        self.model = app.model

        W = self.model.x
        H = self.model.y

        self.header = app.header_canvas
        self.header.config(width=W * self.CS, height=1 * self.CS, bg='#33B5E5', highlightthickness=0)

        self.flags = self.header.create_text(self.CS//4, self.CS//2, anchor='w')
        self.timer = self.header.create_text(W*self.CS - self.CS//4, self.CS//2, anchor='e')
        self.canvas = app.canvas
        self.canvas.config(width=W * self.CS, height=H * self.CS, highlightthickness=0)
        self.cells = [{'x':i % W, 'y':i // W} for i in range(W * H)]
        for i, cell in enumerate(self.cells):
            X = cell['x'] * self.CS
            Y = cell['y'] * self.CS
            clue_str = str(self.model.clues[i]).replace("0", " ")

            cell['frame'] = self.canvas.create_rectangle(X, Y, X + self.CS, Y + self.CS)
            cell['clue'] = self.canvas.create_text(X + self.CS // 2, Y + self.CS // 2, text=clue_str)
            cell['cover'] = self.canvas.create_rectangle(X, Y, X + self.CS, Y + self.CS, fill='grey')
            cell['flag'] = None
            cell['bomb'] = None

        self.update()

    def update_clues(self):
        for i, cell in enumerate(self.cells):
            self.canvas.itemconfig(cell['clue'], text=str(self.model.clues[i]).replace("0", " "))

    def update(self):
        self.header.itemconfig(self.flags, text=f'Flags: {self.model.flags}')
        for i, cell in enumerate(self.cells):
            X = cell['x'] * self.CS
            Y = cell['y'] * self.CS

            if self.model.gridstate[i] == State.KNOWN:
                if cell['cover']:
                    self.canvas.delete(cell['cover'])
                    cell['cover'] = None

            if self.model.gridstate[i] == State.FLAGGED:
                if not cell['flag']:
                    cell['flag'] = self.canvas.create_polygon(*map(lambda x: x*self.CS/20, FLAG_VECTOR), fill='red')
                    self.canvas.move(cell['flag'], X, Y)
            elif cell['flag']:
                self.canvas.delete(cell['flag'])
                cell['flag'] = None

            if self.model.game_over and self.model.clues[i] == self.model.BOMB:
                if self.model.gridstate[i] == State.EXPLODED:
                    cell['bomb'] = self.canvas.create_oval(3, 3, self.CS - 3, self.CS - 3, fill='black')
                    self.canvas.move(cell['bomb'], X, Y)
                if self.model.gridstate[i] == State.UNKNOWN:
                    cell['bomb'] = self.canvas.create_oval(3, 3, self.CS - 3, self.CS - 3, fill='black')
                    self.canvas.move(cell['bomb'], X, Y)

        if self.model.game_over:
            if self.model.win:
                self.draw_ending_text(WIN_TEXT)
            else:
                self.draw_ending_text(LOSS_TEXT)

    def draw_ending_text(self, ending_text):
        X = self.model.x * self.CS // 2
        Y = self.model.y * self.CS // 2
        self.canvas.create_text(X + 1, Y - 1, fill="black", font=("consolas", 22), text=ending_text)
        self.canvas.create_text(X - 1, Y + 1, fill="black", font=("consolas", 22), text=ending_text)
        self.canvas.create_text(X - 1, Y - 1, fill="black", font=("consolas", 22), text=ending_text)
        self.canvas.create_text(X + 1, Y + 1, fill="black", font=("consolas", 22), text=ending_text)
        self.canvas.create_text(X, Y, fill="white", font=("consolas", 22), text=ending_text)

    def update_timer(self):
        self.header.itemconfig(self.timer, text=f'time: {int(self.app.timer.get_time()):03}')


"""
2. 尾部的零
设计一个算法，计算出n阶乘中尾部零的个数

您在真实的面试中是否遇到过这个题？
样例
样例  1:
	输入: 11
	输出: 2

	样例解释:
	11! = 39916800, 结尾的0有2个。

样例 2:
	输入:  5
	输出: 1

	样例解释:
	5! = 120， 结尾的0有1个。

挑战 O(logN)的时间复杂度

链接: https://www.lintcode.com/problem/trailing-zeros/description
"""

class Solution:
    """
    @param: n: An integer
    @return: An integer, denote the number of trailing zeros in n!
    """

    def trailingZeros(self, n):
        count = 0
        temp = n // 5
        while temp != 0:
            count += temp
            temp //= 5
        return count

s = Solution()
print(s.trailingZeros(12345678987654321))

数学/2. 尾部的零.py


from json import load

from matplotlib.colors import Normalize
import matplotlib
from plot_resonance import plot_resonanceDiagram_color
import matplotlib.pyplot as plt
import numpy as np
import os
import glob
import re

from plot_general import myhexbin

class Tune:
    """
    Load and plot tune data
    """

    def __init__(self,
                 home,
                 yearMonDay,
                 hourMinSec,
                 particle,
                 bunchidList,
                 nux,
                 nuy,
                 tuneshift_direction,
                 ncpu,
                 xlim=[0, 1],
                 ylim=[0, 1]):
        self.home = home
        self.yearMonDay = yearMonDay
        self.hourMinSec = hourMinSec
        self.particle = particle
        self.bunchid = bunchidList
        self.nux = nux
        self.nuy = nuy
        self.tuneshift_direction = tuneshift_direction
        self.ncpu = ncpu
        self.xlim = xlim
        self.ylim = ylim

        self.filePath = []
        self.savePath_scatter = []
        self.savePath_hexbin = []

        for bunchid in self.bunchid:
            filePath = self.hourMinSec + '_' + self.particle + '_bunch' + str(
                bunchid) + '_phase_'
            if self.yearMonDay > '2021_0913':
                filePath = os.sep.join([
                    self.home, 'phase', self.yearMonDay, self.hourMinSec,
                    filePath
                ])
            else:
                filePath = os.sep.join([
                    self.home, 'statLumiPara', self.yearMonDay,
                    self.hourMinSec, filePath
                ])

            filePath = glob.glob(filePath + '*')
            filePath.sort()
            # print(filePath)
            tune_turn = []

            for i in range(len(filePath)):
                if os.path.exists(filePath[i]):
                    matchObj = re.match(r'(.*)phase_(.*)turn(.*).csv',
                                        filePath[i])
                    tune_turn.append(matchObj.group(2))
                    savePath = os.sep.join([
                        self.home, 'statLumiPara', self.yearMonDay,
                        self.hourMinSec, 'figure_tuneSpread'
                    ])
                    if not os.path.exists(savePath):
                        os.makedirs(savePath)
                    savePath = os.sep.join([
                        savePath, self.hourMinSec + '_phase_' + self.particle +
                        "_bunch" + str(bunchid) + '_' + tune_turn[i]
                    ])

                    self.filePath.append(filePath[i])
                    self.savePath_scatter.append(savePath + '_scatter')
                    self.savePath_hexbin.append(savePath + '_hexbin')
        # for i in range(len(self.filePath)):
        #     print(self.filePath[i])
        #     print(self.savePath_scatter[i])
        self.nfile = len(self.filePath)
        self.ntask = self.nfile if self.nfile < self.ncpu else self.ncpu
        print('{0:d} files will be drawn'.format(self.nfile))

    def get_phase_limit(self):
        if self.xlim == [0, 1] or self.ylim == [0, 1]:
            if len(self.filePath) > 0:

                print('Cal limit by file: ', self.filePath[0])
                nuX, nuY, tag = np.loadtxt(self.filePath[0],
                                           delimiter=',',
                                           skiprows=1,
                                           usecols=(2, 3, 4),
                                           unpack=True)
                delete_list = []  # 保存需要删除的元素下标，最后统一删除
                for j in range(len(tag)):
                    if tag[j] <= 0:
                        delete_list.append(j)
                    elif np.isnan(nuX[j]) or np.isnan(nuY[j]) or np.isinf(
                            nuX[j]) or np.isinf(nuY[j]):
                        delete_list.append(j)

                nuX = np.delete(nuX, delete_list)
                nuY = np.delete(nuY, delete_list)

                xmin = min(nuX)
                ymin = min(nuY)
                xmax = max(nuX)
                ymax = max(nuY)

                if xmax == xmin:
                    xmin = xmin - (xmin - 0) * 0.2
                    xmax = xmax + (1 - xmax) * 0.2
                if ymax == ymin:
                    ymin = ymin - (ymin - 0) * 0.2
                    ymax = ymax + (1 - ymax) * 0.2

                if self.tuneshift_direction > 0:
                    xgap = abs(xmax - self.nux)
                    ygap = abs(ymax - self.nuy)
                    axmin = self.nux - xgap * 0.3
                    aymin = self.nuy - ygap * 0.3
                    axmax = xmax + xgap * 0.7
                    aymax = ymax + ygap * 0.7
                else:
                    xgap = abs(xmin - self.nux)
                    ygap = abs(ymin - self.nuy)
                    axmin = xmin - xgap * 0.3
                    aymin = ymin - ygap * 0.3
                    axmax = self.nux + xgap * 0.7
                    aymax = self.nuy + ygap * 0.7

                axmin = 0 if axmin < 0 else axmin
                aymin = 0 if aymin < 0 else aymin
                axmax = 1 if axmax > 1 else axmax
                aymax = 1 if aymax > 1 else aymax

                if self.xlim == [0, 1]:
                    self.xlim = [axmin, axmax]
                if self.ylim == [0, 1]:
                    self.ylim = [aymin, aymax]
            else:
                if self.xlim == [0, 1]:
                    if self.nux > 0.5:
                        self.xlim[0] = 0.5
                    else:
                        self.xlim[1] = 0.5
                if self.ylim == [0, 1]:
                    if self.nuy > 0.5:
                        self.ylim[0] = 0.5
                    else:
                        self.ylim[1] = 0.5

        print('xlim: ', self.xlim)
        print('ylim: ', self.ylim)

    def allocate_phase_file(self):
        self.fileIndex = []
        for i in range(self.ncpu):
            self.fileIndex.append([])
        for i in range(self.nfile):
            self.fileIndex[i % self.ncpu].append(i)
        # print(self.fileIndex)

def load_phase(filePath):
    nuX, nuY, tag = np.loadtxt(filePath,
                               delimiter=',',
                               skiprows=1,
                               usecols=(2, 3, 4),
                               unpack=True)

    delete_list = []  # 保存需要删除的元素下标，最后统一删除
    for j in range(len(tag)):
        if tag[j] <= 0:
            delete_list.append(j)
        elif np.isnan(nuX[j]) or np.isnan(nuY[j]) or np.isinf(
                nuX[j]) or np.isinf(nuY[j]):
            delete_list.append(j)

    nuX = np.delete(nuX, delete_list)
    nuY = np.delete(nuY, delete_list)

    return nuX, nuY

def plot_phase_scatter(tune,
                       fig,
                       ax,
                       nuX,
                       nuY,
                       resonanceOrder,
                       myalpha,
                       mysize,
                       resonanceKind='all',
                       myfontsize=10):
    sc = ax.scatter(nuX,
                    nuY,
                    alpha=myalpha,
                    s=mysize,
                    color='red',
                    zorder=resonanceOrder)
    # cbar = fig.colorbar(sc)
    # cbar.ax.tick_params(labelsize=myfontsize)
    plot_resonanceDiagram_color(resonanceOrder, ax, tune.xlim, tune.ylim,
                                resonanceKind, myfontsize)

def plot_phase_hexbin(tune,
                      fig,
                      ax,
                      nuX,
                      nuY,
                      resonanceOrder,
                      myalpha,
                      mygridsize,
                      resonanceKind='all',
                      myfontsize=10):
    # he = ax.hexbin(self.nuX[phaseTime],
    #                self.nuY[phaseTime],
    #                gridsize=mysize,
    #                norm=matplotlib.colors.LogNorm(),
    #                cmap='coolwarm')
    he = myhexbin(ax,
                  nuX,
                  nuY,
                  gridsize=mygridsize,
                  alpha=myalpha,
                  scattersize=1,
                  norm=matplotlib.colors.LogNorm(),
                  cmap='jet')
    cbar = fig.colorbar(he)
    cbar.ax.tick_params(labelsize=myfontsize)

    plot_resonanceDiagram_color(resonanceOrder, ax, tune.xlim, tune.ylim,
                                resonanceKind, myfontsize)

def save_phase_scatter(figure, savepath, mydpi=300):
    # print(savepath)
    figure.savefig(savepath, dpi=mydpi)

def save_phase_hexbin(figure, savepath, mydpi=300):
    # print(savepath)
    figure.savefig(savepath, dpi=mydpi)


import time
from collections import defaultdict
import iso8601
import logging
import arrow

from impactstoryanalytics.widgets.widget import Widget, get_raw_dataclip_data
from impactstoryanalytics.widgets.widget_api_helpers import Converter

logger = logging.getLogger("impactstoryanalytics.widgets.signup_growth_target")

class Signup_growth_target(Widget):
    signups_with_products_url = "https://dataclips.heroku.com/feblvvoknanzuiumyiawutmqdwbo.json"
    raw_signups_data_url = "https://dataclips.heroku.com/qaclkwteaqnrltwxojgzytkgwczs.json"

    def __init__(self):
        Widget.__init__(self)
        self.start_date = arrow.get("2014-02-03")
        self.daily_growth = .0073  # times 7 = 5.11% weekly
        self.signups_total_endpoints = {
            "start": 4255,
            "finish": 10000  # 10k
        }

    def get_data(self):

        signups_with_products_data = get_raw_dataclip_data(self.signups_with_products_url)
        raw_signups_dataclip_data = get_raw_dataclip_data(self.raw_signups_data_url)

        actual_line = {}

        # make the raw signups line
        for row in raw_signups_dataclip_data["values"]:
            day_str = row[0][0:10]
            actual_line[day_str] = [int(row[1])]

        # make the signups-with-products line
        for row in signups_with_products_data["values"]:
            day_str = row[0][0:10]
            actual_line[day_str].append(int(row[1]))

        target_line = self.make_target_line()

        rows = []
        for date in sorted(target_line.keys()):
            actual_values = actual_line.get(date, ["", ""])
            target_values = target_line[date]

            rows.append([date] + actual_values + target_values)

        fields = [
            "date",
            "signups_raw",
            "signups_with_products",
            "target_signups_with_products",
            "target_signups_with_products_total"
        ]

        return {"fields": fields, "values": rows}

    def get_signup_line(self):
        pass

    def make_target_line(self):
        target_line = {}
        datapoint = {
            "date": self.start_date,
            "signups": "",
            "signups_total": self.signups_total_endpoints["start"]
        }

        while datapoint["signups_total"] < self.signups_total_endpoints["finish"]:

            # update the target line
            date_str = datapoint["date"].format('YYYY-MM-DD')
            target_line[date_str] = [
                datapoint["signups"],
                datapoint["signups_total"]
            ]

            # increment values for next iteration
            datapoint["date"] = datapoint["date"].replace(days=+1)
            datapoint["signups"] = self.daily_growth * datapoint["signups_total"]
            datapoint["signups_total"] *= (1 + self.daily_growth)

        return target_line


def counts(str1, str2):
    f = 0
    for x in str1:
        if str2.count(x) > 0:
            f += 1
    return f

# driven code
s1 = 'abcdef'  # first string
s2 = 'defghia'  # second string

# call count function
print(counts(s1, s2))


import os

import torch
from torch import nn

class ResidualGRU(nn.Module):
    def __init__(self, input_size, output_size, hidden_size, num_layers, dropout_p):
        super(ResidualGRU, self).__init__()
        self.input_size = input_size
        self.output_size = output_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.dropout_p = dropout_p

        self.input_layer = nn.Linear(self.input_size, self.hidden_size)
        self.grus = nn.ModuleList([nn.GRU(2 * hidden_size, hidden_size,
                                          batch_first=True) for _ in range(self.num_layers)])
        self.output_layer = nn.Linear(self.hidden_size, output_size)

    def init_hidden(self, batch_size, device):
        return torch.zeros(1, batch_size, self.hidden_size, device=device)

    def forward(self, x, hidden):
        x = nn.SELU()(self.input_layer(x))
        x_prev = x
        for i in range(len(self.grus)):
            gru_input = torch.cat((x_prev, x), dim=-1)
            x_prev = x
            x, hidden = self.grus[i](gru_input, hidden)
        x = self.output_layer(x)
        return x

    def save(self, archive_dir):
        print("Saving model to {}".format(archive_dir))
        if not os.path.exists(archive_dir):
            os.makedirs(archive_dir)
        torch.save({
            "state_dict": self.state_dict(),
            "input_size": self.input_size,
            "output_size": self.output_size,
            "hidden_size": self.hidden_size,
            "num_layers": self.num_layers,
            "dropout_p": self.dropout_p
        }, os.path.join(archive_dir, "model.pt"))

    @staticmethod
    def load(archive_dir, device):
        print("Loading model from {}".format(archive_dir))
        checkpoint = torch.load(os.path.join(archive_dir, "model.pt"), map_location=device)
        model = ResidualGRU(checkpoint["input_size"], checkpoint["output_size"],
                             checkpoint["hidden_size"], checkpoint["num_layers"],
                             checkpoint["dropout_p"])
        model.load_state_dict(checkpoint["state_dict"])
        model = model.to(device)
        return model


# Calculates the Greatest Common Divisor between the 2 numbers.

def gcd(num1, num2):
    div = 1
    i = 1
    while i <= num1 and i <= num2:
        if (num1 % i) == 0 and (num2 % i) == 0:
            div = i
        i += 1
    return div

def main():
    print(gcd(1660, 1590))
    print(gcd(100, 15))
    print(gcd(25, 150))
    print(gcd(2100, 219))

if __name__ == "__main__":
    main()


'''
语法语义
同一语句块中缩进必须一致
三大结构
1,顺序
2.分支，
3.循环
'''

#循环： while :一般不能确定长度的循环

i = 0
while i < 10:
    print('打印:',i)
    i += 1

#如：存钱 存10000 ，年利率为6.7%，第二年连本带息一起存，依此类推，多少年可以翻倍

benqian = 100000
year = 0
while benqian < 200000:
    benqian = benqian * (1 + 0.067)
    year += 1
    print('第{0}年拿了{1}前'.format(year,benqian))

print('第{}年翻倍'.format(year))


#!/usr/bin/python

'''
Update ansibleHosts to register/unregister hosts with ansible

@author : rene.ribaud@hp.com
'''
import sys
import pprint
import argparse
import os
import ConfigParser

# Variable needed

confFile = "/home/sdi/ansible/ansibleHosts"
prog = sys.argv[0]

if __name__ == '__main__':
    pp = pprint.PrettyPrinter(indent=4)

    parser = argparse.ArgumentParser(version = prog + ' version 1.0')

    parser.add_argument('action', action='store',
                        choices=['add','remove'],
                    help='Action to do either add or remove hostname to ansibleHosts')

    parser.add_argument('group', action='store',
                    help='Group to add or remove the hostname')

    parser.add_argument('hostname', action='store',
                    help='hostname to add or remove')

    arguments = parser.parse_args()

    # Load json

    config = ConfigParser.RawConfigParser(allow_no_value=True)
    config.read(confFile)

    if (arguments.action == "add"):
        config.add_section(arguments.group)
        config.set(arguments.group, arguments.hostname)

    if (arguments.action == "remove"):
        config.remove_option(arguments.group, arguments.hostname)

    with open(confFile, 'wb') as configfile:
        config.write(configfile)


from re import search, finditer, findall
from collections import defaultdict

with open('19.in', 'r') as f:
    raw_input = f.read()

TASK_REGEX = '(\\w+) => (\\w+)'

rules = defaultdict(set)
for line in raw_input.split("\n\n")[0].splitlines():
    begin, end = search(TASK_REGEX, line).group(1, 2)
    rules[begin].add(end)

medicine_molecule = raw_input.split("\n\n")[1]

def part1():
    generated = set()
    for r in rules:
        for m in finditer(r, medicine_molecule):
            for to in rules[r]:
                generated.add(medicine_molecule[:m.start()] + to + medicine_molecule[m.start() + len(r):])
    print(len(generated))

def normalize_chomsky(rules):
    rs = defaultdict(set)
    temp_num = 0
    for r in rules:
        for right_side in rules[r]:
            non_terminals = findall('[A-Z][^A-Z]*', right_side)
            if len(non_terminals) == 2:
                rs[r].add("".join(non_terminals))
            else:
                for i in range(len(non_terminals) - 1):
                    if i == 0:
                        left_side = r
                        new_right_side = non_terminals[0] + "Tmp" + str(temp_num)
                    elif i == len(non_terminals) - 2:
                        left_side = "Tmp" + str(temp_num - 1)
                        new_right_side = non_terminals[-2] + non_terminals[-1]
                    else:
                        left_side = "Tmp" + str(temp_num - 1)
                        new_right_side = non_terminals[i] + "Tmp" + str(temp_num)
                    rs[left_side].add(new_right_side)
                    temp_num += 1
    return rs

def search_for_rules(right_side):
    found = set()
    for l, r in rules.items():
        if right_side in r:
            found.add(l)
    return found

def part2():
    global rules
    rules = normalize_chomsky(rules)

    parts = findall('[A-Z][^A-Z]*', medicine_molecule)

    Xs = defaultdict(set)
    for i in range(len(parts)):
        Xs[(i+1, i+1)].add(parts[i])

    Ys = defaultdict(set)
    for i in range(len(parts)):
        Ys[(i+1, i+1)].add(0)

    for subword_length in range(2, len(parts) + 1):
        for start in range(1, len(parts) + 1 - subword_length + 1):
            currently_calculated = (start, start + subword_length - 1)
            for left_side in range(subword_length - 1):
                pyramid_left_element = (start, start + left_side)
                pyramid_right_element = (start + left_side + 1, start + subword_length - 1)
                for a in Xs[pyramid_left_element]:
                    for b in Xs[pyramid_right_element]:
                        for r in search_for_rules(a + b):
                            Xs[currently_calculated].add(r)
                            Ys[currently_calculated].add(min(Ys[pyramid_left_element]) + min(Ys[pyramid_right_element]) + [0, 1]['Tmp' not in r])

    print(min(Ys[(1, len(parts))]))

part1()
part2()


# 가장 긴 증가하는 부분 수열2, 12015번
import sys
from bisect import bisect_left

n = int(sys.stdin.readline().rstrip())
sequence = list(map(int,sys.stdin.readline().rstrip().split()))
def lis(sequence):
    lis = []
    for seq in sequence:
        if len(lis) == 0 or lis[-1] < seq:
            lis.append(seq)
        else:
            index = lowerBound(lis, seq)
            lis[index] = seq
    return len(lis)
def lowerBound(array,target):
    low = 0
    high = len(array)-1
    index = -1
    while low<=high:
        mid = (low+high)//2
        if array[mid]<target:
            low=mid+1
        else:
            high=mid-1
            index = mid
    return index

def lisUsingBisect(sequence):
    lis = []
    for seq in sequence:
        if len(lis) == 0 or lis[-1] < seq:
            lis.append(seq)
        else:
            index = bisect_left(lis,seq)
            lis[index] = seq
    return len(lis)

print(lis(sequence))
print("bisect 이용",lisUsingBisect(sequence))


"""
=============
Demo Axes RGB
=============

RGBAxes to show RGB composite images.
"""

import numpy as np
from matplotlib import cbook
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1.axes_rgb import make_rgb_axes, RGBAxes

def get_rgb():
    Z = cbook.get_sample_data("axes_grid/bivariate_normal.npy", np_load=True)
    Z[Z < 0] = 0.
    Z = Z / Z.max()

    R = Z[:13, :13]
    G = Z[2:, 2:]
    B = Z[:13, 2:]

    return R, G, B

def make_cube(r, g, b):
    ny, nx = r.shape
    R = np.zeros((ny, nx, 3))
    R[:, :, 0] = r
    G = np.zeros_like(R)
    G[:, :, 1] = g
    B = np.zeros_like(R)
    B[:, :, 2] = b

    RGB = R + G + B

    return R, G, B, RGB

def demo_rgb1():
    fig = plt.figure()
    ax = RGBAxes(fig, [0.1, 0.1, 0.8, 0.8], pad=0.0)
    r, g, b = get_rgb()
    ax.imshow_rgb(r, g, b)

def demo_rgb2():
    fig, ax = plt.subplots()
    ax_r, ax_g, ax_b = make_rgb_axes(ax, pad=0.02)

    r, g, b = get_rgb()
    im_r, im_g, im_b, im_rgb = make_cube(r, g, b)
    ax.imshow(im_rgb)
    ax_r.imshow(im_r)
    ax_g.imshow(im_g)
    ax_b.imshow(im_b)

    for ax in fig.axes:
        ax.tick_params(axis='both', direction='in')
        for sp1 in ax.spines.values():
            sp1.set_color("w")
        for tick in ax.xaxis.get_major_ticks() + ax.yaxis.get_major_ticks():
            tick.tick1line.set_markeredgecolor("w")
            tick.tick2line.set_markeredgecolor("w")

demo_rgb1()
demo_rgb2()

plt.show()


import torch
import torch.nn as nn
import torch.nn.functional as F

from ..layers.neural_process import fc_layer

class NeuralProcess(nn.Module):

    def __init__(self, hidden_dim=256):
        super(NeuralProcess, self).__init__()
        self.hidden_dim = 30
        self.h = nn.Sequential(
                fc_layer(3, 400),
                fc_layer(400, 400),
                fc_layer(400, 400),
                )

        self.mean_aggregater = nn.Sequential(
                nn.Linear(400, 400),
                nn.LeakyReLU(),
                nn.Linear(400, self.hidden_dim),
                )
        self.variance_aggregater = nn.Sequential(
                nn.Linear(400, 400),
                nn.LeakyReLU(),
                nn.Linear(400, self.hidden_dim),
                )

        self.g = nn.Sequential(
                fc_layer(2 + self.hidden_dim, 400),
                fc_layer(400, 400),
                fc_layer(400, 400),
                fc_layer(400, 400),
                nn.Linear(400, 1)
                )

    def forward(self, context_x, context_y, target_x, target_y):
        context = torch.cat((context_x, context_y), dim=1)
        context_representations = self.h(context).mean(0)

        context_mean = self.mean_aggregater(context_representations)
        context_std = self.variance_aggregater(context_representations)
        context_var = torch.exp(0.5 * context_std)
        context_z = torch.rand_like(context_var) * context_var + context_mean

        all_x = torch.cat((context_x, target_x), dim=0)
        all_y = torch.cat((context_y, target_y), dim=0)

        all_representations = self.h(torch.cat((all_x, all_y), dim=1)).mean(0)
        all_mean = self.mean_aggregater(all_representations)
        all_std = self.variance_aggregater(all_representations)
        all_var = torch.exp(0.5 * all_std)
        all_z = torch.rand_like(all_var) * all_var + all_mean

        x_domain = all_x.size(0)
        target_x_domain = target_x.size(0)
        context_x_domain = context_x.size(0)

        expanded_all_z = torch.stack([all_z for i in range(target_x_domain)])
        target_logits = F.sigmoid(self.g(
            torch.cat((expanded_all_z, target_x), dim=1)))

        expanded_all_z = torch.stack([all_z for i in range(context_x_domain)])
        context_logits = F.sigmoid(self.g(
            torch.cat((expanded_all_z, context_x), dim=1)))

        expanded_context_z = torch.stack([context_z for i in range(context_x_domain)])
        context_preds_logits = F.sigmoid(self.g(
            torch.cat((expanded_context_z, context_x), dim=1)))

        expanded_target_z = torch.stack([context_z for i in range(target_x_domain)])
        target_preds_logits = F.sigmoid(self.g(
            torch.cat((expanded_target_z, target_x), dim=1)))

        negative_kl_divergence = self._calc_negative_kl_divergence(
                context_mean, context_std, all_mean, all_std).sum()

        return target_logits, context_logits, context_preds_logits, \
            target_preds_logits, negative_kl_divergence

    def _calc_negative_kl_divergence(self, c_mu, c_std, a_mu, a_std):
        c_std = torch.exp(0.5 * c_std)
        a_std = torch.exp(0.5 * a_std)
        return  torch.log(c_std / a_std) + \
            ((a_std ** 2 + (a_mu - c_mu) ** 2) / (2 * c_std ** 2)) - 0.5


import numpy as np
import contextlib
import torch
from collections import deque
import matplotlib.pyplot as plt

# Configures numpy print options
@contextlib.contextmanager
def _printoptions(*args, **kwargs):
    original = np.get_printoptions()
    np.set_printoptions(*args, **kwargs)
    try:
        yield
    finally:
        np.set_printoptions(**original)

class EnvironmentModel:
    def __init__(self, n_states, n_actions, seed=None):
        self.n_states = n_states
        self.n_actions = n_actions

        self.random_state = np.random.RandomState(seed)

    def p(self, next_state, state, action):
        raise NotImplementedError()

    def r(self, next_state, state, action):
        raise NotImplementedError()

    def draw(self, state, action):
        p = [self.p(ns, state, action) for ns in range(self.n_states)]
        next_state = self.random_state.choice(self.n_states, p=p)
        reward = self.r(next_state, state, action)

        return next_state, reward

class Environment(EnvironmentModel):
    def __init__(self, n_states, n_actions, max_steps, pi, seed=None):
        EnvironmentModel.__init__(self, n_states, n_actions, seed)

        self.max_steps = max_steps

        self.pi = pi
        if self.pi is None:
            self.pi = np.full(n_states, 1./n_states)

    def reset(self):
        self.n_steps = 0
        self.state = self.random_state.choice(self.n_states, p=self.pi)

        return self.state

    def step(self, action):
        if action < 0 or action >= self.n_actions:
            raise Exception('Invalid action.')

        self.n_steps += 1
        done = (self.n_steps >= self.max_steps)

        self.state, reward = self.draw(self.state, action)

        return self.state, reward, done

    def render(self, policy=None, value=None):
        raise NotImplementedError()

class FrozenLake(Environment):
    def __init__(self, lake, slip, max_steps, seed=None):
        """
        lake: A matrix that represents the lake. For example:
        lake =  [['&', '.', '.', '.'],
                ['.', '#', '.', '#'],
                ['.', '.', '.', '#'],
                ['#', '.', '.', '$']]
        slip: The probability that the agent will slip
        max_steps: The maximum number of time steps in an episode
        seed: A seed to control the random number generator (optional)
        """
        # start (&), frozen (.), hole (#), goal ($)
        self.lake = np.array(lake)
        self.lake_flat = self.lake.reshape(-1)

        self.size = len(self.lake_flat)
        self.width = int(np.sqrt(self.size))

        self.holeStates = []
        for i in range(len(self.lake_flat)):
            if self.lake_flat[i] == '#': self.holeStates.append(i)

        self.slip = slip

        n_states = self.lake.size + 1
        n_actions = 4

        pi = np.zeros(n_states, dtype=float)
        pi[np.where(self.lake_flat == '&')[0]] = 1.0

        self.absorbing_state = n_states - 1

        Environment.__init__(self, n_states, n_actions, max_steps, pi, seed=seed)

    actions = {'w':0, 'a':1, 's':2, 'd':3}

    def step(self, action):
        state, reward, done = Environment.step(self, action)

        done = (state == self.absorbing_state) or done

        return state, reward, done

    def isAbsorbingState(self, state): return state == self.size
    def isStartState(self, state): return state == np.where(self.lake_flat == '&')[0][0]
    def isGoalState(self, state): return state == np.where(self.lake_flat == '$')[0][0]
    def isHoleState(self, state): return state in self.holeStates

    def p(self, next_state, state, action):
        if self.isAbsorbingState(state) or self.isGoalState(state) or self.isHoleState(state):
            return self.isAbsorbingState(next_state)
        elif self.isAbsorbingState(next_state): return 0

        # Describes [y, x] movement of the action
        actions = [[-1, 0], [0, -1], [1, 0], [0, 1]]
        a = actions[action]
        w = self.width
        ps = np.ndarray((w, w))
        ps.fill(0)
        getCoords = lambda a: (np.array([state//w, state%w]) + np.array(a)).clip(0, self.width - 1)
        ps[getCoords(a)[0], getCoords(a)[1]] = 1 - self.slip
        for _a in actions: ps[getCoords(_a)[0], getCoords(_a)[1]] += self.slip / self.n_actions
        return ps.flatten()[next_state]

    def r(self, _, state, __):
        return 1 if self.isGoalState(state) else 0

    def render(self, policy=None, value=None):
        if policy is None:
            lake = np.array(self.lake_flat)

            if self.state < self.absorbing_state:
                lake[self.state] = '@'

            print(lake.reshape(self.lake.shape))
        else:
            # UTF-8 arrows look nicer, but cannot be used in LaTeX
            # https://www.w3schools.com/charsets/ref_utf_arrows.asp
            actions = ['^', '<', '_', '>']

            print('Lake:')
            print(self.lake)

            print('Policy:')
            policy = np.array([actions[a] for a in policy[:-1]])
            print(policy.reshape(self.lake.shape))

            print('Value:')
            with _printoptions(precision=3, suppress=True):
                print(value[:-1].reshape(self.lake.shape))

def play(env):
    actions = ['w', 'a', 's', 'd']

    state = env.reset()
    env.render()

    done = False
    while not done:
        c = input('\nMove: ')
        if c not in actions:
            raise Exception('Invalid action')

        state, r, done = env.step(actions.index(c))

        env.render()
        print('Reward: {0}.'.format(r))

def policy_evaluation(env, policy, gamma, theta, max_iterations):
    value = np.zeros(env.n_states, dtype=float)
    inner = lambda s, a: sum([env.p(s_, s, a) * (env.r(s_, s, a) + gamma * value[s_]) for s_ in range(env.n_states)])
    outer = lambda s: inner(s, policy[s])
    for _ in range(max_iterations):
        v = value.copy();
        value = np.array([outer(s) for s in range(env.n_states)])
        if max([abs(v[i] - value[i]) for i in range(len(value))]) < theta: break
    print(f'Policy Evaluation: Number of Iterations: {_+1}')
    return value

def policy_improvement(env, value, gamma):
    inner = lambda s, a: sum([env.p(s_, s, a) * (env.r(s_, s, a) + gamma * value[s_]) for s_ in range(env.n_states)])
    outer = lambda s: np.argmax([inner(s, a) for a in range(env.n_actions)])
    return np.array([outer(s) for s in range(env.n_states)])

def policy_iteration(env, gamma, theta, max_iterations, policy=None):
    if policy is None:
        policy = np.zeros(env.n_states, dtype=int)
    else:
        policy = np.array(policy, dtype=int)
    value = policy_evaluation(env, policy, gamma, theta, max_iterations)

    old_policy = None
    i = 0
    while not np.array_equal(policy, old_policy):
        old_policy = policy.copy()
        policy = policy_improvement(env, value, gamma)
        value = policy_evaluation(env, policy, gamma, theta, max_iterations)
        i += 1
    print(f'Policy Iteration: Number of Iterations: {i}')

    return policy, value

def value_iteration(env, gamma, theta, max_iterations, value=None):
    value = np.zeros(env.n_states) if value is None else np.array(value, dtype=np.float)
    # Define math function as outer
    inner = lambda s, a: sum([env.p(s_, s, a) * (env.r(s_, s, a) + gamma * value[s_]) for s_ in range(env.n_states)])
    outer = lambda s: np.max([inner(s, a) for a in range(env.n_actions)])
    # Iterate
    for _ in range(max_iterations):
        _value = value.copy()
        value = np.array([outer(s) for s in range(env.n_states)])
        if max([abs(_value[i] - value[i]) for i in range(len(value))]) < theta: break
    print(f'Value Iteration: Number of Iterations: {_+1}')
    inner = lambda s, a: sum([env.p(s_, s, a) * value[s_] for s_ in range(env.n_states)])
    outer = lambda s: np.argmax([inner(s, a) for a in range(env.n_actions)])
    policy = np.array([outer(s) for s in range(env.n_states)])
    return policy, value

#region Later

# Returns an argmax from input arr, but randomizes tie breakers
def rand_argmax(arr, random_state = None):
    if random_state is None: random_state = np.random.RandomState()
    max_filtered_arr = [i for i in range(len(arr)) if arr[i] == max(arr)]
    return max_filtered_arr[random_state.randint(0, len(max_filtered_arr))]

# Return action a according to an ε-greedy policy based on Q
# Q should be passed to only include state s, and length of it will be n_actions
def eps_greedy_action(q, epsilon, random_state = None):
    if random_state is None: random_state = np.random.RandomState()
    return rand_argmax(q, random_state) if random_state.uniform() > epsilon else random_state.randint(0, len(q))

def sarsa(env, max_episodes, eta, gamma, epsilon, seed=None):
    random_state = np.random.RandomState(seed)

    eta = np.linspace(eta, 0, max_episodes)
    epsilon = np.linspace(epsilon, 0, max_episodes)

    q = np.zeros((env.n_states, env.n_actions))

    print()
    q_iterations = []
    for i in range(max_episodes):
        s = env.reset()
        a = eps_greedy_action(q[s], epsilon[i], random_state)
        done = False
        while not done:
            s_, r, done = env.step(a)
            a_ = eps_greedy_action(q[s_], epsilon[i], random_state)
            q[s, a] += eta[i] * (r + gamma*q[s_, a_] - q[s,a])
            s, a = s_, a_
        q_iterations += [q.max(axis=1).mean()]
    mp = np.convolve(q_iterations, np.ones(20)/20, mode='valid')
    plt.clf()
    plt.cla()
    plt.plot(np.arange(1, len(mp) + 1), mp)
    plt.savefig('plots/sarsa.png')

    policy = q.argmax(axis=1)
    value = q.max(axis=1)

    return policy, value

def q_learning(env, max_episodes, eta, gamma, epsilon, seed=None):
    random_state = np.random.RandomState(seed)

    eta = np.linspace(eta, 0, max_episodes)
    epsilon = np.linspace(epsilon, 0, max_episodes)

    q = np.zeros((env.n_states, env.n_actions))

    q_iterations = []
    for i in range(max_episodes):
        s = env.reset()
        done = False
        while not done:
            a = eps_greedy_action(q[s], epsilon[i], random_state)
            s_, r, done = env.step(a)
            # Q-Learning formula
            q[s, a] += eta[i] * (r + gamma * q[s_].max() - q[s, a])
            s = s_
        q_iterations += [q.max(axis=1).mean()]
    mp = np.convolve(q_iterations, np.ones(20)/20, mode='valid')
    plt.clf()
    plt.cla()
    plt.plot(np.arange(1, len(mp) + 1), mp)
    plt.savefig('plots/qLearning.png')

    policy = q.argmax(axis=1)
    value = q.max(axis=1)

    return policy, value

class LinearWrapper:
    def __init__(self, env):
        self.env = env

        self.n_actions = self.env.n_actions
        self.n_states = self.env.n_states
        self.n_features = self.n_actions * self.n_states

    def encode_state(self, s):
        features = np.zeros((self.n_actions, self.n_features))
        for a in range(self.n_actions):
            i = np.ravel_multi_index((s, a), (self.n_states, self.n_actions))
            features[a, i] = 1.0

        return features

    def decode_policy(self, theta):
        policy = np.zeros(self.env.n_states, dtype=int)
        value = np.zeros(self.env.n_states)

        for s in range(self.n_states):
            features = self.encode_state(s)
            q = features.dot(theta)

            policy[s] = np.argmax(q)
            value[s] = np.max(q)

        return policy, value

    def reset(self):
        return self.encode_state(self.env.reset())

    def step(self, action):
        state, reward, done = self.env.step(action)

        return self.encode_state(state), reward, done

    def render(self, policy=None, value=None):
        self.env.render(policy, value)

def linear_sarsa(env, max_episodes, eta, gamma, epsilon, seed=None):
    random_state = np.random.RandomState(seed)

    eta = np.linspace(eta, 0, max_episodes)
    epsilon = np.linspace(epsilon, 0, max_episodes)

    theta = np.zeros(env.n_features)

    q_iterations = []
    for i in range(max_episodes):
        features = env.reset()
        q = features.dot(theta)
        a = eps_greedy_action(q, epsilon[i], random_state)
        done = False
        while not done:
            features_, r, done = env.step(a)
            delta = r - q[a]
            q = features_.dot(theta)
            a_ = eps_greedy_action(q, epsilon[i], random_state)
            # Update variables
            delta += gamma * q[a_]
            theta += eta[i] * delta * features[a]
            features, a = features_, a_
        q_iterations += [env.decode_policy(theta)[1].mean()]
    mp = np.convolve(q_iterations, np.ones(20)/20, mode='valid')
    plt.clf()
    plt.cla()
    plt.plot(np.arange(1, len(mp) + 1), mp)
    plt.savefig('plots/linear_sarsa.png')

    return theta

def linear_q_learning(env, max_episodes, eta, gamma, epsilon, seed=None):
    random_state = np.random.RandomState(seed)

    eta = np.linspace(eta, 0, max_episodes)
    epsilon = np.linspace(epsilon, 0, max_episodes)

    theta = np.zeros(env.n_features)

    q_iterations = []
    for i in range(max_episodes):
        features = env.reset()
        q = features.dot(theta)
        done = False
        while not done:
            a = eps_greedy_action(q, epsilon[i], random_state)
            features_, r, done = env.step(a)
            delta = r - q[a]
            # q = np.array([sum(theta * features_[a]) for a in range(env.n_actions)])
            q = features_.dot(theta)
            # Update variables
            delta += gamma * q.max()
            theta += eta[i] * delta * features[a]
            features = features_
        q_iterations += [env.decode_policy(theta)[1].mean()]
    mp = np.convolve(q_iterations, np.ones(20)/20, mode='valid')
    plt.clf()
    plt.cla()
    plt.plot(np.arange(1, len(mp) + 1), mp)
    plt.savefig('plots/linear_qLearning.png')

    return theta

class FrozenLakeImageWrapper:
    def __init__(self, env):
        self.env = env

        lake = self.env.lake

        self.n_actions = self.env.n_actions
        self.state_shape = (4, lake.shape[0], lake.shape[1])

        lake_image = [(lake == c).astype(float) for c in ['&', '#', '$']]

        self.state_image = {env.absorbing_state:
                            np.stack([np.zeros(lake.shape)] + lake_image)}
        for state in range(lake.size):
            x = state // env.width
            y = state % env.width
            channel1 = np.zeros(lake.shape)
            channel1[x, y] = 1.0
            self.state_image[state] = np.stack([channel1] + lake_image)

    def encode_state(self, state):
        return self.state_image[state]

    def decode_policy(self, dqn):
        states = np.array([self.encode_state(s) for s in range(self.env.n_states)])
        q = dqn(states).detach().numpy()  # torch.no_grad omitted to avoid import

        policy = q.argmax(axis=1)
        value = q.max(axis=1)

        return policy, value

    def reset(self):
        return self.encode_state(self.env.reset())

    def step(self, action):
        state, reward, done = self.env.step(action)

        return self.encode_state(state), reward, done

    def render(self, policy=None, value=None):
        self.env.render(policy, value)

class DeepQNetwork(torch.nn.Module):
    def __init__(self, env, learning_rate, kernel_size, conv_out_channels,
                 fc_out_features, seed):
        torch.nn.Module.__init__(self)
        torch.manual_seed(seed)

        self.conv_layer = torch.nn.Conv2d(in_channels=env.state_shape[0],
                                          out_channels=conv_out_channels,
                                          kernel_size=kernel_size, stride=1)

        h = env.state_shape[1] - kernel_size + 1
        w = env.state_shape[2] - kernel_size + 1

        self.fc_layer = torch.nn.Linear(in_features=h * w * conv_out_channels,
                                        out_features=fc_out_features)
        self.output_layer = torch.nn.Linear(in_features=fc_out_features,
                                            out_features=env.n_actions)

        self.optimizer = torch.optim.Adam(self.parameters(), lr=learning_rate)

    def forward(self, x):
        x = torch.tensor(x, dtype=torch.float)
        x = self.conv_layer(x)
        x = x.reshape(x.shape[0],-1)
        x = self.fc_layer(x)
        x = self.output_layer(x)
        return x

    def train_step(self, transitions, gamma, tdqn):
        states = np.array([transition[0] for transition in transitions])
        actions = np.array([transition[1] for transition in transitions])
        rewards = np.array([transition[2] for transition in transitions])
        next_states = np.array([transition[3] for transition in transitions])
        dones = np.array([transition[4] for transition in transitions])

        q = self(states)
        q = q.gather(1, torch.Tensor(actions).view(len(transitions), 1).long())
        q = q.view(len(transitions))

        with torch.no_grad():
            next_q = tdqn(next_states).max(dim=1)[0] * (1 - dones)

        target = torch.Tensor(rewards) + gamma * next_q

        loss = ((q - target) ** 2).mean()

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

class ReplayBuffer:
    def __init__(self, buffer_size, random_state):
        self.buffer = deque(maxlen=buffer_size)
        self.random_state = random_state

    def __len__(self):
        return len(self.buffer)

    def append(self, transition):
        self.buffer.append(transition)

    def draw(self, batch_size):
        index = self.random_state.choice(range(len(self.buffer)), batch_size)
        return np.array(self.buffer)[index]

def deep_q_network_learning(env, max_episodes, learning_rate, gamma, epsilon,
                            batch_size, target_update_frequency, buffer_size,
                            kernel_size, conv_out_channels, fc_out_features, seed):
    random_state = np.random.RandomState(seed)
    replay_buffer = ReplayBuffer(buffer_size, random_state)

    dqn = DeepQNetwork(env, learning_rate, kernel_size, conv_out_channels,
                       fc_out_features, seed=seed)
    tdqn = DeepQNetwork(env, learning_rate, kernel_size, conv_out_channels,
                        fc_out_features, seed=seed)

    epsilon = np.linspace(epsilon, 0, max_episodes)

    q_iterations = []
    for i in range(max_episodes):
        state = env.reset()

        done = False
        while not done:
            if random_state.rand() < epsilon[i]:
                action = random_state.choice(env.n_actions)
            else:
                with torch.no_grad():
                    q = dqn(np.array([state]))[0].numpy()

                qmax = max(q)
                best = [a for a in range(env.n_actions) if np.allclose(qmax, q[a])]
                action = random_state.choice(best)

            next_state, reward, done = env.step(action)

            replay_buffer.append((state, action, reward, next_state, done))

            state = next_state

            if len(replay_buffer) >= batch_size:
                transitions = replay_buffer.draw(batch_size)
                dqn.train_step(transitions, gamma, tdqn)

        if (i % target_update_frequency) == 0:
            tdqn.load_state_dict(dqn.state_dict())

        q_iterations += [env.decode_policy(dqn)[1].mean()]
    mp = np.convolve(q_iterations, np.ones(20)/20, mode='valid')
    plt.clf()
    plt.cla()
    plt.plot(np.arange(1, len(mp) + 1), mp)
    plt.savefig('plots/DeepQNetworkLearning.png')
    return dqn
#endregion

def main():
    seed = 0

    # Big lake
    # lake = [['&', '.', '.', '.', '.', '.', '.', '.'],
    #         ['.', '.', '.', '.', '.', '.', '.', '.'],
    #         ['.', '.', '.', '#', '.', '.', '.', '.'],
    #         ['.', '.', '.', '.', '.', '#', '.', '.'],
    #         ['.', '.', '.', '#', '.', '.', '.', '.'],
    #         ['.', '#', '#', '.', '.', '.', '#', '.'],
    #         ['.', '#', '.', '.', '#', '.', '#', '.'],
    #         ['.', '.', '.', '#', '.', '.', '.', '$']]

    # Small lake
    lake = [['&', '.', '.', '.'],
            ['.', '#', '.', '#'],
            ['.', '.', '.', '#'],
            ['#', '.', '.', '$']]

    max_steps = len([cell for row in lake for cell in row])

    env = FrozenLake(lake, slip=0.1, max_steps=max_steps, seed=seed)
    gamma = 0.9

    print('# Model-based algorithms')

    print('')

    print('## Policy iteration')
    # policy, value = policy_iteration(env, gamma, theta=0.001, max_iterations=128)
    # env.render(policy, value)

    print('')

    print('## Value iteration')
    # policy, value = value_iteration(env, gamma, theta=0.001, max_iterations=128)
    # env.render(policy, value)

    print('')

    print('# Model-free algorithms')
    max_episodes = 4000

    print('')

    print('## Sarsa')
    policy, value = sarsa(env, max_episodes, eta=0.5, gamma=gamma,
                          epsilon=0.5, seed=seed)
    env.render(policy, value)

    print('')

    print('## Q-learning')
    policy, value = q_learning(env, max_episodes, eta=0.5, gamma=gamma,
                               epsilon=0.5, seed=seed)
    env.render(policy, value)

    print('')

    linear_env = LinearWrapper(env)

    print('## Linear Sarsa')

    parameters = linear_sarsa(linear_env, max_episodes, eta=0.5, gamma=gamma,
                              epsilon=0.5, seed=seed)
    policy, value = linear_env.decode_policy(parameters)
    linear_env.render(policy, value)

    print('')

    print('## Linear Q-learning')

    parameters = linear_q_learning(linear_env, max_episodes, eta=0.5, gamma=gamma,
                                   epsilon=0.5, seed=seed)
    policy, value = linear_env.decode_policy(parameters)
    linear_env.render(policy, value)

    print('')

    image_env = FrozenLakeImageWrapper(env)

    print('## Deep Q-network learning')

    dqn = deep_q_network_learning(image_env, max_episodes, learning_rate=0.001,
                                  gamma=gamma,  epsilon=0.2, batch_size=32,
                                  target_update_frequency=4, buffer_size=256,
                                  kernel_size=3, conv_out_channels=4,
                                  fc_out_features=8, seed=4)
    policy, value = image_env.decode_policy(dqn)
    image_env.render(policy, value)

main()