import os
from collections import namedtuple

import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.metrics import classification_report
from torch.optim import Adam
from tqdm import tqdm

from data import DataIteratorDistill
from loss import FocalLoss
from model import CNN
from torchtext import data, vocab

from args import get_args, print_args
from config import ConfigBinaryClassification
from config import ConfigBinaryClassificationDistill
from config import ConfigTripleClassification

if __name__ == "__main__":
    args = get_args()
    print_args(args)

    if args.class_num == 2:
        cfg = ConfigBinaryClassificationDistill()
    elif args.class_num == 3:
        cfg = ConfigTripleClassification()
    else:
        raise ValueError("wrong class num")

    device = torch.device("cuda:%d" % args.cuda)
    Data = DataIteratorDistill(config=cfg, train_batchsize=args.batch_size)
    model = torch.load("checkpoints/CNN-29", map_location=device) 

    optimizer = Adam(model.parameters(), lr=args.lr)
    criterion = FocalLoss(classes=args.class_num, device=device).to(device)
    criterion_kv = nn.KLDivLoss().to(device)

    alpha = 0.2
    T = 2
    for epoch in range(args.epoch_num):
        print(epoch)
        for sample in Data.train_iter:
            model.train()
            optimizer.zero_grad()
            output = model(sample.text.permute(1, 0).to(device))
            loss_f = criterion(output, sample.label.to(device))
            output = F.log_softmax(output/T, 1)
            score = torch.cat((sample.pred0.unsqueeze(1).to(device), 
                sample.pred1.unsqueeze(1).to(device)), dim=1)
            score = F.softmax(score/T,1)
            loss_kv = criterion_kv(output, score.to(device)) * T * T
            loss = alpha * loss_f + (1 - alpha) * loss_kv
            #print(loss_f.item(), loss_kv.item())
            loss.backward()
            optimizer.step()

        with torch.no_grad():
            model.eval()
            preds = []
            labels = []
            for sample in Data.valid_iter:
                output = model(sample.text.permute(1, 0).to(device))
                p = output.argmax(1).cpu().tolist()
                l = sample.label.tolist()
                preds += p
                labels += l
            report = classification_report(preds, labels)
            print(report)
            torch.save(model, os.path.join(args.save_dir, args.save_config + str(epoch)))

[((558, 568), 'args.get_args', 'get_args', ([], {}), '()\n', (566, 568), False, 'from args import get_args, print_args\n'), ((573, 589), 'args.print_args', 'print_args', (['args'], {}), '(args)\n', (583, 589), False, 'from args import get_args, print_args\n'), ((810, 845), 'torch.device', 'torch.device', (["('cuda:%d' % args.cuda)"], {}), "('cuda:%d' % args.cuda)\n", (822, 845), False, 'import torch\n'), ((857, 921), 'data.DataIteratorDistill', 'DataIteratorDistill', ([], {'config': 'cfg', 'train_batchsize': 'args.batch_size'}), '(config=cfg, train_batchsize=args.batch_size)\n', (876, 921), False, 'from data import DataIteratorDistill\n'), ((934, 987), 'torch.load', 'torch.load', (['"""checkpoints/CNN-29"""'], {'map_location': 'device'}), "('checkpoints/CNN-29', map_location=device)\n", (944, 987), False, 'import torch\n'), ((633, 668), 'config.ConfigBinaryClassificationDistill', 'ConfigBinaryClassificationDistill', ([], {}), '()\n', (666, 668), False, 'from config import ConfigBinaryClassificationDistill\n'), ((713, 741), 'config.ConfigTripleClassification', 'ConfigTripleClassification', ([], {}), '()\n', (739, 741), False, 'from config import ConfigTripleClassification\n'), ((1059, 1107), 'loss.FocalLoss', 'FocalLoss', ([], {'classes': 'args.class_num', 'device': 'device'}), '(classes=args.class_num, device=device)\n', (1068, 1107), False, 'from loss import FocalLoss\n'), ((1138, 1152), 'torch.nn.KLDivLoss', 'nn.KLDivLoss', ([], {}), '()\n', (1150, 1152), True, 'import torch.nn as nn\n'), ((1501, 1529), 'torch.nn.functional.log_softmax', 'F.log_softmax', (['(output / T)', '(1)'], {}), '(output / T, 1)\n', (1514, 1529), True, 'import torch.nn.functional as F\n'), ((1680, 1703), 'torch.nn.functional.softmax', 'F.softmax', (['(score / T)', '(1)'], {}), '(score / T, 1)\n', (1689, 1703), True, 'import torch.nn.functional as F\n'), ((1949, 1964), 'torch.no_grad', 'torch.no_grad', ([], {}), '()\n', (1962, 1964), False, 'import torch\n'), ((2320, 2356), 'sklearn.metrics.classification_report', 'classification_report', (['preds', 'labels'], {}), '(preds, labels)\n', (2341, 2356), False, 'from sklearn.metrics import classification_report\n')]