import os
from typing import Any, Dict, List, Optional, Union, Callable

import torch
from transformers import BertTokenizer
from transformers import GenerationMixin, LogitsProcessorList, StoppingCriteriaList
from transformers.generation.utils import GenerationConfig, GenerateOutput
from transformers.utils import ModelOutput


class TSGenerationMixin(GenerationMixin):
    tokenizer = BertTokenizer.from_pretrained(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'bert_config'), local_files_only=True)

    @torch.no_grad()
    def generate(
            self,
            inputs: Optional[torch.Tensor] = None,
            text_inputs=None,
            text_input_ids: Optional[torch.Tensor] = None,
            text_attention_mask: Optional[torch.Tensor] = None,
            text_token_type_ids: Optional[torch.Tensor] = None,
            vision_inputs: Optional[torch.Tensor] = None,
            generation_config: Optional[GenerationConfig] = None,
            logits_processor: Optional[LogitsProcessorList] = None,
            stopping_criteria: Optional[StoppingCriteriaList] = None,
            prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
            synced_gpus: Optional[bool] = None,
            assistant_model: Optional["PreTrainedModel"] = None,
            streamer: Optional["BaseStreamer"] = None,
            negative_prompt_ids: Optional[torch.Tensor] = None,
            negative_prompt_attention_mask: Optional[torch.Tensor] = None,
            revin: Optional[bool] = True,
            num_samples: Optional[int] = 1,
            max_output_length: Optional[int] = 96,
            inference_token_len: Optional[int] = None,
            max_text_token_length: Optional[int] = 125,
            **kwargs,
    ) -> Union[GenerateOutput, torch.Tensor]:
        if len(inputs.shape) != 2:
            raise ValueError('Input shape must be: [batch_size, seq_len]')
        if revin:
            means = inputs.mean(dim=-1, keepdim=True)
            stdev = inputs.std(dim=-1, keepdim=True, unbiased=False) + 1e-5
            inputs = (inputs - means) / stdev
        if text_inputs is not None:
            tokenized_text = self._tokenize(text_inputs, max_length=max_text_token_length)
            text_input_ids = tokenized_text['input_ids'].squeeze(0)
            text_attention_mask = tokenized_text['attention_mask'].squeeze(0)
            text_token_type_ids = tokenized_text.get('token_type_ids', torch.zeros_like(text_input_ids)).squeeze(0)

        model_inputs = self.prepare_inputs_for_generation(
            inputs,
            text_input_ids=text_input_ids,
            text_attention_mask=text_attention_mask,
            text_token_type_ids=text_token_type_ids,
            vision_inputs=vision_inputs,
            generation_config=generation_config,
            max_output_length=max_output_length,
            inference_token_len=inference_token_len,
            **kwargs
        )

        outputs = self(**model_inputs, return_dict=True, revin=False, num_samples=num_samples)

        predictions = outputs.logits

        if revin:
            stdev = stdev.unsqueeze(1).repeat(1, num_samples, 1)
            means = means.unsqueeze(1).repeat(1, num_samples, 1)
            predictions = (predictions * stdev) + means

        return predictions

    def prepare_inputs_for_generation(
            self,
            inputs: torch.Tensor,
            text_input_ids: Optional[torch.Tensor] = None,
            text_attention_mask: Optional[torch.Tensor] = None,
            text_token_type_ids: Optional[torch.Tensor] = None,
            vision_inputs: Optional[torch.Tensor] = None,
            generation_config: Optional[GenerationConfig] = None,
            max_output_length: Optional[int] = None,
            inference_token_len: Optional[int] = None,
            **kwargs
    ):
        return {
            "input_ids": inputs,
            "text_input_ids": text_input_ids,
            "text_attention_mask": text_attention_mask,
            "text_token_type_ids": text_token_type_ids,
            "vision_ids": vision_inputs,
            "max_output_length": max_output_length,
            "inference_token_len": inference_token_len,
            **kwargs
        }

    def _tokenize(self, texts, max_length):
        return self.tokenizer(
            texts,
            padding='max_length',
            truncation=True,
            max_length=max_length,
            return_tensors="pt"
        )

    def _update_model_kwargs_for_generation(
            self,
            outputs: ModelOutput,
            model_kwargs: Dict[str, Any],
            horizon_length: int = 1,
            is_encoder_decoder: bool = False,
            standardize_cache_format: bool = False,
    ) -> Dict[str, Any]:
        return model_kwargs