import argparse
import torch
import time
from torch import nn

parser = argparse.ArgumentParser(description='L1 Linear Layer metrics')
parser.add_argument('--emsize', type=int, default=256,
                    help='size of word embeddings')
parser.add_argument('--vocabsize', type=int, default=35_000,
                    help='vocabulary size')
parser.add_argument('--bsz', type=int, default=80,
                    help='batch size')
parser.add_argument('--seqlen', type=int, default=50,
                    help='sequence length')
parser.add_argument('--jbatch', type=int, default=1,
                    help='batching when iterating over j')
parser.add_argument('--cuda', action='store_true',
                    help='use CUDA')
args = parser.parse_args()

print("instantiating model")
model = nn.LSTM(args.emsize, args.emsize)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

print("instantiating tensors")
I = torch.rand(args.seqlen,args.bsz,args.emsize)
H = (torch.rand(1,args.bsz,args.emsize),torch.rand(1,args.bsz,args.emsize))
B = torch.rand(args.emsize,args.vocabsize)
AB = torch.zeros(args.bsz*args.seqlen, args.vocabsize)

if args.cuda:
    print("moving tensors and model to GPU")
    t = time.time()
    I = I.cuda()
    B = B.cuda()
    H = (H[0].cuda(), H[1].cuda())
    AB = AB.cuda()
    model = model.cuda()
    print("time: ", time.time() - t)

print("model forward pass")
model.train()
A,_ = model(I,H)
A = A.view(args.bsz*args.seqlen, args.emsize)

print("Looping through emsize")
torch.cuda.synchronize()
t =- time.time()
for j in range(0, args.emsize, args.jbatch):
    print(torch.cuda.memory_allocated())
    AB += torch.abs(A[:,j:j+args.jbatch].unsqueeze(-1) - B[j:j+args.jbatch,:]).sum(dim=-2)
torch.cuda.synchronize()
t += time.time()
print("time: ", t)