import torch
import torch.nn as nn
import torch.nn.functional as F

import src.data.config as cfg
import src.data.data as data
import src.train.utils as train_utils
import src.train.batch as batch

import src.evaluate.evaluate as evaluate
import src.evaluate.generate as gen
import src.evaluate.sampler as sampling

import utils.utils as utils

from tensorboardX import SummaryWriter


class Trainer(object):
    def __init__(self, opt, meta, data_loader, model, optimizer):
        self.optimizer = optimizer

        self.model = model

        if opt.trainer == "epoch":
            self.epochs = meta.epochs
        self.data_loader = data_loader
        self.opt = opt

        self.losses = {"dev": {}, "test": {}, "train": {}}

        self.top_score = None

        self.lrs = {}

        self.batch_variables = {
            "data": self.data_loader,
            "model": self.model,
            "split": "train"
        }

        self.do_gen = cfg.do_gen
        self.samplers = {}

    def decide_to_save(self):
        to_save = cfg.save and not cfg.toy

        to_save = to_save or cfg.test_save
        print(cfg.save_strategy)
        if cfg.save_strategy == "best":
            if self.top_score[0] != self.opt.train.dynamic.epoch:
                print("DOING IT RIGHT")
                to_save = False
        return to_save

    def save_model(self, tracked_score):
        lrs = {}
        for i, param_group in enumerate(self.optimizer.param_groups):
            lrs[i] = param_group['lr']
        self.lrs[self.opt.train.dynamic.epoch] = lrs

        to_save = self.decide_to_save()

        if to_save:
            data.save_step(
                self.model, self.data_loader.vocab_encoder,
                self.optimizer, self.opt,
                self.opt.train.dynamic.epoch, self.lrs)

    def log_losses(self, opt, losses):
        if (not cfg.toy and cfg.save) or cfg.test_save:
            data.save_eval_file(opt, losses["train"], "losses", split="train")
            data.save_eval_file(opt, losses['dev'], "losses", split="dev")
            data.save_eval_file(opt, losses['test'], "losses", split="test")

    def set_logger(self):
        if cfg.toy:
            self.logger = SummaryWriter(utils.make_name(
                self.opt, prefix="garbage/logs/", eval_=True, do_epoch=False))
        else:
            self.logger = SummaryWriter(utils.make_name(
                self.opt, prefix="logs/", eval_=True, do_epoch=False))
        print("Logging Tensorboard Files at: {}".format(self.logger.logdir))

    def stop_logger(self):
        self.logger.close()

    def run(self):
        self.set_logger()
        self.count = 0
        for epoch in range(self.epochs):
            self.model.train()
            self.opt.train.dynamic.epoch += 1
            self.epoch()

        self.stop_logger()

    def epoch(self):
        nums = self.reset_losses()

        # Initialize progress bar
        bar = utils.initialize_progress_bar(
            self.data_loader.sequences["train"])

        reset = False

        while not reset:
            loss, nums, reset = self.do_forward_pass(nums)
            self.do_backward_pass(loss)
            self.update_parameters()

            bar.update(self.opt.train.dynamic.bs)
            self.count += 1

            for loss_name in self.losses["train"]:
                self.logger.add_scalar(
                    "train/{}".format(loss_name),
                    loss.item() / self.opt.train.dynamic.bs,
                    self.count)

            if cfg.toy and self.counter(nums) > 300:
                break

        with torch.no_grad():
            self.run_evaluation_cycle()

        self.log_losses(self.opt, self.losses)
        self.update_top_score(self.opt)
        self.save_model(self.get_tracked_score())

        self.data_loader.reset_offsets("train")

    def run_evaluation_cycle(self):
        for split in ["dev", "test"]:
            self.evaluator.validate(
                self.opt.train.dynamic.epoch, split,
                self.losses[split])

            if self.do_gen:
                gen.do_gen_run(
                    self.opt, self.generator, self.opt.train.dynamic.epoch,
                    split, self.losses[split])
            iter_num = self.opt.train.dynamic.epoch

            for loss_name in self.losses[split]:
                self.logger.add_scalar(
                    "{}/{}".format(split, loss_name),
                    self.losses[split][loss_name][iter_num],
                    iter_num)

    def clip_gradients(self):
        if self.opt.train.static.clip:
            torch.nn.utils.clip_grad_norm_(
                self.model.parameters(), self.opt.train.static.clip)

    def do_forward_pass(self, nums):
        token_loss, nums, reset = self.batch(
            self.opt, nums, self.losses["train"],
            self.batch_variables)
        return token_loss, nums, reset

    def do_backward_pass(self, loss):
        loss.backward()

    def update_parameters(self):
        if self.opt.model == "lstm":
                self.clip_gradients()
        self.optimizer.step()
        self.optimizer.zero_grad()

    def reset_losses(self):
        loss_names = set([i.rstrip("maicro").rstrip("_") for
                          i in self.losses["train"].keys()])
        return self.initialize_losses(list(loss_names))


class IteratorTrainer(Trainer):
    def __init__(self, opt, meta, data_loader, model, optimizer):
        super(IteratorTrainer, self).__init__(
            opt, meta, data_loader, model, optimizer)

        self.iters = meta.cycle
        self.total_iters = meta.iterations

    def run(self):
        self.set_logger()

        # Initialize progress bar
        bar = utils.set_progress_bar(self.total_iters)

        for cycle_num in range(int(self.total_iters / self.iters)):
            self.model.train()

            self.cycle(bar, cycle_num)

            with torch.no_grad():
                self.run_evaluation_cycle()

            self.log_losses(self.opt, self.losses)
            self.update_top_score(self.opt)
            self.save_model(self.get_tracked_score())

        self.stop_logger()

    def cycle(self, bar, cycle_num):
        nums = self.reset_losses()
        print(self.losses["train"])

        for i in range(1, self.iters + 1):
            # self.model.zero_grad()

            loss, nums, reset = self.do_forward_pass(nums)
            self.do_backward_pass(loss)

            self.update_parameters()
            # print(loss)
            # print(loss.item())
            self.opt.train.dynamic.epoch += 1

            for loss_name in self.losses["train"]:
                self.logger.add_scalar(
                    "train/{}".format(loss_name),
                    loss.item() / self.opt.train.dynamic.bs,
                    self.opt.train.dynamic.epoch)

            bar.update(1)

            if cfg.toy and i > 10:
                break

            if reset:
                self.data_loader.reset_offsets("train")

