from __future__ import annotations
from typing import Callable
from collections import namedtuple

from jinja2 import Template, Environment, meta

import torch
from torch.optim import Adam
from torch.nn import Module
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torch import nn, cat, stack, is_tensor, tensor, Tensor

from accelerate import Accelerator

from einops import rearrange

from torch_einops_utils import (
    pad_sequence,
    safe_cat,
    masked_mean,
    and_masks
)

from ema_pytorch import EMA

from x_transformers import TransformerWrapper

from discrete_continuous_embed_readout import Readout

# default query / demonstration template for in-context learned distillation targets from teacher for student

DEFAULT_STUDENT_PROMPT_TEMPLATE = """
[Instruction]
You are a helpful assistant

[Query]
{{ question }}

[Response]
"""

DEFAULT_TEACHER_PROMPT_TEMPLATE = """
[Task Instructions] You are a helpful assistant. Please answer the question based on the provided logic.

[Expert Demonstration] Question: {{ question }} Expert Reasoning and Answer: {{ answer }}

[Current Task] Question: {{ question }} Answer:
"""

def get_variables_from_template(template):

    env = Environment()

    parsed_template = env.parse(template)

    return set(meta.find_undeclared_variables(parsed_template))

# helpers

def exists(v):
    return v is not None

def default(v, d):
    return v if exists(v) else d

def maybe_cast_tensor(t):
    return t if is_tensor(t) else tensor(t)

# classes

SDFTOutput = namedtuple('SDFTOutput', ('loss', 'response'))

class SDFT(Module):
    def __init__(
        self,
        model: TransformerWrapper,
        tokenizer_encode: Callable[[list[str]], list[Tensor]],
        student_max_response_len,
        student_prompt_template = DEFAULT_STUDENT_PROMPT_TEMPLATE,
        teacher_update_rate = 0.01,
        teacher_prompt_template = DEFAULT_TEACHER_PROMPT_TEMPLATE,
        num_init_student_response_tokens_mask = 0,  # they mentioned some issue where the student starts repeating stuff in the prompt template, where they alleviate by masking out the loss for first few tokens
        eos_id = None, # if set, will mask out any losses after the first eos token id detected in a given sample
    ):
        super().__init__()

        if isinstance(model, TransformerWrapper):
            model.input_not_include_cache = True

        self.student = model

        self.teacher = EMA(
            model,
            beta = 1. - teacher_update_rate,
            include_online_model = False
        )

        # sampling

        self.student_max_response_len = student_max_response_len

        self.discrete_readout = Readout(dim = 0, num_discrete = 1)

        # collection of prompts to list[Int['seq']]

        self.tokenizer_encode = tokenizer_encode

        # store templates

        assert get_variables_from_template(teacher_prompt_template) == {'question', 'answer'}, 'your template must contain only variables `question` and `answer`, embedded like so - {{ question }} ... {{ answer }}'
        self.teacher_prompt_template = Template(teacher_prompt_template)

        assert get_variables_from_template(student_prompt_template) == {'question'}
        self.student_prompt_template = Template(student_prompt_template)

        # end of string

        self.eos_id = eos_id

        # how many initial response tokens to exclude from reverse kl loss

        self.num_init_student_response_tokens_mask = num_init_student_response_tokens_mask

    def parameters(self):
        return self.student.parameters()

    def update_teacher_ema_(self):
        self.teacher.update()

    def forward(
        self,
        questions: list[str],
        answers: list[str],
        student_logit_sample_kwargs: dict = dict()
    ):
        maybe_eos_id, prefix_mask_len = self.eos_id, self.num_init_student_response_tokens_mask

        encode = self.tokenizer_encode
        assert len(questions) == len(answers)

        student_vars = [{'question': question} for question in questions]
        teacher_vars = [{'question': question, 'answer': answer} for question, answer in zip(questions, answers)]

        # ready the prompts for student and teacher

        student_prompts_str = [self.student_prompt_template.render(questions) for questions in student_vars]
        teacher_prompts_str = [self.teacher_prompt_template.render(question_answers) for question_answers in teacher_vars]

        student_prompt_ids = [maybe_cast_tensor(encode(prompt)) for prompt in student_prompts_str]
        teacher_prompt_ids = [maybe_cast_tensor(encode(prompt)) for prompt in teacher_prompts_str]

        student_prompt_ids, student_seq_start_pos = pad_sequence(student_prompt_ids, return_lens = True, left = True, pad_lens = True)
        teacher_prompt_ids, teacher_seq_start_pos = pad_sequence(teacher_prompt_ids, return_lens = True, left = True, pad_lens = True)

        student_cache = None
        teacher_cache = None

        # accumulate

        student_responses = None
        token_kl_div_losses = None

        for _ in range(self.student_max_response_len):

            # forward for logit of student and teacher

            student_logits, student_cache = self.student(student_prompt_ids, cache = student_cache, seq_start_pos = student_seq_start_pos, return_intermediates = True)

            with torch.no_grad():
                self.teacher.eval()
                teacher_logits, teacher_cache = self.teacher(teacher_prompt_ids, cache = teacher_cache, seq_start_pos = teacher_seq_start_pos, return_intermediates = True)

            student_token_logit = student_logits[:, -1:]
            teacher_token_logit = teacher_logits[:, -1:]

            student_token_probs = student_token_logit.softmax(dim = -1)
            teacher_token_log_probs = teacher_token_logit.log_softmax(dim = -1)

            # privileged self distillation via ICL

            token_kl_div = F.kl_div(
                teacher_token_log_probs,
                student_token_probs,
                reduction = 'none'

            ).sum(dim = -1)

            token_kl_div_losses = safe_cat((token_kl_div_losses, token_kl_div), dim = 1)

            # sample

            sampled_action = self.discrete_readout.sample(student_token_logit, **student_logit_sample_kwargs)

            student_responses = safe_cat((student_responses, sampled_action), dim = 1)

            # break if all eos

            if exists(maybe_eos_id) and (student_responses == maybe_eos_id).any(dim = -1).all():
                break

            # set student and teacher tokens to the next sampled token

            student_prompt_ids = sampled_action
            teacher_prompt_ids = sampled_action

        # handle eos

        eos_mask = None

        if exists(maybe_eos_id):
            eos_mask = (student_responses == maybe_eos_id).cumsum(dim = -1) == 0

            eos_mask = F.pad(eos_mask, (1, -1), value = True)

            student_responses.masked_fill_(~eos_mask, -1)

        # handle masking of first few response tokens

        init_tokens_mask = None

        if prefix_mask_len > 0:
            init_tokens_mask = torch.ones_like(student_responses).bool()
            init_tokens_mask[:, :prefix_mask_len] = False

        # maybe masked mean for losses

        mask = and_masks([eos_mask, init_tokens_mask])

        loss = masked_mean(token_kl_div_losses, mask)

        return SDFTOutput(loss, student_responses)

# trainer

class SDFTTrainer(Module):
    def __init__(
        self,
        model: Module,
        dataset: Dataset,
        tokenizer_encode: Callable,
        batch_size = 4,
        grad_accum_steps = 1,
        learning_rate = 2e-5,
        max_grad_norm = 0.5,
        sdft_kwargs: dict = dict(),
        accelerate_kwargs: dict = dict(),
        optim_klass = Adam,
        optim_kwargs: dict = dict()
    ):
        super().__init__()

        self.accelerator = Accelerator(
            gradient_accumulation_steps = grad_accum_steps,
            **accelerate_kwargs
        )

        self.model = SDFT(
            model,
            tokenizer_encode = tokenizer_encode,
            **sdft_kwargs
        )

        self.optimizer = optim_klass(self.model.parameters(), lr = learning_rate, **optim_kwargs)

        self.dataloader = DataLoader(dataset, batch_size = batch_size, shuffle = True)

        self.model, self.optimizer, self.dataloader = self.accelerator.prepare(
            self.model, self.optimizer, self.dataloader
        )

        self.max_grad_norm = max_grad_norm

    def forward(self):
        self.model.train()

        for questions, answers in self.dataloader:
            with self.accelerator.accumulate(self.model):
                output = self.model(questions, answers)

                self.accelerator.backward(output.loss)

                if exists(self.max_grad_norm) and self.accelerator.sync_gradients:
                    self.accelerator.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)

                self.optimizer.step()
                self.optimizer.zero_grad()

                if self.accelerator.sync_gradients:
                    self.model.update_teacher_ema_()

        return output
