import climate
import numpy as np
import sklearn.metrics
import theanets

import models

climate.add_arg('--codebook', metavar='FILE', help='decode classes from numpy FILE')


def main(args):
    X = np.load(args.dataset, mmap_mode='r')
    cut = int(0.9 * len(X))

    D = np.load(args.codebook, mmap_mode='r')
    print('D:', D.shape)
    K = sklearn.metrics.pairwise.euclidean_distances(
        X, D, (D * D).sum(axis=1), squared=True).argmin(axis=1)

    def batch(Z):
        def create():
            inputs = np.zeros((args.batch_size, args.time_steps, len(D)), 'f')
            outputs = np.zeros((args.batch_size, args.time_steps), 'i')
            for b in range(args.batch_size):
                o = np.random.randint(len(Z) - args.time_steps - 1)
                inputs[b, :, Z[o:o+args.time_steps]] = 1
                outputs[b] = Z[o+1:o+1+args.time_steps]
            return [inputs, outputs]
        return create

    net = models.train(
        args, theanets.recurrent.Classifier, batch(K), batch(K))
    speak = D[list(net.predict_sequence(K[cut:cut+200], 1000))]
    np.save('/tmp/speak.npy', speak)


if __name__ == '__main__':
    climate.call(main)
