import time
import math
import os
import numpy as np
import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.nn.functional as F
import torch.optim
from torch.optim import lr_scheduler
from torch.utils.data import Dataset, DataLoader
import Models
import helpers

import pdb
import sys
import gc 

use_gpu = torch.cuda.is_available()
NUM_SAMPLES = 300

def main():
    feat_dim = eval(sys.argv[1])
    NUM_ACTIVITIES = 10
    outfile = '/home-2/fyu12@jhu.edu/work3/fyu/CATARACT/csv_outs/out_' + str(feat_dim) + '.csv'
    data_dir = '/home-2/fyu12@jhu.edu/work3/fyu/CATARACT/lmdb_files/'
    weights_dir = '/home-2/fyu12@jhu.edu/work3/fyu/CATARACT/weights/C3D/trunc_' + str(feat_dim) + '/' + sys.argv[2]
    print(weights_dir)
    model = Models.C3D_trunc(feat_dim)
    model.load_state_dict(torch.load(weights_dir))
    if use_gpu:
        model.cuda()
   
    train_datasets = list()
    for i in range(NUM_ACTIVITIES):
        act = i+1 
        print('Currently loading dataset for activity ' + str(act) + ' in phase train')
        sys.stdout.flush()
        lmdb_dir = data_dir + '/train/' + str(act)
        train_datasets.append(helpers.ActSegmentDataset(lmdb_dir)) 
    train_dataloaders = list()
    for i in range(NUM_ACTIVITIES):
        train_dataloaders.append(DataLoader(train_datasets[i], batch_size = 1,
                                            shuffle = True, num_workers = 1))
    features = np.zeros((NUM_SAMPLES*10, feat_dim))
    for i in range(len(train_datasets)):
        print('Currently on activity: {:d}'.format(i+1))
        for j in range(NUM_SAMPLES):
            if (j+1)%30 == 0:
                print('\tCurrently on segment: {:d}'.format(j+1))
            sample = next(iter(train_dataloaders[i]))[0]
            sample = Variable(sample, requires_grad=False)
            if use_gpu:
                sample = sample.cuda()
            features[i*NUM_SAMPLES+j:] = model(sample).data.cpu().numpy()
        

    np.savetxt(outfile, features)

if __name__ == "__main__":
    main()
