import dataset
import utils
import os
import argparse
import torch.backends.cudnn as cudnn
import torch
import torch.nn.functional as F
import numpy as np
import pandas as pd
import time

parser = argparse.ArgumentParser(description='PET lymphoma classification')

#I/O PARAMS
parser.add_argument('--output', type=str, default='.', help='name of output folder (default: ".")')
parser.add_argument('--splits', type=int, default=None, nargs='+', help='data split range to predict for')

#MODEL PARAMS
parser.add_argument('--chptfolder', type=str, default='results', help='path to trained models (default: "results")')
parser.add_argument('--normalize', action='store_true', default=False, help='normalize images')
parser.add_argument('--nfeat', type=int, default=512, help='number of embedded features (default: 512)')
   
#TRAINING PARAMS
parser.add_argument('--batch_size', type=int, default=200, help='how many images to sample per slide (default: 200)')
parser.add_argument('--workers', default=10, type=int, help='number of data loading workers (default: 10)')

def main():
    # Get user input
    global args
    args = parser.parse_args()               
    # best run per split
    bestrun = pd.read_csv(os.path.join(args.chptfolder,'best_run.csv'))
    # set per split
    if args.splits is not None:
        splits = range(args.splits[0],args.splits[1]+1)
    else:
        splits = range(20)
    t0 = time.time()
    for split in splits:
        # Best checkpoint of this split
        run = bestrun[bestrun.split==split].run.values[0]
        chpnt = os.path.join(args.chptfolder,'checkpoint_split'+str(split)+'_run'+str(run)+'.pth')
        # Get model
        model = utils.model_prediction(chpnt)
        model.cuda()
        cudnn.benchmark = True
        # Test dataset
        _,_,_,dset,_ = dataset.get_datasets_singleview(None,args.normalize,False,split)

        print('  dataset size: {}'.format(len(dset.df)))
        out = predict(dset, model)
        pred_name = 'pred_split'+str(split)+'_run'+str(run)+'.csv'
        out.to_csv(os.path.join(args.output, pred_name), index=False)
        print('  loop time: {:.0f} min'.format((time.time()-t0)/60))
        t0 = time.time()
        
def predict(dset, model):
    # Set loaders
    loader = torch.utils.data.DataLoader(dset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers)    
    # Evaluate on test data
    probs, feats = test(loader, model)
    out = dset.df.copy().reset_index(drop=True)
    out['probs'] = probs
    feats = pd.DataFrame(feats, columns=['F{}'.format(x) for x in range(1,args.nfeat+1)])
    out = pd.concat([out, feats], axis=1)
    return out

def test(loader, model):
    # Set model in test mode
    model.eval()
    # Initialize probability vector
    probs = torch.FloatTensor(len(loader.dataset)).cuda()
    # Initialize features
    feats = np.empty((len(loader.dataset),args.nfeat))
    # Loop through batches
    with torch.no_grad():
        for i, (input,_) in enumerate(loader):
            ## Copy batch to GPU
            input = input.cuda()
            ## Forward pass
            f, output = model.forward_feat(input)
            output = F.softmax(output, dim=1)
            ## Features
            feats[i*args.batch_size:i*args.batch_size+input.size(0),:] = f.detach().cpu().numpy()
            ## Clone output to output vector
            probs[i*args.batch_size:i*args.batch_size+input.size(0)] = output.detach()[:,1].clone()
    return probs.cpu().numpy(), feats

if __name__ == '__main__':
    main()