import gcsfs,logging, os, torch
import pandas as pd
from statistics import mean
from tqdm import trange, tqdm

###
# Base class for Bert classifiers.
###
class 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
            batche = []
            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]

                    # Statistics
                    batche.append(loss.item())

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

                    loss.backward()

                    tr_loss += loss.item()

                    if (step + 1) % self.config['gradient_accumulation_steps'] == 0:
                        batch_losses.append(mean(batche))
                        # 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))

    def save_batch_losses(self, losses):
        path = os.path.join(self.config['output_dir'], self.config['experiment_name'])
        if path[:2] != "gs":
            if not os.path.exists(path):
                os.makedirs(path)

        losses.to_csv(os.path.join(self.config['output_dir'], self.config['experiment_name'], "batch_loss.csv"))

    def set_eval_mode(self):
        self.model.eval()

    def load_from_checkpoint(self):

        if 'load_from_checkpoint' in self.config:
            file_path = os.path.join(self.config['output_dir'], "checkpoints", self.config['load_from_checkpoint'])

            checkpoint = torch.load(file_path)
            self.epochs = checkpoint['epoch']
            self.model.load_state_dict(checkpoint['bert_dict'])
            self.optimizer.load_state_dict(checkpoint['optimizer'])

            # work around - for some reason reloading an optimizer that worked with CUDA tensors
            # causes an error - see https://github.com/pytorch/pytorch/issues/2830
            for state in self.optimizer.state.values():
                for k, v in state.items():
                    if isinstance(v, torch.Tensor):
                        if self.config['device'] == 'gpu':
                            state[k] = v.cuda()
                        else:
                            state[k] = v

    def load_object_from_location(self, checkpoint_file):
        if checkpoint_file[:2] != "gs":
            return torch.load(checkpoint_file)
        else:

            fs = gcsfs.GCSFileSystem()
            with fs.open(checkpoint_file, mode='rb') as f:
                return torch.load(f)

    def load_from_checkpoint(self, checkpoint_file):
        file_path = os.path.join(self.config['output_dir'], self.config['experiment_name'],"checkpoints", checkpoint_file)
        checkpoint = self.load_object_from_location(file_path)

        self.epochs = checkpoint['epoch']
        self.model.load_state_dict(checkpoint['bert_dict'])
        self.optimizer.load_state_dict(checkpoint['optimizer'])

        # work around - for some reason reloading an optimizer that worked with CUDA tensors
        # causes an error - see https://github.com/pytorch/pytorch/issues/2830
        for state in self.optimizer.state.values():
            for k, v in state.items():
                if isinstance(v, torch.Tensor):
                    if self.config['device'] == 'gpu':
                        state[k] = v.cuda()
                    else:
                        state[k] = v

    def save_object_to_location(self, object):

        if self.config['output_dir'][:2] != "gs":
            if not os.path.exists(
                    os.path.join(self.config['output_dir'], self.config['experiment_name'], "checkpoints")):
                os.makedirs(os.path.join(self.config['output_dir'], self.config['experiment_name'], "checkpoints"))
            torch.save(object,
                       os.path.join(self.config['output_dir'], self.config['experiment_name'], "checkpoints",
                                    str(self.epochs)))
        else:
            fs = gcsfs.GCSFileSystem()
            file_name = os.path.join(self.config['output_dir'], self.config['experiment_name'], "checkpoints",
                                    str(self.epochs))
            with fs.open(file_name, mode='wb') as f:
                return torch.save(object, f)

    def save(self):
        checkpoint = {
            'epoch': self.epochs + 1,
            'bert_dict': self.model.state_dict(),
            'optimizer': self.optimizer.state_dict(),
        }
        self.save_object_to_location(checkpoint)
        logging.info("Saved model")