from .io import preprocess
from .train import train
from .network import mlp


def autoencode(count_matrix,
               kfold=None,
               dimreduce=True,
               reconstruct=True,
               mask=None,
               type='normal',
               activation='relu',
               testset=False,
               learning_rate=1e-2,
               hidden_size=(256,64,256),
               l2_coef=0.,
               hidden_dropout=0.1,
               epochs=200,
               batch_size=32,
               optimizer=None,
               init='glorot_uniform',
               **kwargs):

    x = preprocess(count_matrix, kfold=kfold, mask=mask, testset=testset)

    net = mlp(x['shape'][1],
              hidden_size=hidden_size,
              l2_coef=l2_coef,
              hidden_dropout=hidden_dropout,
              activation=activation,
              init=init,
              masking=(mask is not None),
              loss_type=type)

    losses = train(x, net['model'], net['loss'],
                   learning_rate=learning_rate,
                   epochs=epochs, batch_size=batch_size,
                   optimizer=optimizer, **kwargs)

    net['losses'] = losses

    if dimreduce:
        net['reduced'] = net['encoder'].predict(count_matrix)
        net['reduced_linear'] = net['encoder_linear'].predict(count_matrix)
    if reconstruct:
        net['reconstructed'] = net['model'].predict(count_matrix)

    return net

@@ -1,12 +1,23 @@
+import os%0A%0A
 from .io imp
@@ -30,16 +30,29 @@
 eprocess
+, save_matrix
 %0Afrom .t
@@ -600,16 +600,51 @@
 iform',%0A
+               output_folder=None,%0A
         
@@ -1413,88 +1413,1075 @@
     
-if reconstruct:%0A        net%5B'reconstructed'%5D = net%5B'model'%5D.predict(count_matrix
+    if 'dispersion' in net%5B'extra_models'%5D:%0A            net%5B'dispersion'%5D = net%5B'extra_models'%5D%5B'dispersion'%5D()%0A        if 'pi' in net%5B'extra_models'%5D:%0A            net%5B'pi'%5D = net%5B'extra_models'%5D%5B'pi'%5D.predict(count_matrix)%0A%0A        if output_folder:%0A            save_matrix(net%5B'reduced'%5D, os.path.join(output_folder,%0A                                                     'reduced.tsv'))%0A            save_matrix(net%5B'reduced_linear'%5D, os.path.join(output_folder,%0A                                                     'reduced_linear.tsv'))%0A            if 'dispersion' in net:%0A                save_matrix(net%5B'dispersion'%5D, os.path.join(output_folder,%0A                                                            'dispersion.tsv'))%0A            if 'pi' in net:%0A                save_matrix(net%5B'pi'%5D, os.path.join(output_folder,%0A                                                    'pi.tsv'))%0A%0A    if reconstruct:%0A        net%5B'mean'%5D = net%5B'model'%5D.predict(count_matrix)%0A        if output_folder:%0A            save_matrix(net%5B'mean'%5D, os.path.join(output_folder, 'mean.tsv')
 )%0A%0A 

Save everything to a folder if given by the user