class Net(nn.Module):
    def __init__(self, name, freeze_layers=True):
        super(Net, self).__init__()
        model = self.feature_extractor(name)
        
        if freeze_layers:
            for params in model.parameters():
                params.requires_grad = False

        self.basemodel = model
        # these are the fully connected layers
        self.feedforward1 = nn.Linear(1000, 256)
        self.dropout1 = nn.Dropout(0.2)
        self.feedforward2 = nn.Linear(256, 32)
        self.dropout2 = nn.Dropout(0.2)
        self.feedforward3 = nn.Linear(32, 2)
    
    def forward(self, x):
        x = self.basemodel(x)
        x = F.relu(self.dropout1(self.feedforward1(x)))
        x = F.relu(self.dropout2(self.feedforward2(x)))
        x = self.feedforward3(x)
        return x

    def feature_extractor(self, name):
        if name == 'vgg16':
            model = models.vgg16(pretrained=True)
        elif name == 'resnet18':
            model = models.resnet18(pretrained=True)
        elif name == 'resnet50':
            model = models.resnet18(pretrained=True)
        elif name == 'resnet50':
            model = models.alexnet(pretrained=True)
        elif name == 'squeezenet1_0':
            model = models.squeezenet1_0(pretrained=True)
        elif name == 'densenet161':
            model = models.densenet161(pretrained=True)
        else:
            raise ValueError('Should be one of vgg16,resnet50,resnet18, alexnet, squeezenet1_0, densenet161')
        return model



device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)
def train_model(model, criterion, optimizer, num_epochs=25, dataset_sizes=dataset_sizes, name=''):
    start = datetime.now()
    
    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0
    best_loss = np.inf
    best_tpr = 0.0
    best_fpr = np.inf
    save_every = 5
    
    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch+1, num_epochs))
        print('-'*10)
        
        for phase in ['train', 'valid']:
            if phase == 'train':
                model.train()
            else:
                model.eval()
                
            running_loss = 0.0
            running_corrects = 0
            running_tp = 0
            running_fp = 0
            running_fn = 0
            running_tn = 0
            
            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)
                # https://discuss.pytorch.org/t/why-do-we-need-to-set-the-gradients-manually-to-zero-in-pytorch/4903/6
                optimizer.zero_grad()
                # feed forward process to find out the loss
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)



                    # criterion is the loss
                    loss = criterion(outputs, labels)
                    
                    # then back propogate the loss  
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()
                # running loss    
                running_loss += loss.item()*inputs.size(0)
                running_corrects += torch.sum(preds == labels.data).item()
                tp, fp, fn, tn = confusion_matrix(preds, labels)
                running_tp += tp
                running_fp += fp
                running_fn += fn
                running_tn += tn
                
            epoch_loss = running_loss/dataset_sizes[phase]
            epoch_acc = running_corrects/dataset_sizes[phase]
            epoch_tpr = running_tp/(running_tp+running_fn)
            epoch_fpr = running_fp/(running_fp+running_tn)
            
            print('{} Loss: {:.4f} Acc: {:.4f} TPR: {:.4f} FPR: {:.4f}'.format(phase, epoch_loss, epoch_acc, 
                                                                  epoch_tpr,
                                                                  epoch_fpr))
            
            if phase == 'valid' and epoch_loss < best_loss:
                best_loss = epoch_loss
                best_acc = epoch_acc
                best_tpr = epoch_tpr
                best_fpr = epoch_fpr
                best_model_wts = copy.deepcopy(model.state_dict())
       
        if (epoch+1)%save_every == 0:
            model = model.to(torch.device('cpu'))
            state = {'optimizer': optimizer,
                     'optimizer_state_dict': optimizer.state_dict(),
                    'criterion': criterion,
                     'criterion_state_dict': criterion.state_dict(),
                    'model_state_dict': model.state_dict(),
                    'model': model,
                    'epoch': epoch+1}
            torch.save(state, name)
            model = model.to(device)
            print('Model Checkpointed')
        
        print()
        
    time_elapsed = datetime.now() - start

    print('Training complete in '+str(time_elapsed))
    print('Best Valid Loss: {:.4f} Acc: {:.4f} TPR: {:.4f} FPR: {:.4f}'.format(best_loss, best_acc,
                                                                              best_tpr, best_fpr))

    model.load_state_dict(best_model_wts)
    return model


num_epochs=50
#for m in ['vgg16', 'resnet50', 'resnet18', 'alexnet', 'squeezenet1_0', 'densenet161'][4:]:
for m in ['resnet50']:
    for freeze in [False]: #[True, False]:
        print('-'*100)
        print()
        print('Model: '+m+' Frozen Layers: '+str(freeze))
        name='./'+m+str(freeze)+str(50)+'.pth'
        net = Net(m, freeze_layers=freeze)
        # The batch size should be larger than the number of GPUs used.
        net = nn.DataParallel(net)
        net.to(device)
        
        criterion = nn.CrossEntropyLoss()
        parameters = filter(lambda p: p.requires_grad, net.parameters())
        # differential learning rates for different layers
        optimizer = optim.Adam(parameters)
        
        net = train_model(net, criterion, optimizer, num_epochs=num_epochs, name=name)
        
        net = net.to(torch.device('cpu'))
        state = {'optimizer': optimizer,
                 'optimizer_state_dict': optimizer.state_dict(),
                'criterion': criterion,
                 'criterion_state_dict': criterion.state_dict(),
                'model_state_dict': net.state_dict(),
                'model': net,
                'epoch': num_epochs}
#         torch.save(model_wtd.state_dict(), './Checkpoints/resnetdev2.pth')
#         torch.save(state, './'+m+str(freeze)+'_24April_jackets_single_data_added_'+'50.pth')
        torch.save(state, './'+m+str(freeze)+'13feb_organic_non_organic'+'50.pth')
        print('Best Model Checkpointed')
        
        print()
        print('-'*100)