# -*- coding: utf-8 -*-
"""
Support function for net handling
"""

import time

import numpy as np
import scipy.sparse as sp

############## return cor list (raw) with integer values ( = float * 1000)

def return_net_list(Z_cor_mat):

    t2 = time.time()

    print(Z_cor_mat.shape)

    x_sig,y_sig = np.where(Z_cor_mat != 0.0)

    net_list = np.array(np.column_stack((x_sig + 1,y_sig + 1,Z_cor_mat[x_sig,y_sig]*1000)),dtype = int)

    print(net_list.shape)

    t3 = time.time()

    print("Sparse Weighted correlation thresholding computation took " + str(t3-t2) + "s")

    return net_list

def return_int_net_list(int_mat,min_int = 0):

    t2 = time.time()

    x_sig,y_sig = np.where(int_mat > min_int)

    net_list = np.array(np.column_stack((x_sig + 1,y_sig + 1,int_mat[x_sig,y_sig])),dtype = 'int64')

    print(net_list.shape)

    t3 = time.time()

    print("Sparse Weighted correlation thresholding computation took " + str(t3-t2) + "s")

    return net_list

##################################### Formatting data for external community detection algorithm (radatools) ##############################

def export_List_net_from_list(Z_List_file,Z_list):

    print(Z_list.shape)

    #print "saving file " + Z_List_file

    np.savetxt(Z_List_file,Z_list,fmt = '%d %d %d')

#def read_List_net_file(Z_List_file):

    #Z_list = np.loadtxt(Z_List_file,dtype = 'int64')

    #return Z_list

#def read_mod_file(mod_file):

    #community_vect = np.loadtxt(mod_file,dtype = 'int', delimiter = '\t')

    #return community_vect[:,1]

def read_lol_file(lol_file):

    with open(lol_file,'r') as f :

        lines = f.readlines()[4:]

        line_nb_elements = lines[0]

        nb_elements = int(line_nb_elements.split(': ')[1])

        #print nb_elements

        community_vect = np.empty((nb_elements),dtype = int)

        lines = lines[3:]

        #print lines

        for i,line in enumerate(lines):

            try:
                nb_nodes,index_nodes = line.split(': ')
                print (nb_nodes,index_nodes)

                if int(nb_nodes) > 1:
                    index_nodes = np.array(list(map(int,index_nodes.split(' '))),dtype = int) -1

                    #print i,index_nodes
                    community_vect[index_nodes] = i

                else :
                    community_vect[int(index_nodes) -1] = i

            except ValueError:
                print("Warning, reading lol file ")

        f.close()

    return community_vect

def read_Pajek_corres_nodes(Pajek_net_file):

    with open(Pajek_net_file,'r') as f :

        lines = f.readlines()

        line_nb_elements = lines[0]

        nb_elements = int(line_nb_elements.split(' ')[1])

        #print nb_elements

        node_corres = np.empty((nb_elements),dtype = 'int')

        lines = lines[1:(nb_elements+1)]

        #print lines

        for i,line in enumerate(lines):
            #print line

            new_index,old_index = line.split(' ')

            #print i+1, new_index, old_index

            #node_corres[i] = old_index

            if (i+1) == int(new_index) :
                node_corres[i] = int(old_index)-1
            else:
                print("Warning, incompatible indexes {} {}".format(new_index,i+1))

        f.close()

    return node_corres

####### from modified Pajek file, read coords

#def read_Pajek_rel_coords(Pajek_net_file):

    #with open(Pajek_net_file,'r') as f :

        #lines = f.readlines()

        #line_nb_elements = lines[0]

        #nb_elements = int(line_nb_elements.split(' ')[1])

        #print nb_elements

        #node_rel_coords = np.empty((nb_elements,3),dtype = 'float')

        #node_lines = lines[1:(nb_elements+1)]

        ##print lines

        #for i,line in enumerate(node_lines):
            ##print line

            #node_line = line.split(' ')

            ##print node_line

            #node_rel_coords[i,0] = node_line[2]
            #node_rel_coords[i,1] = node_line[3]
            #node_rel_coords[i,2] = node_line[4]

            ##print node_rel_coords[i,:]

        #f.close()

    #return node_rel_coords

    ###########" return corres_nodes and sparse matrix from pajek file

def read_Pajek_corres_nodes_and_sparse_matrix(Pajek_net_file):

    with open(Pajek_net_file,'r') as f :

        lines = f.readlines()

        line_nb_elements = lines[0]

        nb_elements = int(line_nb_elements.split(' ')[1])

        print(nb_elements)

        node_corres = np.empty((nb_elements),dtype = 'int')

        node_lines = lines[1:(nb_elements+1)]

        #print lines

        for i,line in enumerate(node_lines):
            #print line

            new_index,old_index = line.split(' ')

            #print i+1, new_index, old_index

            #node_corres[i] = old_index

            if (i+1) == int(new_index) :
                node_corres[i] = int(old_index)-1
            else:
                print("Warning, incompatible indexes {} {}".format(new_index,i+1))

        list_sparse_matrix = [line.strip().split(' ') for line in lines[(nb_elements+2):]]

        np_list_sparse_matrix = np.transpose(np.array(list_sparse_matrix,dtype = 'int64'))

        sparse_matrix = sp.coo_matrix((np_list_sparse_matrix[2,:],(np_list_sparse_matrix[0,:]-1,np_list_sparse_matrix[1,:]-1)), shape = (nb_elements,nb_elements))

        f.close()

    return node_corres,sparse_matrix

#####################" compute modular Network
def compute_modular_network(sparse_matrix,community_vect):

    mod_mat = np.empty(sparse_matrix.todense().shape)

    mod_mat[:] = np.NAN

    for u,v,w in zip(sparse_matrix.row,sparse_matrix.col,sparse_matrix.data):

        if (community_vect[u] == community_vect[v]):

            mod_mat[u,v] = community_vect[u]
        else:

            mod_mat[u,v] = -1

    return mod_mat

### read strength from Network_Properties node results
def get_strength_values_from_info_nodes_file(info_nodes_file):

    from pandas.io.parsers import read_csv

    info_nodes = read_csv(info_nodes_file, sep = "\t")

    #print info_nodes

    return info_nodes['Strength'].values

def get_strength_pos_values_from_info_nodes_file(info_nodes_file):

    from pandas.io.parsers import read_csv

    info_nodes = read_csv(info_nodes_file, sep = "\t")

    print(info_nodes)

    return info_nodes['Strength_Pos']

def get_strength_neg_values_from_info_nodes_file(info_nodes_file):

    from pandas.io.parsers import read_csv

    info_nodes = read_csv(info_nodes_file, sep = "\t")

    print(info_nodes)

    return info_nodes['Strength_Neg']

def get_degree_pos_values_from_info_nodes_file(info_nodes_file):

    from pandas.io.parsers import read_csv

    info_nodes = read_csv(info_nodes_file, sep = "\t")

    print(info_nodes)

    return info_nodes['Degree_Pos']

def get_degree_neg_values_from_info_nodes_file(info_nodes_file):

    from pandas.io.parsers import read_csv

    info_nodes = read_csv(info_nodes_file, sep = "\t")

    print(info_nodes)

    return info_nodes['Degree_Neg']

##################################### Formatting data for external community detection algorithm (Louvain_Traag) ##############################

def export_Louvain_net_from_list(Z_Louvain_file,Z_list,coords):

    print(np.array(Z_list).shape)

    #print sig_x,sig_y
    print("column_stack")
    tab_edges = np.column_stack((np.array(Z_list),np.repeat(1,repeats = len(Z_list))))

    print(tab_edges.shape)

    print("file")

    with open(Z_Louvain_file,'w') as f:

        #### write node list
        nb_nodes = coords.shape[0]

        print("Nb nodes: " + str(nb_nodes))

        coords_list = coords.tolist()

        f.write('>\n')

        for node in range(nb_nodes):

            ##print node + coord label
            f.write(str(node+1) + ' ' + '_'.join(map(str,coords_list[node])) + '\n')

        #### write slice list
        f.write('>\n')
        f.write('1 1\n')

        #### write edge list
        f.write('>\n')

        np.savetxt(f,tab_edges,fmt = '%d %d %d %d')


<reponame>angus-g/sosfilt
import numpy as np
from scipy.signal import signaltools
from ._sosfilt import _sosfilt

def _validate_sos(sos):
    sos = np.atleast_2d(sos)

    if sos.ndim == 2:
        return sos, -1, 1

    n_filters, n_sections, m = sos.shape

    if m != 6:
        raise ValueError("last dim of sos array must be size 6")
    if not (sos[..., 3] == 1).all():
        raise ValueError("sos[..., 3] should be all ones")

    return sos, n_sections, n_filters

def _validate_nfilters(x, axis, n_filters):
    # with 1 filter, we'll broadcast it as usual
    if n_filters == 1:
        return

    # otherwise, the number of filters must match the
    # product of the non-filtered axes of x
    x_shape = list(x.shape)
    x_shape.pop(axis)
    if n_filters != np.product(x_shape):
        raise ValueError("n_filters must match product of non-filtered axes")

def _validate_ntaps(ntaps):
    first_tap = ntaps[0]
    if not np.all(ntaps == first_tap):
        raise ValueError("all filters must have the same number of taps")

    return first_tap

def sosfilt_zi(sos):
    sos = np.asarray(sos)
    if sos.ndim != 3 or sos.shape[2] != 6:
        raise ValueError("sos must be shape (n_filters, n_sections, 6)")

    if sos.dtype.kind in "bui":
        sos = sos.astype(np.float64)

    n_filters, n_sections = sos.shape[:2]
    zi = np.empty((n_filters, n_sections, 2), dtype=sos.dtype)
    scale = np.ones(n_filters, dtype=sos.dtype)
    for section in range(n_sections):
        b = sos[:, section, :3]
        a = sos[:, section, 3:]

        # lfilter_zi solves zi = A*zi + B
        # where A = scipy.linalg.companion(a).T
        # and   B = b[1:] - a[1:]*b[0]
        #
        # because a[0] = 1 for a sos filter, we have
        # A = [ 1 + a[1], -1;
        #         a[2],    1]
        # so zi[0] = (B[0] + B[1]) / (1 + a[1] + a[2])
        # and zi[1] = B[1] - a[2] * zi[0]
        #
        # we can pretty easily write this in a vectorised
        # way over all n_filters!

        B = b[:, 1:] - a[:, 1:] * b[:, [0]]
        zi[:, section, 0] = B.sum(axis=-1) / a.sum(axis=-1)
        zi[:, section, 1] = B[:, 1] - a[:, 2] * zi[:, section, 0]
        zi[:, section, :] *= scale[:, None]

        scale *= b.sum(axis=-1) / a.sum(axis=-1)

    return zi

def sosfilt(sos, x, axis=-1, zi=None):
    x = signaltools._validate_x(x)

    sos, n_sections, n_filters = _validate_sos(sos)

    x_zi_shape = list(x.shape)
    x_zi_shape[axis] = 2
    x_zi_shape = tuple([n_sections] + x_zi_shape)
    inputs = [sos, x]

    if zi is not None:
        inputs.append(np.asarray(zi))

    dtype = np.result_type(*inputs)
    if dtype.char not in "fdgFDGO":
        raise NotImplementedError(f"input type '{dtype}' not supported")

    if zi is not None:
        zi = np.array(zi, dtype)
        if zi.shape != x_zi_shape:
            raise ValueError("invalid zi shape")
        return_zi = True
    else:
        zi = np.zeros(x_zi_shape, dtype=dtype)
        return_zi = False

    axis = axis % x.ndim
    x = np.moveaxis(x, axis, -1)
    zi = np.moveaxis(zi, [0, axis + 1], [-2, -1])
    x_shape, zi_shape = x.shape, zi.shape
    x = np.reshape(x, (-1, x.shape[-1]))
    x = np.array(x, dtype, order="C")  # make a modifiable copy
    zi = np.ascontiguousarray(np.reshape(zi, (-1, n_sections, 2)))
    sos = sos.astype(dtype, copy=False)

    _sosfilt(sos, x, zi)

    x.shape = x_shape
    x = np.moveaxis(x, -1, axis)

    if return_zi:
        zi.shape = zi_shape
        zi = np.moveaxis(zi, [-2, -1], [0, axis + 1])
        out = (x, zi)
    else:
        out = x

    return out

def sosfiltfilt(sos, x, axis=-1, padtype="odd", padlen=None):
    """
    A forward-backward digital filter using cascaded second-order sections.

    Parameters
    ----------
    sos : array_like
        An array of second-order filter coefficients. It must have either
        the shape ``(n_filters, n_sections, 6)`` or ``(n_sections, 6)``.
        In the latter case, the single filter will be broadcast over the
        whole input array. In the former case, `n_filters` must match
        the product of the non-filter axes of `x`. Additionally,
        `n_sections` must be the same for all filters.
    x : array_like
        The array of data to be filtered.
    axis : int, optional
        The axis of `x` to which the filter is applied.
        Default is -1.
    padtype : str or None, optional
        Must be 'odd', 'even', 'constant', or None. Determines the type
        of extension to use for the padded signal to which the filter is
        applied. If None, no padding is used. The default is 'odd'.
    padlen : int or None, optional
        The number of elements by which to extend `x` at both ends of
        `axis` before applying the filter.
    """

    # filter validation
    sos, n_sections, n_filters = _validate_sos(sos)

    if n_filters == 1:
        # defer to scipy's implementation for the usual case of a single filter
        return signaltools.sosfiltfilt(
            sos, x, axis=axis, padtype=padtype, padlen=padlen
        )

    # input validation
    x = signaltools._validate_x(x)
    _validate_nfilters(x, axis, n_filters)

    # padding validation
    ntaps = np.ones(n_filters, dtype=int) * (2 * n_sections + 1)
    ntaps -= np.minimum((sos[..., 2] == 0).sum(axis=1), (sos[..., 5] == 0).sum(axis=1))
    ntaps = _validate_ntaps(ntaps)
    edge, ext = signaltools._validate_pad(padtype, padlen, x, axis, ntaps=ntaps)

    # filter initial conditions
    zi = sosfilt_zi(sos)

    # to handle multiple filters, we might want zi.shape = (n_sections, ..., 2, ...)
    zi = np.swapaxes(zi, 0, 1)  # => (n_sections, n_filters, 2)
    zi_shape = list(x.shape)
    zi_shape[axis] = 2
    # we need to swap axis to the end first, so it picks up the end of zi
    zi_shape[axis], zi_shape[-1] = zi_shape[-1], zi_shape[axis]
    zi_shape = [n_sections] + zi_shape
    zi = zi.reshape(zi_shape)  # should look like (n_sections, ..., 2)
    # now we need to swap axis back from the end to where it should be
    zi = np.swapaxes(zi, -1, (axis % x.ndim) + 1)

    # forward filter
    x_0 = signaltools.axis_slice(ext, stop=1, axis=axis)
    (y, zf) = sosfilt(sos, ext, axis=axis, zi=zi * x_0)

    # backward filter
    y_0 = signaltools.axis_slice(y, start=-1, axis=axis)
    (y, zf) = sosfilt(
        sos, signaltools.axis_reverse(y, axis=axis), axis=axis, zi=zi * y_0
    )

    # reshaping
    y = signaltools.axis_reverse(y, axis=axis)
    if edge > 0:
        y = signaltools.axis_slice(y, start=edge, stop=-edge, axis=axis)

    return y


import numpy as np
import matplotlib.pyplot as plt
import scipy.fftpack as fftp
import scipy.optimize as opt
import sys
import os
import string
from scipy import signal

from bcolors import *
from funct import *

def periodogram(x,y,z,t,name):
    vel=x*x+y*y
    components=[vel,z*z]
    F=[]
    A=[]

    for i in range(0,len(components)):
        window=np.kaiser(components[i].shape[-1],5)
        vel=components[i]#*window
        f,a=DO_FFT(vel,20)
        F.append(f)
        A.append(a)

    plt.figure(1,figsize=(11,7))
    plt.title("Spectral Gap - <NAME>")
    plt.subplot(121)
    plt.title("FFT horizontal velocity")
    plt.xlabel("Frequency [Hz]")
    plt.ylabel("Power Spectrum [dB]")
    plt.loglog(F[0],A[0],'k', markersize=0.1)

    plt.subplot(122)
    plt.title("FFT vertical velocity")
    plt.xlabel("Frequency [Hz]")
    plt.ylabel("Power Spectrum [dB]")
    plt.loglog(F[1],A[1],'r',markersize=0.1)

    try:
        plt.savefig("graph/"+name+"_FFT.png", figuresize=(8,6), dpi=320, format="png")
        print_ok("Graph saved in: "+"graph/"+name+"_FFT.png")
    except IOError as IoE:
        print_fail("I/O Error! Erro number = {0}; {1}".format(IoE.errno,IoE.strerror))
        exit(IoE.errno)


<filename>PAM_repeat1000times-regression.py<gh_stars>1-10
import pandas as pd
import numpy as np
import xgboost as xgb
import shap
from scipy.stats import pearsonr

from sklearn.preprocessing import StandardScaler
from sklearn import metrics
from sklearn.model_selection import KFold, GridSearchCV

import time
import random

from utils import data_handler, plotter

#handle warnings
import warnings
warnings.filterwarnings("ignore")

np.random.seed(3)

def test(y_pred,y_true,end_ptr):
    '''Evaluate the performance of PAM through 3 metrics;
       And find the best predicted condition for the next PAM trial

       Inputs:
           X_test      - for testing
           y_true      - ground true for X_test
           end_ptr     - range of [0,end_ptr] in y_pred and y_true are values observable to PAM

       Uses:(global variables)
           pred_result - of shape (total_samp,), storing all the results predicted by PAM
           true_results- of shape (total_samp,), the same order as of pred_results, storing all the ground truth

       Outputs:

    '''
    #All sample metrics
    r2 = metrics.r2_score(y_true,y_pred)
    mse = metrics.mean_squared_error(y_true,y_pred)
    [pear,p_value] = pearsonr(y_true,y_pred)

    #metrics based on observable samples
    y_true_s = y_true[0:(end_ptr+1)]
    y_pred_s = y_pred[0:(end_ptr+1)]
    r2_s = metrics.r2_score(y_true_s,y_pred_s)
    mse_s = metrics.mean_squared_error(y_true_s,y_pred_s)
    [pear_s,p_value_s] = pearsonr(y_true_s,y_pred_s)
    return [r2,pear,p_value,mse,r2_s,pear_s,p_value_s,mse_s] #, best_pos_ind, best_prob

# data handling
X_df,Y_df = data_handler.load_XY(1)
X = X_df.as_matrix()
Y = Y_df.as_matrix() / 100

# setup and construct initial training dataset
# cross validation settup
inner_nsplits = 10
init_train_size = 20
totalSamp = X.shape[0]

Y_global_max = np.max(Y)
all_ind = np.random.permutation(list(range(0,totalSamp)))
all_ind_wo_max = list(range(0,totalSamp))
all_ind_wo_max.remove(0)

# PAM guided synthesis
def PAM_regression(save_csv= False, verbose=False, to_break=True, title='cqd_PAM_',batch=1):

    ## start PAM guided synthesis...
    init_time = time.time()
    Nc = 0

    #construct initial training set
    results_mat = np.zeros(((totalSamp-init_train_size),12))

    train_ind = random.sample(all_ind_wo_max, init_train_size)
    test_ind = [x for x in all_ind if x not in train_ind]
    if(verbose):
        print('initial training set indexes',train_ind)

    # set up result storage to compute eval metrics, in the order of PAM
    #  ignore the initial training set, as it is not determined by PAM
    pred_results = np.zeros(totalSamp-init_train_size)
    true_results = np.zeros(totalSamp-init_train_size)

    # setup the hyperparameter range for tuning
    tuned_parameters = dict(learning_rate=[0.01],
                            n_estimators=[300,500,700], #100,,300,400,500
                            colsample_bylevel = [0.5,0.7,0.9],
                          gamma=[0,0.2], #0,0.1,0.2,0.3,0.4
                          max_depth =[3,7,11], # [3,7,11]]
                          reg_lambda = [0.1,1,10], #[0.1,1,10]
                         # reg_alpha = [1],
                           subsample=[0.4,0.7,1])

    j=0
    loop_count = 0
    mean_y_only_init = np.mean(Y[train_ind])
    std_y_only_init = np.std(Y[train_ind])

    while(j<totalSamp-init_train_size):
        inner_cv = KFold(n_splits=inner_nsplits,shuffle=True, random_state=j)
        X_train = X[train_ind]
        Y_train = Y[train_ind]
        X_test = X[test_ind]
        Y_test = Y[test_ind]

        last_max = np.max(Y_train)

        # GradientBoost
        reg = xgb.XGBRegressor(objective="reg:linear",min_child_weight=1,**{'tree_method':'exact'},
                                 silent=True,n_jobs=4,random_state=3,seed=3);

        gb_clf = GridSearchCV(reg,tuned_parameters, cv=inner_cv,scoring='r2',verbose=0,n_jobs=4)
        gb_clf.fit(X_train, Y_train)
        y_pred = gb_clf.predict(X_test)

        # choose the batch of conditions with best predicted yield
        best_pos_ind = np.argsort(-y_pred)[:batch]
        best_prob = y_pred[best_pos_ind]
        next_ind = np.array(test_ind)[best_pos_ind]

        # update results storage
        train_size = len(Y_train)
        temp = list(range(0,len(y_pred)))
        ind_notbest = [x for x in temp if x not in best_pos_ind]

        start_ptr = j
        end_ptr = np.min([start_ptr + batch, totalSamp-init_train_size])
        pred_results[start_ptr:end_ptr] = best_prob
        pred_results[end_ptr:totalSamp-init_train_size]= y_pred[ind_notbest]
        true_results[start_ptr:end_ptr] = Y_test[best_pos_ind]
        true_results[end_ptr:totalSamp-init_train_size] = Y_test[ind_notbest]

        pred_metrics = test(pred_results,true_results,end_ptr-1)

        # calculate results
        next_best_true_ind = next_ind[np.argmax(Y_test[best_pos_ind])]
        next_best_y_true = np.max(Y_test[best_pos_ind])
        result_list = [train_size,next_best_true_ind,next_best_y_true,best_prob[0],] + pred_metrics
        results_mat[loop_count,:] = np.array(result_list)

        loop_count = loop_count +1
        j = j + batch

        if(verbose):
            print(loop_count,'->',j,', best_next_ind=',next_best_true_ind, ' best_Y_true=',"{0:.6f}".format(next_best_y_true),' train_max=',"{0:.6f}".format(last_max),' r2=',pred_metrics[0])

        train_ind = [*train_ind , *next_ind]
        test_ind = [x for x in test_ind if x not in next_ind]

        ## critical point
        if(next_best_y_true==Y_global_max and Nc == 0):
            Nc = j+init_train_size
            if(to_break):
                break

    saved_title = '-'
    if(save_csv):
        results = pd.DataFrame(data=results_mat[0:j,:],columns=['sample_size','pred_ind','best_pred_result','y_true','r2','pearson','p_value','mse','r2_s','pearson_s','p_value_s','mse_s'])
        saved_title = data_handler.save_csv(results,title=title)

    # compute stats
    mean_y_wo_init =  np.mean(true_results[0:j])
    std_y_wo_init = np.std(true_results[0:j])

    mean_y_w_init = np.mean(Y[train_ind])
    std_y_w_init = np.std(Y[train_ind])

    run_time = (time.time() - init_time)/60

    return [saved_title, Nc,mean_y_wo_init,std_y_wo_init,mean_y_w_init,std_y_w_init,mean_y_only_init,std_y_only_init, run_time]

outer_loop = 10
inner_loop = 100

print('start PAM for ',str(outer_loop*inner_loop),' times...')
# save the results some repetitions for backup
for j in range(0,outer_loop):

    #PAM_results = np.zeros((inner_loop,9))
    init_time = time.time()
    res_arr = []

    for i in range(0,inner_loop):

        loop_count = j*inner_loop + i
        result = PAM_regression(save_csv= False, verbose=False, to_break = True, title='cqd_PAM_'+str(loop_count)+'th_loop_')
        res_arr.append(result)
        print(str(loop_count),' -> ',str(result[0]),'  time=',result[len(result)-1])

    PAM_df = pd.DataFrame(data=res_arr, columns=['file-name','num_experiments','mean_y_wo_init','std_y_wo_init','mean_y_w_init','std_y_w_init','mean_y_only_init','std_y_only_init', 'run_time'])
    saved_path = data_handler.save_csv(PAM_df,title='cqd_PAM_'+str(inner_loop)+'times_')
    print('total = ',str((time.time()-init_time)/3600),'  hrs  >>-------saved')


<reponame>nids2001/UncertainSCI
from itertools import combinations

import numpy as np
from scipy import special as sp
from scipy.special import comb

from UncertainSCI.utils.prob import discrete_sampling

def hyperbolic_cross_indices(d, k):
    """
    Returns indices associated with a d-dimensional (isotropic)
    hyperbolic cross index space up to degree k.
    """

    assert k >= 0
    assert d >= 1

    if d == 1:
        lambdas = range(k+1)
        return lambdas

    lambdas = np.zeros([1, d], dtype=int)

    # First add all indices with sparsity 1
    for q in range(d):
        temp = np.zeros([k-1, d], dtype=int)
        temp[:, q] = np.arange(1, k, dtype=int)
        lambdas = np.vstack([lambdas, temp])

    # Now determine the maximum 0-norm the entries can be. I.e., for
    # which values of p is 2^p <= k+1?
    pmax = int(np.floor(np.log(k+1)/np.log(2)))

    # For each sparsity p, populate with all possible indices of that
    # sparsity
    for p in range(2, pmax+1):
        # Determine all possible locations where nonzero entries can occur
        combs = combinations(range(d), p)
        combs = np.array([row for row in combs], dtype=int)

        # Now we have 2^p < k+1, i.e., an index with nonzero entries
        # np.ones([p 1]) is ok.
        # Keep incrementing these entries until product exceeds k+1
        possible_indices = np.ones([1, p])
        ind = 0

        while ind < possible_indices.shape[0]:
            # Add any possibilities that are children of
            # possible_indices[ind,:]

            lambd = possible_indices[ind, :]
            for q in range(p):
                temp = lambd.copy()
                temp[q] += 1
                if np.prod(temp+1) <= k+1:
                    possible_indices = np.vstack([possible_indices, temp])

            ind += 1

        possible_indices = np.vstack({tuple(row) for row in possible_indices})
        arow = lambdas.shape[0]
        lambdas = np.vstack([lambdas,
                             np.zeros([combs.shape[0]*possible_indices.shape[0],
                                       d],
                                      dtype=int)])

    # Now for each combination, we put in possible_indices
        for c in range(combs.shape[0]):
            i1 = arow
            i2 = arow + possible_indices.shape[0]

            lambdas[i1:i2, combs[c, :]] = possible_indices

            arow = i2

    return lambdas

def total_degree_indices(d, k):
    # Returns multi-indices associated with d-variate polynomials of
    # degree less than or equal to k. Each row is a multi-index, ordered
    # in total-degree-graded reverse lexicographic ordering.

    assert d > 0
    assert k >= 0

    if d == 1:
        return np.arange(k+1, dtype=int).reshape([k+1, 1])

    # total degree indices up to degree k in d-1 dimensions:
    lambdasd1 = total_degree_indices(d-1, k)
    # lambdasd1 should already be sorted by total degree, which is
    # assumed below

    lambdas = np.zeros([np.round(int(comb(d+k, d))), d], dtype=int)

    i0 = 0
    for qk in range(0, k+1):

        n = int(np.round(comb(d-1+(k-qk), d-1)))
        i1 = i0 + n

        lambdas[i0:i1, 0] = qk
        lambdas[i0:i1, 1:] = lambdasd1[:n, :]
        i0 = i1

    # My version of numpy < 1.12, so I don't have np.flip :(
    # degrees = np.cumsum(np.flip(lambdas,axis=1), axis=1)
    degrees = np.cumsum(np.fliplr(lambdas), axis=1)

    ind = np.lexsort(degrees.transpose())
    lambdas = lambdas[ind, :]

    return lambdas

def degree_encompassing_N(d, N):
    # Returns the smallest degree k such that nchoosek(d+k,d) >= N

    k = 0
    while np.round(comb(d+k, d)) < N:
        k += 1

    return k

def total_degree_indices_N(d, N):
    # Returns the first N ( > 0) d-dimensional multi-indices when ordered by
    # total degree graded reverse lexicographic ordering.

    assert N > 0

    return total_degree_indices(d, degree_encompassing_N(d, N))[:N, :]

def tensor_product(d, k):
    # Returns multi-indices associated with d-variate polynomials of
    # degree less than or equal to k. Each row is a multi-index, ordered
    # in p degree, where p = Inf

    from itertools import product

    Ival = np.empty(shape=[0, d], dtype=int)

    for t in (_ for _ in product(range(k+1), repeat=d)):
        Ival = np.vstack((Ival, np.asarray(t)))

    return Ival

def multi_indices_degree(d, k, p):
    # Returns multi-indices associated with d-variate polynomials of
    # degree less than or equal to k. Each row is a multi-index, ordered
    # in p degree, p could be any positive number including numpy.inf

    if p < 1:
        lambdas = total_degree_indices(d, k)
        norm = (np.sum(lambdas**p, axis=1))**(1/p)
        norm = np.round(norm, decimals=8)
        flags = (norm <= k)
        lambdas = lambdas[flags]

    elif p == np.inf:
        lambdas = tensor_product(d, k)

    elif p == 1:
        lambdas = total_degree_indices(d, k)

    else:
        lambdas = tensor_product(d, k)
        norm = (np.sum(lambdas**p, axis=1))**(1/p)
        norm = np.round(norm, decimals=8)
        flags = (norm <= k)
        lambdas = lambdas[flags]

    return lambdas

def pdjk(d, k):
    j = np.arange(k+1)
    p = np.exp(np.log(d) + sp.gammaln(k+1) - sp.gammaln(j+1) +
               sp.gammaln(j+d) - sp.gammaln(k+d+1))
    assert np.abs(sum(p)-1) < 1e-8
    return p

def sampling_total_degree_indices(N, d, k):

    """
    Chooses N random multi-indices (with the uniform probability law) from the
    set of d-variate multi-indices whose total degree is k and less

    Parameters
    ------
    param1: N
    Numebr of chosen random multi-indices
    param2: d
    dimension of variables
    param3L k
    total degree of variables

    Returns
    ------
    The output lambdas is an N x d matrix, with each row containing one of
    these multi-indices
    """
    lambdas = np.zeros((N, d))

    degrees = discrete_sampling(N, pdjk(d, k), np.arange(k+1)).T

    for i in range(1, d):
        for n in range(1, N+1):
            lambdas[n-1, i-1] = \
                discrete_sampling(1, pdjk(d-i, degrees[n-1]),
                                  np.arange(degrees[n-1], 0-1e-8, -1))

        degrees = degrees - lambdas[:, i-1]

    lambdas[:, d-1] = degrees

    return lambdas

class MultiIndexSet():
    def __init__(self, dim=None):
        self.dim = dim
        self.indices = np.zeros([0, self.dim])
        self.adaptive = False

    def get_indices(self):
        return self.indices

    def size(self):
        return self.indices.shape[0]

    def max_univariate_degree(self, dim=None):
        """
        Returns the maximum index value along a given dimension. If no
        dimension is specified, returns the maximum degree across all
        dimensions.

        Args:
            dim: A dimension integer taking values between 0 and self.dim-1,
              optional.
        Returns:
            deg: Maximum index value (integer).
        """

        if dim is None:
            return np.max(self.get_indices())
        else:
            return np.max(self.get_indices()[:, dim])

    def zero_indices(self, dim_indices=None):
        """
        Identifies indices in the index set whose entries in the dimensions
        dim_indices are 0.

        Args:
            dim_list: list-like iterable containing dimension indices.
        Returns:
            flags: A boolean numpy vector indicating which indices in the index
              set satisfy the conditions.
        """

        assert all([0 <= dim <= self.dim-1 for dim in dim_indices])

        if dim_indices is None:
            return self.get_indices()

        indices = self.get_indices()

        return np.linalg.norm(indices[:, dim_indices], axis=1) == 0

    def isamember(self, trial_indices):
        """
        Determines if input indices are members of the current index set.

        Args:
            trial_indices: An :math:`K \\times d` numpy array, where each row
              corresponds to an index.

        Returns:
            member: A numpy boolean array of size :math:`K` indicating if the
              rows of trial_indices are part of the current index set.
        """

        K, d = trial_indices.shape
        assert self.dim == d, \
               "Input index array should have {0:d} columns".format(self.dim)

        M = self.indices.shape[0]
        member = np.zeros(K, dtype=bool)
        for m in range(K):

            index = trial_indices[m, :]

            matches = np.ones(M, dtype=bool)
            for q in range(d):
                matches[matches] = index[q] == self.indices[matches, q]

            if np.any(matches):
                member[m] = True

        return member

    def get_margin(self):
        """
        Computes the margin of the index set :math:`\\Lambda`. In :math:`d`
        dimensions, this is defined as the set of indices :math:`\\lambda \\in
        N_0^d \\backslash \\Lambda` such that

        .. math::

          \\lambda - e_j \\in \\Lambda

        for some :math:`j = 1, \\ldots, d`.

        Returns:
            margin: A numpy array of size :math:`M \\times d` where each row
              contains an index in the margin.
        """

        # Do this in a brute-force manner:
        # - search for leaves of the current index set as margin candidates
        # - weed out leaves that are not in the margin

        M, d = self.indices.shape
        margin = np.zeros([0, d], dtype=self.indices.dtype)

        for m in range(M):
            candidates = np.tile(self.indices[m, :], [d, 1]) +\
                         np.eye(d, dtype=self.indices.dtype)
            membership_flags = ~self.isamember(candidates)
            margin = np.unique(
                      np.append(margin, candidates[membership_flags, :], axis=0),
                      axis=0)

        return margin

    def get_reduced_margin(self):
        """
        Computes the reduced margin of the index set :math:`\\Lambda`. In
        :math:`d` dimensions, this is defined as the set of indices
        :math:`\\lambda \\in N_0^d \\backslash \\Lambda` such that

        .. math::

          \\lambda - e_j \\in \\Lambda

        for every :math:`j = 1, \\ldots, d` satisfying :math:`\\lambda_j \\neq
        0`.

        Returns:
            margin: A numpy array of size :math:`M \\times d` where each row
              contains an index in the margin.
        """

        # We'll sequentially test elements in the margin

        margin = self.get_margin()
        K, d = margin.shape
        reduced_margin_inds = []

        for k in range(K):
            candidates = np.tile(margin[k, :], [d, 1]) - np.eye(d)
            candidates = candidates[~np.any(candidates < 0, axis=1), :]

            if np.all(self.isamember(candidates)):
                reduced_margin_inds.append(k)

        return margin[reduced_margin_inds, :]

    def augment(self, indices):
        """
        Augments the index set with the given indices.
        """

        K, d = indices.shape
        assert d == self.dim, \
               "Input index array should have {0:d} columns".format(self.dim)

        membership_flags = self.isamember(indices)
        if np.any(~membership_flags):
            self.adaptive = True
            self.indices = np.append(self.indices,
                                     indices[~membership_flags, :], axis=0)

class LpSet(MultiIndexSet):
    def __init__(self, dim=1, order=0, p=1):
        assert dim > 0 and order >= 0 and p >= 0

        super().__init__(dim=dim)
        self.dim = dim
        self.order = order
        self.p = p
        self.indices = self.get_indices()

    def get_indices(self):
        if self.p < 1:
            lambdas = total_degree_indices(self.dim, self.order)
            norm = (np.sum(lambdas**self.p, axis=1))**(1/self.p)
            norm = np.round(norm, decimals=8)
            flags = (norm <= self.order)
            lambdas = lambdas[flags]

        elif self.p == np.inf:
            lambdas = tensor_product(self.dim, self.order)

        elif self.p == 1:
            lambdas = total_degree_indices(self.dim, self.order)

        else:
            lambdas = tensor_product(self.dim, self.order)
            norm = (np.sum(lambdas**self.p, axis=1))**(1/self.p)
            norm = np.round(norm, decimals=8)
            flags = (norm <= self.order)
            lambdas = lambdas[flags]

        return lambdas

class TotalDegreeSet(MultiIndexSet):
    def __init__(self, dim=1, order=0):
        assert dim > 0 and order >= 0

        super().__init__(dim=dim)

        self.dim, self.order = dim, order
        self.indices = self.get_indices()

    def get_indices(self):
        if self.adaptive:
            return super().get_indices()
        else:
            return total_degree_indices(self.dim, self.order)

class HyperbolicCrossSet(MultiIndexSet):
    def __init__(self, dim=1, order=0):
        assert dim > 0 and order >= 0

        super().__init__(dim=dim)

        self.dim, self.order = dim, order
        self.indices = self.get_indices()

    def get_indices(self):
        if self.adaptive:
            return super().get_indices()
        else:
            return hyperbolic_cross_indices(self.dim, self.order)

if __name__ == "__main__":

    pass


# coding: utf-8

# In[1]:

import matplotlib.pyplot as plt
import sgbm
import stdev
import simple_return as sr
import duration
import numpy as np
import pandas as pd
import statistics as sc
from scipy.stats import kurtosis, skew
import pickle
import os
import multiprocessing
from multiprocessing import Pool

# In[2]:

x=pd.ExcelFile('/home/anindya/Parallel_code/Squeeze/Stat.xlsx')
page=x.parse(0)
N=page.N
mu=page.tdrift
tvola=page.tvola
p=0.15

# In[3]:

with open('/home/anindya/Parallel_code/Squeeze/vol.pkl','rb') as f:
    vol = pickle.load(f)

# In[4]:

seed=range(0,201)
dt=5./(250*360)
sdt=np.sqrt(dt)
name="/home/anindya/Parallel_code/Squeeze/SMGBM/Group1/"

# In[5]:

def loop(k,theta,r,b):
    GM_m=[]
    GS_m=[]
    GW_m=[]
    GK_m=[]
    l=[]
    if(tvola[k]-((p*(np.percentile(vol[k],r))))>0):
        si1=np.percentile(vol[k],r)
    else:
        si1=tvola[k]/p
    si2=((tvola[k]-(p*si1))/(1-p))
    for j in range(200):
        sm=sgbm.sgbm(dt,N[k],p*100,mu[k],si1,si2,theta,b,seed[j])
        ret_m=sr.s_ret(np.array(sm,dtype=float))
        ret_m=np.array(ret_m)
        L=len(ret_m)
        n=20
        new_ret_m=[np.array(ret_m[i:i+n]) for i in range(L-n)]
        Ln=len(new_ret_m)
        new_std_m=np.array([stdev.sd(new_ret_m[i]) for i in range(Ln)])
        volatility_m= new_std_m/sdt
        dur_m=duration.duration(np.array(volatility_m))
        dur_m=np.array(dur_m,dtype=float)
        GM_m.append(np.mean(dur_m))
        GS_m.append(stdev.sd(dur_m))
        GW_m.append(skew(dur_m))
        GK_m.append(kurtosis(dur_m,fisher=False))
        l.append(len(dur_m))
    return (GM_m,GS_m,GW_m,GK_m,l)
#     return (GM_m,GS_m)

# In[6]:

pool = Pool()

for k in range(0,6,1):
    name1=name+"I0"+str(k+1)
    if not os.path.exists(name1):
        os.mkdir(name1)
    for b in range(2,15,1):
        name2=name1+"/shape"+str(b)
        if not os.path.exists(name2):
            os.mkdir(name2)
        for i in range(5,17,1):
            t1=[]
            t2=[]
            t3=[]
            t4=[]
            L1=[]
            theta=i
            args=[]
            for r in range(0,31,1):
                args.append((k,theta,r,b))
            for one,two,three,four,ll in pool.starmap(loop, args):
                t1.append(one)
                t2.append(two)
                t3.append(three)
                t4.append(four)
                L1.append(ll)
            strname=name2+"/theta"+str(i)+".xlsx"
            writer=pd.ExcelWriter(strname,engine='xlsxwriter')
            for w in range(0,31,1):
                df=pd.DataFrame({'T1':t1[w],'T2':t2[w],'T3':t3[w],'T4':t4[w],'Len':L1[w]},index=range(1,201))
                vv="Vola"+str((w))
                df.to_excel(writer,sheet_name=vv)
            writer.save()

pool.close()


# Copyright (c) Facebook, Inc. and its affiliates.

# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import os
import cv2
import torch
import numpy as np
import argparse
import torchvision
from PIL import Image
from tqdm import tqdm
from pathlib import Path
from datetime import datetime
from retry.api import retry_call
from torch.utils import data
from torchvision import transforms
from part_selector import Trainer as Trainer_selector
from part_generator import Trainer as Trainer_cond_unet
from scipy.ndimage.morphology import distance_transform_edt

COLORS = {'initial':1-torch.cuda.FloatTensor([45, 169, 145]).view(1, -1, 1, 1)/255., 'eye':1-torch.cuda.FloatTensor([243, 156, 18]).view(1, -1, 1, 1)/255., 'none':1-torch.cuda.FloatTensor([149, 165, 166]).view(1, -1, 1, 1)/255.,
        'beak':1-torch.cuda.FloatTensor([211, 84, 0]).view(1, -1, 1, 1)/255., 'body':1-torch.cuda.FloatTensor([41, 128, 185]).view(1, -1, 1, 1)/255., 'details':1-torch.cuda.FloatTensor([171, 190, 191]).view(1, -1, 1, 1)/255.,
        'head':1-torch.cuda.FloatTensor([192, 57, 43]).view(1, -1, 1, 1)/255., 'legs':1-torch.cuda.FloatTensor([142, 68, 173]).view(1, -1, 1, 1)/255., 'mouth':1-torch.cuda.FloatTensor([39, 174, 96]).view(1, -1, 1, 1)/255.,
        'tail':1-torch.cuda.FloatTensor([69, 85, 101]).view(1, -1, 1, 1)/255., 'wings':1-torch.cuda.FloatTensor([127, 140, 141]).view(1, -1, 1, 1)/255.}

class Initialstroke_Dataset(data.Dataset):
    def __init__(self, folder, image_size):
        super().__init__()
        self.folder = folder
        self.image_size = image_size
        self.paths = [p for p in Path(f'{folder}').glob(f'**/*.png')]
        self.transform = transforms.Compose([
            transforms.ToTensor(),
        ])

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

    def __getitem__(self, index):
        path = self.paths[index]
        img = self.transform(Image.open(path))
        return img

    def sample(self, n):
        sample_ids = [np.random.randint(self.__len__()) for _ in range(n)]
        samples = [self.transform(Image.open(self.paths[sample_id])) for sample_id in sample_ids]
        return torch.stack(samples).cuda()

def load_latest(model_dir, name):
    model_dir = Path(model_dir)
    file_paths = [p for p in Path(model_dir / name).glob('model_*.pt')]
    saved_nums = sorted(map(lambda x: int(x.stem.split('_')[1]), file_paths))
    if len(saved_nums) == 0:
        return
    name = saved_nums[-1]
    print(f'continuing from previous epoch - {name}')
    return name

def noise(n, latent_dim):
    return torch.randn(n, latent_dim).cuda()

def noise_list(n, layers, latent_dim):
    return [(noise(n, latent_dim), layers)]

def mixed_list(n, layers, latent_dim):
    tt = int(torch.rand(()).numpy() * layers)
    return noise_list(n, tt, latent_dim) + noise_list(n, layers - tt, latent_dim)

def image_noise(n, im_size):
    return torch.FloatTensor(n, im_size, im_size, 1).uniform_(0., 1.).cuda()

def evaluate_in_chunks(max_batch_size, model, *args):
    split_args = list(zip(*list(map(lambda x: x.split(max_batch_size, dim=0), args))))
    chunked_outputs = [model(*i) for i in split_args]
    if len(chunked_outputs) == 1:
        return chunked_outputs[0]
    return torch.cat(chunked_outputs, dim=0)

def evaluate_in_chunks_unet(max_batch_size, model, map_feats, *args):
    split_args = list(zip(*list(map(lambda x: x.split(max_batch_size, dim=0), args))))
    split_map_feats = list(zip(*list(map(lambda x: x.split(max_batch_size, dim=0), map_feats))))
    chunked_outputs = [model(*i, j) for i, j in zip(split_args, split_map_feats)]
    if len(chunked_outputs) == 1:
        return chunked_outputs[0]
    return torch.cat(chunked_outputs, dim=0)

def styles_def_to_tensor(styles_def):
    return torch.cat([t[:, None, :].expand(-1, n, -1) for t, n in styles_def], dim=1)

def gs_to_rgb(image, color):
    image_rgb = image.repeat(1, 3, 1, 1)
    return 1-image_rgb*color

@torch.no_grad()
def generate_truncated(S, G, style, noi, trunc_psi = 0.75, num_image_tiles = 8, bitmap_feats=None, batch_size=8):
    latent_dim = G.latent_dim
    z = noise(2000, latent_dim)
    samples = evaluate_in_chunks(batch_size, S, z).cpu().numpy()
    av = np.mean(samples, axis = 0)
    av = np.expand_dims(av, axis = 0)

    w_space = []
    for tensor, num_layers in style:
        tmp = S(tensor)
        av_torch = torch.from_numpy(av).cuda()
        # import ipdb;ipdb.set_trace()
        tmp = trunc_psi * (tmp - av_torch) + av_torch
        w_space.append((tmp, num_layers))

    w_styles = styles_def_to_tensor(w_space)
    generated_images = evaluate_in_chunks_unet(batch_size, G, bitmap_feats, w_styles, noi)
    return generated_images.clamp_(0., 1.)

@torch.no_grad()
def generate_part(model, partial_image, partial_rgb, color=None, percentage=20, num=0, num_image_tiles=8, trunc_psi=1., save_img=False, results_dir='../results', evolvement=False):
    model.eval()
    ext = 'png'
    num_rows = np.sqrt(num_image_tiles)
    latent_dim = model.G.latent_dim
    image_size = model.G.image_size
    num_layers = model.G.num_layers
    if percentage == 'eye':
        n_eye = 10
        generated_partial_images_candidates = []
        scores = torch.zeros(n_eye)
        for _ in range(n_eye):
            latents_z = noise_list(num_image_tiles, num_layers, latent_dim)
            n = image_noise(num_image_tiles, image_size)
            image_partial_batch = partial_image[:, -1:, :, :]
            bitmap_feats = model.Enc(partial_image)
            generated_partial_images = generate_truncated(model.S, model.G, latents_z, n, trunc_psi = trunc_psi, bitmap_feats=bitmap_feats)
            generated_partial_images_candidates.append(generated_partial_images)
        generated_partial_images_candidates = torch.cat(generated_partial_images_candidates, 0)
        # eye size rank
        n_pixels = generated_partial_images_candidates.sum(-1).sum(-1).sum(-1) # B
        for rank, i_eye in enumerate(torch.argsort(n_pixels, descending=True)):
            scores[i_eye] += (rank+1)/n_eye
        # eye distance rank
        initial_stroke = partial_image[:, :1].cpu().data.numpy()
        initial_stroke_dt = torch.cuda.FloatTensor(distance_transform_edt(1-initial_stroke))
        dt_pixels = (generated_partial_images_candidates*initial_stroke_dt).sum(-1).sum(-1).sum(-1) # B
        for rank, i_eye in enumerate(torch.argsort(dt_pixels, descending=False)): # the smaller the better
            if n_pixels[i_eye] > 3:
                scores[i_eye] += (rank+1)/n_eye
        generated_partial_images = generated_partial_images_candidates[torch.argsort(scores, descending=True)[0]].unsqueeze(0)
    else:
        # latents and noise
        latents_z = noise_list(num_image_tiles, num_layers, latent_dim)
        n = image_noise(num_image_tiles, image_size)
        image_partial_batch = partial_image[:, -1:, :, :]
        bitmap_feats = model.Enc(partial_image)
        generated_partial_images = generate_truncated(model.S, model.G, latents_z, n, trunc_psi = trunc_psi, bitmap_feats=bitmap_feats)
    # regular
    generated_partial_images = generate_truncated(model.S, model.G, latents_z, n, trunc_psi = trunc_psi, bitmap_feats=bitmap_feats)
    generated_partial_rgb = gs_to_rgb(generated_partial_images, color)
    generated_images = generated_partial_images + image_partial_batch
    generated_rgb = 1 - ((1-generated_partial_rgb)+(1-partial_rgb))
    if save_img:
        torchvision.utils.save_image(generated_partial_rgb, os.path.join(results_dir, f'{str(num)}-{percentage}-comp.{ext}'), nrow=num_rows)
        torchvision.utils.save_image(generated_rgb, os.path.join(results_dir, f'{str(num)}-{percentage}.{ext}'), nrow=num_rows)
    return generated_partial_images.clamp_(0., 1.), generated_images.clamp_(0., 1.), generated_partial_rgb.clamp_(0., 1.), generated_rgb.clamp_(0., 1.)

def train_from_folder(
    data_path = '../../data',
    results_dir = '../../results',
    models_dir = '../../models',
    n_part = 1,
    image_size = 128,
    network_capacity = 16,
    batch_size = 3,
    num_image_tiles = 8,
    trunc_psi = 0.75,
    generate_all=False,
):
    min_step = 299
    name_eye='short_bird_creative_sequential_r6_partstack_aug_eye_unet2_largeaug'
    load_from = load_latest(models_dir, name_eye)
    load_from = min(min_step, load_from)
    model_eye = Trainer_cond_unet(name_eye, results_dir,  models_dir, n_part=n_part, batch_size=batch_size, image_size=image_size, network_capacity=network_capacity)
    model_eye.load_config()
    model_eye.GAN.load_state_dict(torch.load('%s/%s/model_%d.pt'%(models_dir, name_eye, load_from)))

    name_head='short_bird_creative_sequential_r6_partstack_aug_head_unet2'
    load_from = load_latest(models_dir, name_head)
    load_from = min(min_step, load_from)
    model_head = Trainer_cond_unet(name_head, results_dir,  models_dir, n_part=n_part, batch_size=batch_size, image_size=image_size, network_capacity=network_capacity)
    model_head.load_config()
    model_head.GAN.load_state_dict(torch.load('%s/%s/model_%d.pt'%(models_dir, name_head, load_from)))

    name_body='short_bird_creative_sequential_r6_partstack_aug_body_unet2'
    load_from = load_latest(models_dir, name_body)
    load_from = min(min_step, load_from)
    model_body = Trainer_cond_unet(name_body, results_dir,  models_dir, n_part=n_part, batch_size=batch_size, image_size=image_size, network_capacity=network_capacity)
    model_body.load_config()
    model_body.GAN.load_state_dict(torch.load('%s/%s/model_%d.pt'%(models_dir, name_body, load_from)))

    name_beak='short_bird_creative_sequential_r6_partstack_aug_beak_unet2'
    load_from = load_latest(models_dir, name_beak)
    load_from = min(min_step, load_from)
    model_beak = Trainer_cond_unet(name_beak, results_dir,  models_dir, n_part=n_part, batch_size=batch_size, image_size=image_size, network_capacity=network_capacity)
    model_beak.load_config()
    model_beak.GAN.load_state_dict(torch.load('%s/%s/model_%d.pt'%(models_dir, name_beak, load_from)))

    name_legs='short_bird_creative_sequential_r6_partstack_aug_legs_unet2'
    load_from = load_latest(models_dir, name_legs)
    load_from = min(min_step, load_from)
    model_legs = Trainer_cond_unet(name_legs, results_dir,  models_dir, n_part=n_part, batch_size=batch_size, image_size=image_size, network_capacity=network_capacity)
    model_legs.load_config()
    model_legs.GAN.load_state_dict(torch.load('%s/%s/model_%d.pt'%(models_dir, name_legs, load_from)))

    name_wings='short_bird_creative_sequential_r6_partstack_aug_wings_unet2'
    load_from = load_latest(models_dir, name_wings)
    load_from = min(min_step, load_from)
    model_wings = Trainer_cond_unet(name_wings, results_dir,  models_dir, n_part=n_part, batch_size=batch_size, image_size=image_size, network_capacity=network_capacity)
    model_wings.load_config()
    model_wings.GAN.load_state_dict(torch.load('%s/%s/model_%d.pt'%(models_dir, name_wings, load_from)))

    name_mouth='short_bird_creative_sequential_r6_partstack_aug_mouth_unet2'
    load_from = load_latest(models_dir, name_mouth)
    load_from = min(min_step, load_from)
    model_mouth = Trainer_cond_unet(name_mouth, results_dir,  models_dir, n_part=n_part, batch_size=batch_size, image_size=image_size, network_capacity=network_capacity)
    model_mouth.load_config()
    model_mouth.GAN.load_state_dict(torch.load('%s/%s/model_%d.pt'%(models_dir, name_mouth, load_from)))

    name_tail='short_bird_creative_sequential_r6_partstack_aug_tail_unet2'
    load_from = load_latest(models_dir, name_tail)
    load_from = min(min_step, load_from)
    model_tail = Trainer_cond_unet(name_tail, results_dir,  models_dir, n_part=n_part, batch_size=batch_size, image_size=image_size, network_capacity=network_capacity)
    model_tail.load_config()
    model_tail.GAN.load_state_dict(torch.load('%s/%s/model_%d.pt'%(models_dir, name_tail, load_from)))

    name_selector='short_bird_creative_selector_aug'
    load_from = load_latest(models_dir, name_selector)
    part_selector = Trainer_selector(name_selector, results_dir, models_dir, n_part=n_part, batch_size = batch_size, image_size = image_size, network_capacity = network_capacity)
    part_selector.load_config()
    part_selector.clf.load_state_dict(torch.load('%s/%s/model_%d.pt'%(models_dir, name_selector, load_from)))

    if not os.path.exists(results_dir):
        os.mkdir(results_dir)
    inital_dir = '%s/bird_short_test_init_strokes_%d'%(data_path, image_size)
    dataset = Initialstroke_Dataset(inital_dir, image_size=image_size)
    dataloader = data.DataLoader(dataset, num_workers=5, batch_size=batch_size, drop_last=False, shuffle=False, pin_memory=True)
    # import ipdb;ipdb.set_trace()

    models = [model_eye, model_head, model_body, model_beak, model_legs, model_wings, model_mouth, model_tail]
    target_parts = ['eye', 'head', 'body', 'beak', 'legs', 'wings', 'mouth', 'tail', 'none']
    part_to_id = {'initial': 0, 'eye': 1, 'head': 4, 'body': 3, 'beak': 2, 'legs': 5, 'wings': 8, 'mouth': 6, 'tail': 7}
    max_iter = 10
    if generate_all:
        generation_dir = os.path.join(results_dir, 'DoodlerGAN_all')
        if not os.path.exists(generation_dir):
            os.mkdir(generation_dir)
            os.mkdir(os.path.join(generation_dir, 'bw'))
            os.mkdir(os.path.join(generation_dir, 'color_initial'))
            os.mkdir(os.path.join(generation_dir, 'color'))
        for count, initial_strokes in enumerate(dataloader):
            initial_strokes = initial_strokes.cuda()
            start_point = len(os.listdir(os.path.join(generation_dir, 'bw')))
            print('%d sketches generated'%start_point)
            for i in range(batch_size):
                samples_name = f'generated-{start_point+i}'
                stack_parts = torch.zeros(1, 10, image_size, image_size).cuda()
                initial_strokes_rgb = gs_to_rgb(initial_strokes[i], COLORS['initial'])
                stack_parts[:, 0] = initial_strokes[i, 0]
                stack_parts[:, -1] = initial_strokes[i, 0]
                partial_rgbs = initial_strokes_rgb.clone()
                prev_part = []
                for iter_i in range(max_iter):
                    outputs = part_selector.clf.D(stack_parts)
                    part_rgbs = torch.ones(1, 3, image_size, image_size).cuda()
                    select_part_order = 0
                    select_part_ids = torch.topk(outputs, k=8, dim=0)[1]
                    select_part_id = select_part_ids[select_part_order].item()
                    select_part = target_parts[select_part_id]
                    while (select_part == 'none' and iter_i < 6 or select_part in prev_part):
                        select_part_order += 1
                        if select_part_order > 7:
                            import ipdb;ipdb.set_trace()
                        select_part_id = select_part_ids[select_part_order].item()
                        select_part = target_parts[select_part_id]
                    if select_part == 'none':
                        break
                    prev_part.append(select_part)
                    sketch_rgb = partial_rgbs
                    stack_part = stack_parts.clone()
                    select_model = models[select_part_id]
                    part, partial, part_rgb, partial_rgb = generate_part(select_model.GAN, stack_part, sketch_rgb, COLORS[select_part], select_part, samples_name, 1, results_dir=results_dir, trunc_psi=0.1)
                    stack_parts[0, part_to_id[select_part]] = part[0, 0]
                    partial_rgbs[0] = partial_rgb[0]
                    stack_parts[0, -1] = partial[0, 0]
                    part_rgbs[0] = part_rgb[0]
                initial_colored_full = np.tile(np.max(stack_parts.cpu().data.numpy()[:, 1:-1], 1), [3, 1, 1])
                initial_colored_full = 1-np.max(np.stack([1-initial_strokes_rgb.cpu().data.numpy()[0], initial_colored_full]), 0)
                cv2.imwrite(os.path.join(generation_dir, 'bw', f'{str(samples_name)}.png'), (1-stack_parts[0, -1].cpu().data.numpy())*255.)
                cv2.imwrite(os.path.join(generation_dir, 'color_initial', f'{str(samples_name)}-color.png'), cv2.cvtColor(initial_colored_full.transpose(1, 2, 0)*255., cv2.COLOR_RGB2BGR))
                cv2.imwrite(os.path.join(generation_dir, 'color', f'{str(samples_name)}-color.png'), cv2.cvtColor(partial_rgbs[0].cpu().data.numpy().transpose(1, 2, 0)*255., cv2.COLOR_RGB2BGR))
    else:
        now = datetime.now()
        timestamp = now.strftime("%m-%d-%Y_%H-%M-%S")
        stack_parts = torch.zeros(num_image_tiles*num_image_tiles, 10, image_size, image_size).cuda()
        initial_strokes = dataset.sample(num_image_tiles*num_image_tiles).cuda()
        initial_strokes_rgb = gs_to_rgb(initial_strokes, COLORS['initial'])
        stack_parts[:, 0] = initial_strokes[:, 0]
        stack_parts[:, -1] = initial_strokes[:, 0]
        partial_rgbs = initial_strokes_rgb.clone()
        partial_rgbs_variation = initial_strokes_rgb.clone()
        prev_parts = [[] for _ in range(num_image_tiles**2)]
        samples_name = f'generated-{timestamp}-{min_step}'
        for iter_i in range(max_iter):
            outputs = part_selector.clf.D(stack_parts)
            part_rgbs = torch.ones(num_image_tiles*num_image_tiles, 3, image_size, image_size).cuda()
            for i in range(num_image_tiles**2):
                prev_part = prev_parts[i]
                select_part_order = 0
                select_part_ids = torch.topk(outputs[i], k=9, dim=0)[1]
                select_part_id = select_part_ids[select_part_order].item()
                select_part = target_parts[select_part_id]
                while (select_part == 'none' and iter_i < 6 or select_part in prev_part):
                    select_part_order += 1
                    select_part_id = select_part_ids[select_part_order].item()
                    select_part = target_parts[select_part_id]
                if select_part == 'none':
                    break
                prev_parts[i].append(select_part)
                sketch_rgb = partial_rgbs[i].clone().unsqueeze(0)
                stack_part = stack_parts[i].unsqueeze(0)
                select_model = models[select_part_id]
                part, partial, part_rgb, partial_rgb = generate_part(select_model.GAN, stack_part, sketch_rgb, COLORS[select_part], select_part, samples_name, 1, results_dir=results_dir, trunc_psi=0.1)
                stack_parts[i, part_to_id[select_part]] = part[0, 0]
                stack_parts[i, -1] = partial[0, 0]
                partial_rgbs[i] = partial_rgb[0]
                part_rgbs[i] = part_rgb[0]
            torchvision.utils.save_image(partial_rgbs, os.path.join(results_dir, f'{str(samples_name)}-{str(min_step)}-round{iter_i}.png'), nrow=num_image_tiles)
            torchvision.utils.save_image(part_rgbs, os.path.join(results_dir, f'{str(samples_name)}-{str(min_step)}-part-round{iter_i}.png'), nrow=num_image_tiles)
        torchvision.utils.save_image(1-stack_parts[:, -1:], os.path.join(results_dir, f'{str(samples_name)}-{str(min_step)}-final_pred.png'), nrow=num_image_tiles)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--data_dir", type=str, default='../data')
    parser.add_argument("--results_dir", type=str, default='../results/creative_bird_generation')
    parser.add_argument("--models_dir", type=str, default='../models')
    parser.add_argument('--n_part', type=int, default=10)
    parser.add_argument('--image_size', type=int, default=64)
    parser.add_argument('--network_capacity', type=int, default=16)
    parser.add_argument('--batch_size', type=int, default=100)
    parser.add_argument('--num_image_tiles', type=int, default=8)
    parser.add_argument('--trunc_psi', type=float, default=1.)
    parser.add_argument('--generate_all', action='store_true')

    args = parser.parse_args()
    print(args)

    train_from_folder(args.data_dir, args.results_dir, args.models_dir, args.n_part, args.image_size, args.network_capacity,
        args.batch_size, args.num_image_tiles, args.trunc_psi, args.generate_all)