import pandas as pd

import torch

from tqdm import trange, tqdm

from classifiers.standard.classifier import Classifier

class SequenceClassifier(Classifier):

    def train(self, datareader):

        device = torch.device(self.config['device'])

        self.model.train()

        self.model.to(device)

        batch_losses = []

        for _ in trange(self.epochs, int(self.config['epochs']), desc="Epoch"):

            tr_loss = 0

            with tqdm(datareader.get_train(), desc="Iteration") as t:

                for step, batch in enumerate(t):

                    batch = tuple(t.to(device) for t in batch)

                    input_ids, input_mask, segment_ids, label_ids = batch

                    loss =  self.model(input_ids, labels=label_ids)[0]

                    loss = loss / self.config['gradient_accumulation_steps']

                    loss.backward()

                    tr_loss += loss.item()

                    if (step + 1) % self.config['gradient_accumulation_steps'] == 0:

                        # Update the model gradients

                        #torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)

                        self.optimizer.step()

                        self.optimizer.zero_grad()

            # save a checkpoint here

            self.save()

            self.epochs = self.epochs+1

        self.save_batch_losses(pd.DataFrame(batch_losses))