import climate
import lmj.plot as plt
import numpy as np
import seaborn as sns
import theanets

logging = climate.get_logger('rica')

import models

climate.add_arg('--codebook', metavar='FILE', help='save codebook to FILE')
climate.add_arg('--frames', type=int, metavar='T', help='train on sequences of T frames')
climate.add_arg('--overcomplete', type=float, default=2, metavar='K',
                help='learn a Kx overcomplete codebook')


def main(args):
    data = np.load(args.dataset, mmap_mode='r')
    N = data.shape[1]
    T = args.frames
    K = int(N * T * args.overcomplete)

    def batches():
        batch = np.zeros((args.batch_size, N * T), 'f')
        for b in range(args.batch_size):
            o = np.random.randint(len(data) - T - 1)
            batch[b] = data[o:o+T].ravel()
        return [batch]

    net = theanets.Autoencoder([N * T, (K, 'linear'), (N * T, 'tied')])
    net.train(batches,
              monitors={'hid1:out': (-0.1, -0.01, 0.01, 0.1)},
              **models.kwargs_from_args(args))

    D = net.find('hid1', 'w').get_value().T

    R = 6
    C = 18
    dimg = np.zeros((R * (N + 1) - 1, C * (T + 1) - 1), float)
    bimg = np.zeros((R * (N + 1) - 1, C * (T + 1) - 1), float)
    idx = abs(D).max(axis=1).argsort()[::-1]
    for r in range(R):
        for c in range(C):
            o = np.random.randint(len(data) - T - 1)
            dimg[r*(N+1):r*(N+1)+N, c*(T+1):c*(T+1)+T] = data[o:o+T].T
            bimg[r*(N+1):r*(N+1)+N, c*(T+1):c*(T+1)+T] = D[idx[r * C + c]].reshape((T, N)).T

    _, (dax, bax) = plt.subplots(1, 2)
    dax.imshow(dimg, cmap='coolwarm')
    sns.despine(ax=dax, left=True, bottom=True)
    dax.set_xticks([])
    dax.set_yticks([])
    bax.imshow(bimg, cmap='coolwarm')
    sns.despine(ax=bax, left=True, bottom=True)
    bax.set_xticks([])
    bax.set_yticks([])
    plt.show()

    logging.info('%s: saving %s %s', args.codebook, D.shape, D.dtype)
    np.save(args.codebook, D)


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