# -*- coding: utf-8 -*-
"""vit

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1L3_6O4U1Y9r0-tTuL04vPMKKl70Ha_qz
"""

import patchdata
import model
import test
import torch
import torch.optim as optim
import torch.nn as nn
import argparse

# parser함수를 이용하여 Terminal에서 바로 실행할 수 있도록 하였다
# 여러 Hyper Parameter를 터미널에서 실행 시 인자로 받는다
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Vision Transformer')
    parser.add_argument('--img_size', default=32, type=int, help='image size')
    parser.add_argument('--patch_size', default=4, type=int, help='patch size')
    parser.add_argument('--batch_size', default=128, type=int, help='batch size')

    # save 하한선. 아무리 Loss값이 더 낮아졌다고 해도 이 이상의 정확도가 아니면 그냥 저장하지 않도록 한다
    parser.add_argument('--save_acc', default=50, type=int, help='val acc')
    parser.add_argument('--epochs', default=50, type=int, help='training epoch')
    parser.add_argument('--lr', default=2e-3, type=float, help='learning rate')
    parser.add_argument('--drop_rate', default=.1, type=float, help='drop rate')

    # L2 Regularization
    parser.add_argument('--weight_decay', default=0, type=float, help='weight decay')
    parser.add_argument('--num_classes', default=10, type=int, help='number of classes')

    # Latent Vector Dimension : Encoder진입 전 Embedded되어 변화된 차원을 의미한다. 즉 단어의 임베딩 차원
    parser.add_argument('--latent_vec_dim', default=128, type=int, help='latent dimension')

    # Attention head개수
    parser.add_argument('--num_heads', default=8, type=int, help='number of heads')

    # Encoder 몇 번 거치도록 할 것인지
    parser.add_argument('--num_layers', default=12, type=int, help='number of layers in transformer')
    parser.add_argument('--dataname', default='cifar10', type=str, help='data name')
    parser.add_argument('--mode', default='train', type=str, help='train or evaluation')

    # Pretrained model을 선택할 것인지
    parser.add_argument('--pretrained', default=0, type=int, help='pretrained model')
    args = parser.parse_args()
    print(args)

    latent_vec_dim = args.latent_vec_dim

    # Encoder말단에 보면 FC Layer가 2개가 있어서 다음과 같이 연산이 진행된다
    # (N + 1) * D > (N + 1) * D_hid > (N + 1) * D
    # 이 때 은닉층 node수인 D_hid를 의미한다
    mlp_hidden_dim = int(latent_vec_dim/2)

    # patch_size = (H * W) / (P * P)로 정해져 있다
    num_patches = int((args.img_size * args.img_size) / (args.patch_size * args.patch_size))
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    # Image Patches
    # 이미지를 patch로 분해한 후 Flatten하여 Linear Projection(Embedding직전) 단계로 들어가기 직전의 Data로 만든다
    d = patchdata.Flattened2Dpatches(dataname=args.dataname, img_size=args.img_size, patch_size=args.patch_size,
                                     batch_size=args.batch_size)
    trainloader, valloader, testloader = d.patchdata()
    image_patches, _ = iter(trainloader).next()

    # Model을 불러온다
    vit = model.VisionTransformer(patch_vec_size=image_patches.size(2), num_patches=image_patches.size(1),
                                  latent_vec_dim=latent_vec_dim, num_heads=args.num_heads, mlp_hidden_dim=mlp_hidden_dim,
                                  drop_rate=args.drop_rate, num_layers=args.num_layers, num_classes=args.num_classes).to(device)

    # 만약 pre-trained model을 사용한다고 하면
    # 훈련된 weight를 불러온다
    if args.pretrained == 1:
        vit.load_state_dict(torch.load('./model.pth'))

    if args.mode == 'train':
        # Loss and optimizer
        criterion = nn.CrossEntropyLoss() # 이미지 분류이므로 Cross Entropy사용하기
        optimizer = optim.Adam(vit.parameters(), lr=args.lr, weight_decay=args.weight_decay)

        #optimizer = torch.optim.SGD(vit.parameters(), lr=args.lr, momentum=0.9)
        #scheduler = optim.lr_scheduler.OneCycleLR(optimizer, max_lr=args.lr, steps_per_epoch=len(trainloader), epochs=args.epochs)

        # Train
        n = len(trainloader) # len(trainloader)하면 데이터가 총 몇 batch인가가 나온다
        best_acc = args.save_acc
        for epoch in range(args.epochs):
            running_loss = 0
            # 이 for loop가 하나의 epoch이다
            # 한 번 loop돌 때마다 batch_size개의 이미지가 나온다
            for img, labels in trainloader:
                optimizer.zero_grad()
                outputs, _ = vit(img.to(device))
                loss = criterion(outputs, labels.to(device))
                loss.backward()
                optimizer.step()
                running_loss += loss.item() # 한 batch loop의 평균 손실함수값이 반환
                #scheduler.step()

            # running_loss에는 각 batch의 평균 손실함수 값이 더해져있다
            # 여기서 epoch의 평균 손실함수를 구하려면 running_loss / number of batch해야한다
            # 주의! 여기서 number of batch != batch_size임을 주의하기
            train_loss = running_loss / n # 한 epoch동안의 평균 손실값
            val_acc, val_loss = test.accuracy(valloader, vit)
            # if epoch % 5 == 0:
            print('[%d] train loss: %.3f, validation loss: %.3f, validation acc %.2f %%' % (epoch, train_loss, val_loss, val_acc))

            if val_acc > best_acc:
                best_acc = val_acc
                # print('[%d] train loss: %.3f, validation acc %.2f - Save the best model' % (epoch, train_loss, val_acc))
                torch.save(vit.state_dict(), './model.pth')

    else:
        test_acc, test_loss = test.accuracy(testloader, vit)
        print('test loss: %.3f, test acc %.2f %%' % (test_loss, test_acc))