import cPickle as pickle

import string

import sys

import time

from itertools import izip

import lasagne as nn

import numpy as np

import theano

from datetime import datetime, timedelta

import utils

import logger

import theano.tensor as T

import buffering

from configuration import config, set_configuration

import pathfinder

# theano.config.warn_float64 = 'raise'

if len(sys.argv) < 2:

    sys.exit("Usage: train.py <configuration_name>")

config_name = sys.argv[1]

set_configuration('configs_seg_patch', config_name)

expid = utils.generate_expid(config_name)

print

print "Experiment ID: %s" % expid

print

# metadata

metadata_dir = utils.get_dir_path('models', pathfinder.METADATA_PATH)

metadata_path = metadata_dir + '/%s.pkl' % expid

# logs

logs_dir = utils.get_dir_path('logs', pathfinder.METADATA_PATH)

sys.stdout = logger.Logger(logs_dir + '/%s.log' % expid)

sys.stderr = sys.stdout

print 'Build model'

model = config().build_model()

all_layers = nn.layers.get_all_layers(model.l_out)

all_params = nn.layers.get_all_params(model.l_out)

num_params = nn.layers.count_params(model.l_out)

print '  number of parameters: %d' % num_params

print string.ljust('  layer output shapes:', 36),

print string.ljust('#params:', 10),

print 'output shape:'

for layer in all_layers:

    name = string.ljust(layer.__class__.__name__, 32)

    num_param = sum([np.prod(p.get_value().shape) for p in layer.get_params()])

    num_param = string.ljust(num_param.__str__(), 10)

    print '    %s %s %s' % (name, num_param, layer.output_shape)

train_loss = config().build_objective(model, deterministic=False)

valid_loss = config().build_objective(model, deterministic=True)

learning_rate_schedule = config().learning_rate_schedule

learning_rate = theano.shared(np.float32(learning_rate_schedule[0]))

updates = config().build_updates(train_loss, model, learning_rate)

x_shared = nn.utils.shared_empty(dim=len(model.l_in.shape))

y_shared = nn.utils.shared_empty(dim=len(model.l_target.shape))

idx = T.lscalar('idx')

givens_train = {}

givens_train[model.l_in.input_var] = x_shared[idx * config().batch_size:(idx + 1) * config().batch_size]

givens_train[model.l_target.input_var] = y_shared[idx * config().batch_size:(idx + 1) * config().batch_size]

givens_valid = {}

givens_valid[model.l_in.input_var] = x_shared

givens_valid[model.l_target.input_var] = y_shared

# theano functions

iter_train = theano.function([idx], train_loss, givens=givens_train, updates=updates)

iter_get_predictions = theano.function([idx], nn.layers.get_output(model.l_out), givens=givens_train,

                                       on_unused_input='ignore')

iter_get_targets = theano.function([idx], nn.layers.get_output(model.l_target), givens=givens_train,

                                   on_unused_input='ignore')

iter_get_inputs = theano.function([idx], nn.layers.get_output(model.l_in), givens=givens_train,

                                  on_unused_input='ignore')

iter_validate = theano.function([], valid_loss, givens=givens_valid)

if config().restart_from_save:

    print 'Load model parameters for resuming'

    resume_metadata = utils.load_pkl(config().restart_from_save)

    nn.layers.set_all_param_values(model.l_out, resume_metadata['param_values'])

    start_chunk_idx = resume_metadata['chunks_since_start'] + 1

    chunk_idxs = range(start_chunk_idx, config().max_nchunks)

    lr = np.float32(utils.current_learning_rate(learning_rate_schedule, start_chunk_idx))

    print '  setting learning rate to %.7f' % lr

    learning_rate.set_value(lr)

    losses_eval_train = resume_metadata['losses_eval_train']

    losses_eval_valid = resume_metadata['losses_eval_valid']

else:

    chunk_idxs = range(config().max_nchunks)

    losses_eval_train = []

    losses_eval_valid = []

    start_chunk_idx = 0

train_data_iterator = config().train_data_iterator

valid_data_iterator = config().valid_data_iterator

print

print 'Data'

print 'n train: %d' % train_data_iterator.nsamples

print 'n validation: %d' % valid_data_iterator.nsamples

print 'n chunks per epoch', config().nchunks_per_epoch

print

print 'Train model'

chunk_idx = 0

start_time = time.time()

prev_time = start_time

tmp_losses_train = []

# use buffering.buffered_gen_threaded()

for chunk_idx, (x_chunk_train, y_chunk_train, id_train) in izip(chunk_idxs, buffering.buffered_gen_threaded(

        train_data_iterator.generate())):

    if chunk_idx in learning_rate_schedule:

        lr = np.float32(learning_rate_schedule[chunk_idx])

        print '  setting learning rate to %.7f' % lr

        print

        learning_rate.set_value(lr)

    # load chunk to GPU

    x_shared.set_value(x_chunk_train)

    y_shared.set_value(y_chunk_train)

    # make nbatches_chunk iterations

    chunk_train_losses = []

    for b in xrange(config().nbatches_chunk):

        loss = iter_train(b)

        chunk_train_losses.append(loss)

        tmp_losses_train.append(loss)

    print chunk_idx, np.mean(chunk_train_losses)

    if ((chunk_idx + 1) % config().validate_every) == 0:

        print

        print 'Chunk %d/%d' % (chunk_idx + 1, config().max_nchunks)

        # calculate mean train loss since the last validation phase

        mean_train_loss = np.mean(tmp_losses_train)

        print 'Mean train loss: %7f' % mean_train_loss

        losses_eval_train.append(mean_train_loss)

        tmp_losses_train = []

        # load validation data to GPU

        tmp_losses_valid = []

        for i, (x_chunk_valid, y_chunk_valid, ids_batch) in enumerate(

                buffering.buffered_gen_threaded(valid_data_iterator.generate(),

                                                buffer_size=2)):

            x_shared.set_value(x_chunk_valid)

            y_shared.set_value(y_chunk_valid)

            l_valid = iter_validate()

            print i, l_valid

            tmp_losses_valid.append(l_valid)

        # calculate validation loss across validation set

        valid_loss = np.mean(tmp_losses_valid)

        # TODO: taking mean is not correct if chunks have different sizes!!!

        print 'Validation loss: ', valid_loss

        losses_eval_valid.append(valid_loss)

        now = time.time()

        time_since_start = now - start_time

        time_since_prev = now - prev_time

        prev_time = now

        est_time_left = time_since_start * (config().max_nchunks - chunk_idx + 1.) / (chunk_idx + 1. - start_chunk_idx)

        eta = datetime.now() + timedelta(seconds=est_time_left)

        eta_str = eta.strftime("%c")

        print "  %s since start (%.2f s)" % (utils.hms(time_since_start), time_since_prev)

        print "  estimated %s to go (ETA: %s)" % (utils.hms(est_time_left), eta_str)

        print

    if ((chunk_idx + 1) % config().save_every) == 0:

        print

        print 'Chunk %d/%d' % (chunk_idx + 1, config().max_nchunks)

        print 'Saving metadata, parameters'

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

            pickle.dump({

                'configuration_file': config_name,

                'git_revision_hash': utils.get_git_revision_hash(),

                'experiment_id': expid,

                'chunks_since_start': chunk_idx,

                'losses_eval_train': losses_eval_train,

                'losses_eval_valid': losses_eval_valid,

                'param_values': nn.layers.get_all_param_values(model.l_out)

            }, f, pickle.HIGHEST_PROTOCOL)

            print '  saved to %s' % metadata_path

            print