from typing import Callable, Iterator, List, Optional, Union, Iterable, TypeVar, Dict, Generator
from transformersrl.store import SampleStore
from transformersrl.types import PPOSample, ActorCriticOutput
from transformers import GenerationMixin, PreTrainedTokenizer, PreTrainedTokenizerFast
from transformers.utils import ModelOutput
import itertools
import torch
import dataclasses
from torchtyping import TensorType
import contextlib
from transformersrl.pad import left_pad_sequence
from torch.nn.utils.rnn import pad_sequence

__all__ = ['generate_ppo_samples']

Tokenizer = Union[PreTrainedTokenizerFast, PreTrainedTokenizer]
T = TypeVar('T')


def chunk_n(iterable: Iterable[T], n: int) -> Iterator[List[T]]:
    """
    chunk_n 把一个iterable转换成 iterable of batched samples。n为最大的batch size
    :param iterable:
    :param n:
    :return:
    """
    it = iter(iterable)
    while True:
        chunk = list(itertools.islice(it, n))
        if not chunk:
            return
        yield chunk


@contextlib.contextmanager
def gpt_pad_left_special_tokenizer(tokenizer: Tokenizer):
    """
    GPT left padding的tokenizer的context manager。
    """
    if tokenizer.verbose:
        tokenizer.verbose = False
        verbose_disabled = True
    else:
        verbose_disabled = False
    # save the original pad token，if original pad token is None, we also disable verbose to supress warning
    pre_pad_token = tokenizer.pad_token
    if verbose_disabled:
        tokenizer.verbose = True
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "left"
    yield tokenizer
    tokenizer.pad_token = pre_pad_token
    tokenizer.padding_side = "right"


@dataclasses.dataclass
class GenerateSample:
    """
    Generation模型生成的样本
    """
    query: str
    response: str
    query_ids: TensorType["query_length", torch.long]
    response_ids: TensorType["response_length", torch.long]


def _generate_response_by_generation_model(
        query: Iterator[str],
        tokenizer: Tokenizer,
        generation_model: GenerationMixin,
        generate_batch_size: int,
        max_query_length: int,
        max_response_length: int,
        generate_kwargs: dict,
) -> Generator[GenerateSample, None, None]:
    """
    用generation model生成response
    """
    for query_text in chunk_n(query, generate_batch_size):
        generate_device = generation_model.device
        cpu = torch.device("cpu")
        with gpt_pad_left_special_tokenizer(tokenizer) as pad_left_tokenizer:
            cpu_query_batch = pad_left_tokenizer(query_text, padding=True, add_special_tokens=False,
                                                 max_length=max_query_length, return_tensors='pt')
        input_ids = cpu_query_batch["input_ids"].to(generate_device)
        attention_mask = cpu_query_batch["attention_mask"].to(generate_device)
        eos_id = tokenizer.eos_token_id
        with torch.no_grad():
            generation_output: ModelOutput = generation_model.generate(
                input_ids=input_ids,
                # max_length is the total length of the query and response
                max_length=max_response_length + input_ids.shape[1],
                attention_mask=attention_mask,
                return_dict_in_generate=True,
                **generate_kwargs,
            )
            sequences_ids = generation_output["sequences"][:, input_ids.shape[1]:].to(cpu)
            num_return_sequences: int = generate_kwargs['num_return_sequences']
            sequences = tokenizer.batch_decode(sequences_ids, skip_special_tokens=True)

            for i, query in enumerate(query_text):
                for j in range(num_return_sequences):
                    # un-padding query_ids and response_ids
                    response_ids = sequences_ids[i * num_return_sequences + j]
                    t = (response_ids == eos_id).nonzero().tolist()
                    if len(t) != 0:
                        response_ids = response_ids[:t[0][0]]

                    query_ids = cpu_query_batch['input_ids'][i]
                    t = (query_ids == eos_id).nonzero().tolist()
                    if len(t) != 0:
                        query_ids = query_ids[t[-1][0] + 1:]

                    yield GenerateSample(
                        query=query,
                        response=sequences[i * num_return_sequences + j],
                        query_ids=query_ids,
                        response_ids=response_ids,
                    )


@dataclasses.dataclass
class ScoredSample:
    """
    评分后的样本
    """
    sample: GenerateSample
    score: float


def _score_generated_sample(sample: Iterator[GenerateSample],
                            scorer: Callable[[List[str], List[str]], List[float]],
                            scorer_batch_size: int) -> Generator[ScoredSample, None, None]:
    """
    评分生成的样本
    """
    for batch in chunk_n(sample, scorer_batch_size):
        scores = scorer([sample.query for sample in batch], [sample.response for sample in batch])
        for sample, score in zip(batch, scores):
            yield ScoredSample(sample, score)


def _call_ref_model(scored_sample: Iterator[ScoredSample],
                    reference_model: Callable[..., ActorCriticOutput],
                    reference_batch_size: int,
                    eos_token_id: int) -> Generator[PPOSample, None, None]:
    """
    用reference model对样本进行评分
    """

    for batch in chunk_n(scored_sample, reference_batch_size):
        query_ids = left_pad_sequence([sample.sample.query_ids for sample in batch], padding_value=eos_token_id)
        response_ids = pad_sequence([sample.sample.response_ids for sample in batch], padding_value=eos_token_id,
                                    batch_first=True)
        max_query_len = query_ids.shape[1]
        input_ids = torch.cat([query_ids, response_ids], dim=1)
        attention_mask = input_ids.not_equal(eos_token_id).long()
        device = reference_model.device
        input_ids = input_ids.to(device)
        attention_mask = attention_mask.to(device)
        with torch.no_grad():
            reference_model.eval()
            output = reference_model(input_ids=input_ids, attention_mask=attention_mask)
            output.logits = output.logits.to("cpu")
            output.values = output.values.to("cpu")

            for i, sample in enumerate(batch):
                reward = torch.zeros(size=sample.sample.response_ids.shape, dtype=torch.float)
                reward[-1] = sample.score

                yield PPOSample(
                    query=sample.sample.query_ids,
                    response=sample.sample.response_ids,
                    logits=output.logits[i][max_query_len: max_query_len + len(sample.sample.response_ids)],
                    values=output.values[i][max_query_len: max_query_len + len(sample.sample.response_ids)],
                    reward=reward,
                )


def _adjust_reward(ppo_samples: Iterator[PPOSample],
                   adjust_reward: Optional[Callable[[List[PPOSample]], None]],
                   adjust_reward_batch_size: int) -> Generator[PPOSample, None, None]:
    """
    调整reward
    """
    if adjust_reward is None:
        for sample in ppo_samples:
            yield sample
    else:
        for batch in chunk_n(ppo_samples, adjust_reward_batch_size):
            adjust_reward(batch)
            for sample in batch:
                yield sample


def generate_ppo_samples(
        query: Iterator[str],
        tokenizer: Tokenizer,
        scorer: Callable[[List[str], List[str]], List[float]],
        generation_model: GenerationMixin,
        reference_model: Callable[..., ActorCriticOutput],
        max_query_length: int,
        max_response_length: int,
        adjust_reward: Optional[Callable[[List[PPOSample]], None]] = None,
        generate_batch_size: int = 8,
        generate_kwargs: Optional[Dict] = None,
        scorer_batch_size: int = 8,
        reference_batch_size: int = 8,
        adjust_reward_batch_size: int = 8,
) -> SampleStore[PPOSample]:
    """
    生成PPO samples for training

    :param query: query iterator
    :param tokenizer: 分词器
    :param scorer: 评价函数
    :param generation_model: 生成模型
    :param reference_model: 参考模型, 用于计算参考模型的 prob, value
    :param max_query_length: query最大长度
    :param max_response_length: response最大长度
    :param adjust_reward: 调整reward的函数, 用于增加 kl散度约束。可为空
    :param generate_batch_size:  生成过程的batch size
    :param generate_kwargs: 生成过程的参数
    :param scorer_batch_size: 评价过程的batch size
    :param reference_batch_size: 参考模型的batch size
    :param adjust_reward_batch_size: 调整reward的batch size
    :return: 生成的PPO样本
    """
    # 设置默认生成配置
    if generate_kwargs is None:
        generate_kwargs = dict()
    if 'num_beams' not in generate_kwargs:
        generate_kwargs['num_beams'] = 4
    if 'num_return_sequences' not in generate_kwargs:
        generate_kwargs['num_return_sequences'] = 4
    if 'do_sample' not in generate_kwargs:
        generate_kwargs['do_sample'] = False
    if 'early_stopping' not in generate_kwargs:
        generate_kwargs['early_stopping'] = True

    store = SampleStore[PPOSample]()

    # 生成PPOSample的过程主要分为4步
    # 1. 使用 generation model生成 response 文本
    # 2. 使用 scorer评价 query + response 文本
    # 3. 调用 reference model计算 prob, value
    # 4. 调整 reward， 例如增加KL散度约束等
    #
    # NOTE: 这四个步骤目前都可以设置独立的batch size，故使用generator做中间过程
    # NOTE: 每一步的输出尽量干净。即使某一步内部计算的时候需要padding，也应该在输出的时候去掉padding
    # NOTE: 因为可能生成很多PPOSample，所以中间结果都放到CPU中
    # NOTE: 如果速度不够，可以考虑多线程的生产者消费者模式

    generated_samples = _generate_response_by_generation_model(
        query=query, tokenizer=tokenizer, generation_model=generation_model,
        max_response_length=max_response_length, max_query_length=max_query_length,
        generate_batch_size=generate_batch_size, generate_kwargs=generate_kwargs)
    scored_samples = _score_generated_sample(generated_samples, scorer, scorer_batch_size)
    ppo_samples = _call_ref_model(scored_samples, reference_model, reference_batch_size, tokenizer.eos_token_id)
    for ppo_sample in _adjust_reward(ppo_samples, adjust_reward, adjust_reward_batch_size):
        store.append(ppo_sample)
    return store
