import tensorflow as tf

from lottery_ticket.foundations import model_fc
from lottery_ticket.foundations import paths
from lottery_ticket.foundations import pruning
from lottery_ticket.foundations import save_restore
from lottery_ticket.foundations import trainer
from lottery_ticket.foundations.experiment_base import ExperimentBase
from lottery_ticket.mnist_fc import constants

class Experiment(ExperimentBase):
  def __init__(self, trial):
    self.output_dir = paths.trial(paths.experiment(constants.EXPERIMENT_PATH, 'big_two_layer'), trial)

  def train_once(self, iteration, presets=None, masks=None):
    tf.reset_default_graph()
    sess = tf.Session()
    dataset = dataset_mnist.DatasetMnist(
        constants.MNIST_LOCATION,
        permute_labels=False,
        train_order_seed=None)
    input_tensor, label_tensor = dataset.placeholders
    hyperparameters = {'layers': [(1000, tf.nn.relu), (500, tf.nn.relu), (10, None)]}
    model = model_fc.ModelFc(hyperparameters, input_tensor, label_tensor, presets=presets, masks=masks)
    params = {
        'test_interval': 100,
        'save_summaries': True,
        'save_network': True,
    }
    return trainer.train(
