from torch.utils.data import BatchSampler, DataLoader, RandomSampler
from transformers import get_linear_schedule_with_warmup
from lm_seqs_dataset import LmSeqsDataset
from torch.optim import AdamW
from tqdm import tqdm
from torch import nn
from torch.nn import functional
import torch

from typing import List
import math
import os


def produce_average_teacher_logits(teacher_logits):
    return torch.stack(teacher_logits, dim=0).sum(dim=0) / len(teacher_logits)


def produce_averaged_hidden_states(teacher_hidden_states):
    averaged_hidden_states = []
    for hs1, hs2 in zip(teacher_hidden_states[0], teacher_hidden_states[1]):
        averaged_hidden_states.append((hs1 + hs2) / 2.0)

    return tuple(averaged_hidden_states)


class Distiller:
    def __init__(
        self, params, dataset: LmSeqsDataset, token_probs: torch.tensor, student: nn.Module,
            teachers: List[nn.Module]
    ):
        print("Initializing Distiller")
        self.params = params
        self.output_path = params.output_path

        self.student = student
        self.teachers = teachers

        self.student_config = student.config
        self.student_vocab_size = student.config.vocab_size

        sampler = RandomSampler(dataset)
        sampler = BatchSampler(sampler=sampler, batch_size=params.batch_size, drop_last=False)

        self.dataloader = DataLoader(dataset=dataset, batch_sampler=sampler, collate_fn=dataset.batch_sequences)

        self.temperature = params.temperature
        assert self.temperature > 0.0

        self.alpha_ce = params.alpha_ce
        self.alpha_mlm = params.alpha_mlm
        self.alpha_mse = params.alpha_mse
        self.alpha_cos = params.alpha_cos

        self.mlm_mask_prop = params.mlm_mask_prop
        assert 0.0 <= self.mlm_mask_prop <= 1.0
        assert params.word_mask + params.word_keep + params.word_rand == 1.0
        self.pred_probs = torch.FloatTensor([params.word_mask, params.word_keep, params.word_rand])
        self.pred_probs = self.pred_probs.to("cuda")
        self.token_probs = token_probs.to("cuda")

        self.epoch = 0
        self.n_iter = 0
        self.n_total_iter = 0
        self.n_sequences_epoch = 0
        self.total_loss_epoch = 0
        self.last_loss = 0
        self.last_loss_ce = 0
        self.last_loss_mlm = 0
        if self.alpha_mse > 0.0:
            self.last_loss_mse = 0
        if self.alpha_cos > 0.0:
            self.last_loss_cos = 0
        self.last_log = 0

        self.ce_loss_fct = nn.KLDivLoss(reduction="batchmean")
        self.lm_loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
        if self.alpha_mse > 0.0:
            self.mse_loss_fct = nn.MSELoss(reduction="sum")
        if self.alpha_cos > 0.0:
            self.cosine_loss_fct = nn.CosineEmbeddingLoss(reduction="mean")

        print("--- Initializing model optimizer")
        assert params.gradient_accumulation_steps >= 1
        self.num_steps_epoch = len(self.dataloader)
        num_train_optimization_steps = (
            int(self.num_steps_epoch / params.gradient_accumulation_steps * params.n_epoch) + 1
        )

        no_decay = ["bias", "LayerNorm.weight"]
        optimizer_grouped_parameters = [
            {
                "params": [
                    p for n, p in student.named_parameters() if not any(nd in n for nd in no_decay) and p.requires_grad
                ],
                "weight_decay": params.weight_decay,
            },
            {
                "params": [
                    p for n, p in student.named_parameters() if any(nd in n for nd in no_decay) and p.requires_grad
                ],
                "weight_decay": 0.0,
            },
        ]
        print(
            "------ Number of trainable parameters (student): %i"
            % sum([p.numel() for p in self.student.parameters() if p.requires_grad])
        )
        print("------ Number of parameters (student): %i" % sum([p.numel() for p in self.student.parameters()]))
        self.optimizer = AdamW(
            optimizer_grouped_parameters, lr=params.learning_rate, eps=params.adam_epsilon, betas=(0.9, 0.98)
        )

        warmup_steps = math.ceil(num_train_optimization_steps * params.warmup_prop)
        self.scheduler = get_linear_schedule_with_warmup(
            self.optimizer, num_warmup_steps=warmup_steps, num_training_steps=num_train_optimization_steps
        )

    def prepare_batch_mlm(self, batch):
        """
        Prepare the batch: from the token_ids and the lengths, compute the attention mask and the masked label for MLM.

        Input:
        ------
            batch: `Tuple`
                token_ids: `torch.tensor(bs, seq_length)` - The token ids for each of the sequence. It is padded.
                lengths: `torch.tensor(bs)` - The lengths of each of the sequences in the batch.

        Output:
        -------
            token_ids: `torch.tensor(bs, seq_length)` - The token ids after the modifications for MLM.
            attn_mask: `torch.tensor(bs, seq_length)` - The attention mask for the self-attention.
            mlm_labels: `torch.tensor(bs, seq_length)` - The masked language modeling labels. There is a -100 where there is nothing to predict.
        """
        token_ids, lengths = batch
        assert token_ids.size(0) == lengths.size(0)

        attn_mask = torch.arange(token_ids.size(1), dtype=torch.long, device=lengths.device) < lengths[:, None]

        bs, max_seq_len = token_ids.size()
        mlm_labels = token_ids.new(token_ids.size()).copy_(token_ids)

        x_prob = self.token_probs[token_ids.flatten()]
        n_tgt = math.ceil(self.mlm_mask_prop * lengths.sum().item())
        tgt_ids = torch.multinomial(x_prob / x_prob.sum(), n_tgt, replacement=False)
        pred_mask = torch.zeros(
            bs * max_seq_len, dtype=torch.bool, device=token_ids.device
        )  # previously `dtype=torch.uint8`, cf pytorch 1.2.0 compatibility
        pred_mask[tgt_ids] = 1
        pred_mask = pred_mask.view(bs, max_seq_len)

        pred_mask[token_ids == self.params.special_tok_ids["pad_token"]] = 0

        _token_ids_real = token_ids[pred_mask]
        _token_ids_rand = _token_ids_real.clone().random_(self.student_vocab_size)
        _token_ids_mask = _token_ids_real.clone().fill_(self.params.special_tok_ids["mask_token"])
        probs = torch.multinomial(self.pred_probs, len(_token_ids_real), replacement=True)
        _token_ids = (
            _token_ids_mask * (probs == 0).long()
            + _token_ids_real * (probs == 1).long()
            + _token_ids_rand * (probs == 2).long()
        )
        token_ids = token_ids.masked_scatter(pred_mask, _token_ids)

        mlm_labels[~pred_mask] = -100  # previously `mlm_labels[1-pred_mask] = -1`, cf pytorch 1.2.0 compatibility

        # sanity checks
        assert 0 <= token_ids.min() <= token_ids.max() < self.student_vocab_size

        return token_ids, attn_mask, mlm_labels

    def train(self):
        print("Starting training")
        self.student.train()
        for t in self.teachers:
            t.eval()

        for _ in range(self.params.n_epoch):
            print(f"--- Starting epoch {self.epoch}/{self.params.n_epoch-1}")

            progress_bar = tqdm(self.dataloader, desc=f"Epoch {self.epoch}")
            for batch in progress_bar:
                batch = tuple(b.to(f"cuda") for b in batch)

                token_ids, attn_mask, lm_labels = self.prepare_batch_mlm(batch=batch)

                token_ids.to(device='cuda')
                attn_mask.to(device='cuda')
                lm_labels.to(device='cuda')

                self.step(input_ids=token_ids, attention_mask=attn_mask, lm_labels=lm_labels)

                progress_bar.update()
                progress_bar.set_postfix(
                    {"Last_loss": f"{self.last_loss:.2f}", "Avg_cum_loss": f"{self.total_loss_epoch/self.n_iter:.2f}"}
                )
            progress_bar.close()

            print(f"--- Ending epoch {self.epoch}/{self.params.n_epoch-1}")
            self.end_epoch()

        print("Save very last checkpoint as `pytorch_model.bin`.")
        self.save_checkpoint(checkpoint_name="pytorch_model.bin")
        print("Training is finished")

    def step(self, input_ids: torch.tensor, attention_mask: torch.tensor, lm_labels: torch.tensor):
        """
        Input:
        ------
        input_ids: `torch.tensor(bs, seq_length)` - The token ids.
        attention_mask: `torch.tensor(bs, seq_length)` - The attention mask for self attention.
        lm_labels: `torch.tensor(bs, seq_length)` - The language modeling labels (mlm labels for MLM and clm labels for CLM).
        """
        student_outputs = self.student(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True)
        # (bs, seq_length, voc_size)
        with torch.no_grad():
            teacher_outputs = []
            for t in self.teachers:
                teacher_outputs.append(t(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True))
                # (bs, seq_length, voc_size)

        student_logits, student_hidden_states = student_outputs["logits"], student_outputs["hidden_states"]
        teachers_logits = []
        teachers_hidden_states = []
        for outputs in teacher_outputs:
            teachers_logits.append(outputs["logits"])
            teachers_hidden_states.append(outputs["hidden_states"])

        averaged_teacher_logits = produce_average_teacher_logits(teachers_logits)
        averaged_teacher_hidden_states = produce_averaged_hidden_states(teachers_hidden_states)

        assert student_logits.size() == averaged_teacher_logits.size()

        if self.params.restrict_ce_to_mask:
            mask = (lm_labels > -1).unsqueeze(-1).expand_as(student_logits)
            # (bs, seq_length, voc_size)
        else:
            mask = attention_mask.unsqueeze(-1).expand_as(student_logits)
            # (bs, seq_length, voc_size)

        student_logits_slct = torch.masked_select(student_logits, mask)
        # (bs * seq_length * voc_size) modulo the 1s in mask
        student_logits_slct = student_logits_slct.view(-1, student_logits.size(-1))
        # (bs * seq_length, voc_size) modulo the 1s in mask
        averaged_teacher_logits_slct = torch.masked_select(averaged_teacher_logits, mask)
        # (bs * seq_length * voc_size) modulo the 1s in mask
        averaged_teacher_logits_slct = averaged_teacher_logits_slct.view(-1, student_logits.size(-1))
        # (bs * seq_length, voc_size) modulo the 1s in mask
        assert averaged_teacher_logits_slct.size() == student_logits_slct.size()

        loss_ce = (
            self.ce_loss_fct(
                nn.functional.log_softmax(student_logits_slct / self.temperature, dim=-1),
                nn.functional.softmax(averaged_teacher_logits_slct / self.temperature, dim=-1),
            )
            * (self.temperature ** 2)
        )
        loss = self.alpha_ce * loss_ce

        loss_mlm = self.lm_loss_fct(student_logits.view(-1, student_logits.size(-1)), lm_labels.view(-1))
        loss += self.alpha_mlm * loss_mlm

        if self.alpha_mse > 0.0:
            loss_mse = self.mse_loss_fct(student_logits_slct,
                                         averaged_teacher_logits_slct) / student_logits_slct.size(0)
            # Reproducing batchmean reduction
            loss += self.alpha_mse * loss_mse

        if self.alpha_cos > 0.0:
            s_hidden_states = student_hidden_states[-1]  # (bs, seq_length, dim)
            t_hidden_states = averaged_teacher_hidden_states[-1]  # (bs, seq_length, dim)
            mask = attention_mask.unsqueeze(-1).expand_as(s_hidden_states)  # (bs, seq_length, dim)
            assert s_hidden_states.size() == t_hidden_states.size()
            dim = s_hidden_states.size(-1)

            s_hidden_states_slct = torch.masked_select(s_hidden_states, mask)  # (bs * seq_length * dim)
            s_hidden_states_slct = s_hidden_states_slct.view(-1, dim)  # (bs * seq_length, dim)
            t_hidden_states_slct = torch.masked_select(t_hidden_states, mask)  # (bs * seq_length * dim)
            t_hidden_states_slct = t_hidden_states_slct.view(-1, dim)  # (bs * seq_length, dim)

            target = s_hidden_states_slct.new(s_hidden_states_slct.size(0)).fill_(1)  # (bs * seq_length,)
            loss_cos = self.cosine_loss_fct(s_hidden_states_slct, t_hidden_states_slct, target)
            loss += self.alpha_cos * loss_cos

        self.total_loss_epoch += loss.item()
        self.last_loss = loss.item()
        self.last_loss_ce = loss_ce.item()
        if self.alpha_mlm > 0.0:
            self.last_loss_mlm = loss_mlm.item()
        if self.alpha_mse > 0.0:
            self.last_loss_mse = loss_mse.item()
        if self.alpha_cos > 0.0:
            self.last_loss_cos = loss_cos.item()

        self.optimize(loss)

        self.n_sequences_epoch += input_ids.size(0)

    def optimize(self, loss):
        if (loss != loss).data.any():
            print("NaN detected")
            exit()

        if self.params.gradient_accumulation_steps > 1:
            loss = loss / self.params.gradient_accumulation_steps

        loss.backward()

        self.iter()
        if self.n_iter % self.params.gradient_accumulation_steps == 0:
            nn.utils.clip_grad_norm_(self.student.parameters(), self.params.max_grad_norm)
            self.optimizer.step()
            self.optimizer.zero_grad()
            self.scheduler.step()

    def iter(self):
        self.n_iter += 1
        self.n_total_iter += 1

        if self.n_total_iter % self.params.checkpoint_interval == 0:
            self.save_checkpoint()

    def end_epoch(self):
        print(f"{self.n_sequences_epoch} sequences have been trained during this epoch.")

        self.save_checkpoint(checkpoint_name=f"model_epoch_{self.epoch}.pth")

        self.epoch += 1
        self.n_sequences_epoch = 0
        self.n_iter = 0
        self.total_loss_epoch = 0

    def save_checkpoint(self, checkpoint_name: str = "checkpoint.pth"):
        mdl_to_save = self.student.module if hasattr(self.student, "module") else self.student
        mdl_to_save.config.save_pretrained(self.output_path)
        state_dict = mdl_to_save.state_dict()
        torch.save(state_dict, os.path.join(self.output_path, checkpoint_name))
