
from __future__ import print_function, division

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import  numpy as np
import  torchvision
from torchvision import datasets,models,transforms
import matplotlib.pyplot as plt
import  time
import  os
import  copy
plt.ion()

#load the data
data_transforms={
    'train':transforms.Compose(
        [
            transforms.RandomResizedCrop(224),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])

        ]
    ),
    'val':transforms.Compose(
        [
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]
    )
}

data_dir='data/hymenoptera_data'
image_datasets={x:datasets.ImageFolder(os.path.join(data_dir,x),data_transforms[x])
                for x in ['train','val']}
dataLoaders={x: torch.utils.data.DataLoader(image_datasets[x],batch_size=4,shuffle=True,num_workers=4)
             for x in ['train','val']}

dataset_sizes={x: len(image_datasets[x]) for x in ['train','val']}
class_names=image_datasets['train'].classes
device=torch.device("cuda:0" if torch.cuda.is_available() else 'cpu')


#visualize the images
def imshow(inp,title=None):
    inp=inp.numpy().transpose((1,2,0))
    mean=np.array([0.485, 0.456, 0.406])
    std=np.array([0.229,0.224,0.225])
    inp=inp*std+mean
    inp=np.clip(inp,0,1)
    plt.imshow(inp)
    if title is not None:
        plt.title(title)
    plt.pause(5.001)



def run():
    inputs,classes=next(iter(dataLoaders['train']))
    out=torchvision.utils.make_grid(inputs)
    imshow(out,title=[class_names[x] for x in classes])

def train_model(model,criterion, optimizer, scheduler, num_epochs=25):
    since=time.time()
    best_model_wts=copy.deepcopy(model.state_dict())
    best_acc=0.0

    for epoch in range(num_epochs):
        print(f'epoch {epoch+1}/{num_epochs}\n')
        print('-'*20)

        for phase in ['train','val']:
            if phase=='train':
                model.train()
            else:
                model.eval()

            running_loss=0.0
            running_corrects=0

            #iterate over data
            for inputs, labels in dataLoaders[phase]:
                inputs=inputs.to(device)
                labels=labels.to(device)

                #zero gradients
                optimizer.zero_grad()

                #forward
                with torch.set_grad_enabled(phase=='train'):
                    outputs=model(inputs)
                    _,preds=torch.max(outputs,1)
                    loss=criterion(outputs,labels)

                    if phase=='train':
                        loss.backward()
                        optimizer.step()

                running_loss+=loss.item()*inputs.size(0)
                running_corrects+=torch.sum(preds==labels.data)

            if phase=='train':
                scheduler.step()

            #statistics
            epoch_loss=running_corrects/dataset_sizes[phase]
            epoch_acc=running_corrects.double()/dataset_sizes[phase]
            print(f'phase: {phase} loss:{epoch_loss}  acc:{epoch_acc}\n')

            if phase=='val' and epoch_acc>best_acc:
                best_acc=epoch_acc
                best_model_wts=copy.deepcopy(model.state_dict())

    #training complete
    time_elapsed=time.time()-since
    print(f'training complete in {time_elapsed//60}m {time_elapsed%60}s \n')
    print(f'Best acc: {best_acc}')

    #return the model with best acc
    model.load_state_dict(best_model_wts)
    return  model

def visualize_model(model,num_images=6):
    was_training=model.training
    model.eval()
    images_so_far=0
    fig=plt.figure()

    with torch.no_grad():
        for i, (inputs, labels) in enumerate(dataLoaders['val']):
            inputs=inputs.to(device)
            labels=labels.to(device)
            outputs=model(inputs)
            _,preds=torch.max(outputs,1)

            for j in range(inputs.size()[0]):
                images_so_far+=1
                ax=plt.subplot(num_images//2,2,images_so_far)
                ax.axis('off')
                ax.set_title(f'predicted:{class_names[preds[j]]}')
                imshow(inputs.cpu().data[j])

                if images_so_far==num_images:
                    model.train(mode=was_training)
                    return
        model.train(mode=was_training)



if __name__=='__main__':
    model_ft = models.resnet18(pretrained=True)
    num_ftrs = model_ft.fc.in_features
    model_ft.fc = nn.Linear(num_ftrs, 2)
    model_ft = model_ft.to(device)
    criterion = nn.CrossEntropyLoss()

    # all parameters are being optimized
    optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)

    # Decay LR by a factor of 0.1 every 7 epochs
    exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)
    model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler, num_epochs=25)
    visualize_model(model_ft)