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, **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,
              masking=(mask is not None),
              loss_type=type)

    model, encoder, decoder, loss, extras = net['model'], net['encoder'], \
                                            net['decoder'], net['loss'], \
                                            net['extra_models']

    losses = train(x, model, loss,
                   learning_rate=learning_rate,
                   epochs=epochs, batch_size=batch_size,
                   optimizer=optimizer, **kwargs)

    ret = {'model':   model,
           'encoder': encoder,
           'decoder': decoder,
           'extra_models': extras,
           'losses':  losses}

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

    return ret

@@ -104,20 +104,50 @@
 rix,
- kfold=None,
+%0A               kfold=None,%0A              
  dim
@@ -158,16 +158,31 @@
 ce=True,
+%0A              
  reconst
@@ -221,23 +221,53 @@
 one,
- type='normal',
+%0A               type='normal',%0A              
  act
@@ -281,16 +281,31 @@
 ='relu',
+%0A              
  testset
@@ -346,16 +346,31 @@
 te=1e-2,
+%0A              
  hidden_
@@ -387,16 +387,31 @@
 64,256),
+%0A              
  l2_coef
@@ -453,20 +453,50 @@
 0.1,
- epochs=200,
+%0A               epochs=200,%0A              
  bat
@@ -537,16 +537,31 @@
 er=None,
+%0A              
  **kwarg
@@ -902,136 +902,36 @@
     
-model, encoder, decoder, loss, extras = net%5B'model'%5D, net%5B'encoder'%5D, %5C%0A                                           
+losses = train(x,
  net%5B'
-decoder
+model
 '%5D, 
@@ -946,110 +946,8 @@
 s'%5D,
- %5C%0A                                            net%5B'extra_models'%5D%0A%0A    losses = train(x, model, loss,
 %0A   
@@ -1107,178 +1107,111 @@
     
-ret = %7B'model':   model,%0A           'encoder': encoder,%0A           'decoder': decoder,%0A           'extra_models': extras,%0A           'losses':  losses%7D%0A%0A    if dimreduce:
+net%5B'losses'%5D = losses%0A%0A    if dimreduce:%0A        net%5B'reduced'%5D = net%5B'encoder'%5D.predict(count_matrix)
 %0A   
@@ -1211,25 +1211,25 @@
 ix)%0A        
-r
+n
 et%5B'reduced'
@@ -1231,20 +1231,41 @@
 uced
-'%5D = encoder
+_linear'%5D = net%5B'encoder_linear'%5D
 .pre
@@ -1311,17 +1311,17 @@
         
-r
+n
 et%5B'reco
@@ -1334,21 +1334,28 @@
 ted'%5D = 
+net%5B'
 model
+'%5D
 .predict
@@ -1381,12 +1381,12 @@
  return 
-r
+n
 et%0A

Produce both linear and nonlinear dim. reduced representations.