import argparse
import json
import os
import os.path as osp

import torch
import torch.nn as nn
import torchvision.datasets as datasets
import torchvision.transforms as transforms

from datasets.image_folder import ImageFolder
from models.resnet import ResNet
from utils.common import ensure_path, pick_vectors, set_gpu

DIR_PATH = os.path.dirname(os.path.realpath(__file__))


def get_save_path(_path, pred_file_path):
    save_path = os.path.join(DIR_PATH, _path)
    pred_file_name = os.path.basename(pred_file_path)
    save_path = os.path.join(save_path, pred_file_name)

    if not os.path.exists(save_path):
        os.makedirs(save_path)

    return save_path


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--pred")
    parser.add_argument("--train-dir")
    parser.add_argument("--save-path", default="save/resnet-fit")
    parser.add_argument("--gpu", default="0")
    args = parser.parse_args()

    save_path = get_save_path(args.save_path, args.pred)

    pred = torch.load(args.pred)
    test_sets = json.load(
        open(osp.join(DIR_PATH, "materials/imagenet-testsets.json"), "r")
    )
    train_wnids = test_sets["train"]

    train_dir = args.train_dir
    normalize = transforms.Normalize(
        mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
    )

    print("number of train classes {}".format(len(train_wnids)))
    train_dataset = ImageFolder(train_dir, train_wnids, stage="train")
    loader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=64,
        shuffle=True,
        num_workers=4,
        pin_memory=True,
        sampler=None,
    )

    print("total training images : {}".format(len(train_dataset)))

    # get the pred wnids from train wnids
    pred_wnids = pred["wnids"]
    pred_vectors = pred["pred"]
    pred_dic = dict(zip(pred_wnids, pred_vectors))
    pred_vectors = pick_vectors(pred_dic, train_wnids, is_tensor=True)

    model = ResNet("resnet50", 1000)
    sd = model.resnet_base.state_dict()
    sd.update(torch.load(osp.join(DIR_PATH, "materials/resnet50-base.pth")))
    model.resnet_base.load_state_dict(sd)

    fcw = pred_vectors

    # confirm if the pred vectors are 1000
    fcw.size(0) == 1000
    fcw.size(1) == 2049

    model.fc.weight = nn.Parameter(fcw[:, :-1])
    model.fc.bias = nn.Parameter(fcw[:, -1])

    model = model.cuda()
    model.train()

    #
    optimizer = torch.optim.SGD(
        model.resnet_base.parameters(), lr=0.0001, momentum=0.9
    )
    loss_fn = nn.CrossEntropyLoss().cuda()

    keep_ratio = 0.9975
    trlog = {}
    trlog["loss"] = []
    trlog["acc"] = []

    for epoch in range(0, 20):

        ave_loss = None
        ave_acc = None

        for i, (data, label) in enumerate(loader, 1):
            data = data.cuda()
            label = label.cuda()

            logits = model(data)
            loss = loss_fn(logits, label)

            _, pred = torch.max(logits, dim=1)
            acc = torch.eq(pred, label).type(torch.FloatTensor).mean().item()

            if i == 1:
                ave_loss = loss.item()
                ave_acc = acc
            else:
                ave_loss = ave_loss * keep_ratio + loss.item() * (
                    1 - keep_ratio
                )
                ave_acc = ave_acc * keep_ratio + acc * (1 - keep_ratio)

            print(
                "epoch {}, {}/{}, loss={:.4f} ({:.4f}), acc={:.4f} ({:.4f})".format(
                    epoch, i, len(loader), loss.item(), ave_loss, acc, ave_acc
                )
            )

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        trlog["loss"].append(ave_loss)
        trlog["acc"].append(ave_acc)

        torch.save(trlog, osp.join(save_path, "trlog"))

        torch.save(
            model.resnet_base.state_dict(),
            osp.join(save_path, "epoch-{}.pth".format(epoch)),
        )
