import torch
from tqdm import tqdm
import torch.optim as optim
from src.models.gunet.utils.dataset import GraphData
import pandas as pd


class Trainer:
    def __init__(self, args, net, G_data, save_path, log_path):
        self.args = args
        self.net = net
        self.feat_dim = G_data.feat_dim
        self.fold_idx = G_data.fold_idx
        self.init(args, G_data.train_gs, G_data.test_gs)
        # specify the model save and logging paths
        self.save_path = save_path
        self.log_path = log_path
        self.model_file_name = f'gunet_{args.data}_{args.seed}.pt'
        if torch.cuda.is_available():
            self.net.cuda()
        self.training_log = []

    def init(self, args, train_gs, test_gs):
        print('#train: %d, #test: %d' % (len(train_gs), len(test_gs)))
        train_data = GraphData(train_gs, self.feat_dim)
        test_data = GraphData(test_gs, self.feat_dim)
        self.train_d = train_data.loader(self.args.batch, True)
        self.test_d = test_data.loader(self.args.batch, False)
        self.optimizer = optim.Adam(
            self.net.parameters(), lr=self.args.lr, amsgrad=True,
            weight_decay=0.0008)

    def to_cuda(self, gs):
        if torch.cuda.is_available():
            if type(gs) == list:
                return [g.cuda() for g in gs]
            return gs.cuda()
        return gs

    def run_epoch(self, epoch, data, model, optimizer):
        losses, accs, n_samples = [], [], 0
        for batch in tqdm(data, desc=str(epoch), unit='b'):
            cur_len, gs, hs, ys = batch
            gs, hs, ys = map(self.to_cuda, [gs, hs, ys])
            loss, acc = model(gs, hs, ys)
            losses.append(loss*cur_len)
            accs.append(acc*cur_len)
            n_samples += cur_len
            if optimizer is not None:
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

        avg_loss, avg_acc = sum(losses) / n_samples, sum(accs) / n_samples
        return avg_loss.item(), avg_acc.item()

    def train(self):
        max_acc = 0.0
        train_str = 'Train epoch %d: loss %.5f acc %.5f'
        test_str = 'Test epoch %d: loss %.5f acc %.5f max %.5f'
        # line_str = '%d:\t%.5f\n'
        for e_id in range(self.args.num_epochs):
            self.net.train()
            train_loss, train_acc = self.run_epoch(
                e_id, self.train_d, self.net, self.optimizer)
            print(train_str % (e_id, train_loss, train_acc))

            with torch.no_grad():
                self.net.eval()
                loss, acc = self.run_epoch(e_id, self.test_d, self.net, None)
            max_acc = max(max_acc, acc)
            print(test_str % (e_id, loss, acc, max_acc))

            if acc == max(max_acc, acc) and self.save_path is not None:
                print('Best model')
                torch.save(self.net.state_dict(), f'{self.save_path}/{self.model_file_name}')
                # torch.save(self.net.state_dict(), f'{self.save_dir}/model.pickle')
            self.training_log.append([e_id, train_loss, loss, train_acc, acc])

        if self.log_path is not None:
            logs = pd.DataFrame(self.training_log, columns=['epoch', 'train_loss', 'valid_loss', 'train_acc', 'valid_acc'])
            logs.to_csv(f'{self.log_path}/gunet_{self.args.data}_{self.args.seed}.csv')
        # with open(self.args.acc_file, 'a+') as f:
        #     f.write(line_str % (self.fold_idx, max_acc))
