from collections import defaultdict
from functools import partial
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
from transformers import (
    Trainer,
    PreTrainedModel,
    PreTrainedTokenizerBase,
    EvalPrediction,
    TrainerCallback,
    TrainingArguments
)
from transformers.utils import is_sagemaker_mp_enabled
import torch
import torch.nn as nn
from torch.utils.data import Dataset
import src.custom_utils as custom_utils
from src.custom_args import ModelArguments
import torch.distributed as dist
import torch.nn.functional as F
from src.train.sampling_utils import MbeirDatataskBalanceSampler
from src.train.train_utils import compute_metrics_retrieval

logger = custom_utils.get_logger(__name__)


MODEL_OUTPUT_KEYS_TO_LOG = ['r_at_1', 'logit_scale', 'logit_pos', 'logit_neg', 'loss_logit_pos']
NO_DECAY_PARAMETERS = {
    'token_embedding',
    'position_embedding',
    'patch_embedding',
    'class_embedding',
    'query_hidden_states',
    'query_state_ids',
    'logit_scale'
}

class CustomTrainer(Trainer):
    def __init__(
        self,
        model: Union[PreTrainedModel, nn.Module] = None,
        args: TrainingArguments = None,
        data_collator: Optional[Callable] = None,
        train_dataset: Optional[Dataset] = None,
        eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None,
        tokenizer: Optional[PreTrainedTokenizerBase] = None,
        model_init: Optional[Callable[[], PreTrainedModel]] = None,
        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
        callbacks: Optional[List[TrainerCallback]] = None,
        optimizers: Tuple[torch.optim.Optimizer,
                          torch.optim.lr_scheduler.LambdaLR] = (None, None),
        preprocess_logits_for_metrics: Optional[Callable[[
            torch.Tensor, torch.Tensor], torch.Tensor]] = None,
        model_args: Optional[ModelArguments] = None,
    ):
        super().__init__(
            model=model,
            args=args,
            data_collator=data_collator,
            train_dataset=train_dataset,
            eval_dataset=eval_dataset,
            tokenizer=tokenizer,
            model_init=model_init,
            compute_metrics=compute_metrics,
            callbacks=callbacks,
            optimizers=optimizers,
            preprocess_logits_for_metrics=preprocess_logits_for_metrics
        )

        self.train_logs: Dict[str, List] = defaultdict(list)
        self.all_query_features = []
        self.all_passage_features = []
        self.all_eval_labels = []

        self.model_args = model_args
        
        self._static_graph_set = False

    def create_optimizer(self):
        """
        Setup the optimizer.

        We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the
        Trainer's init through `optimizers`, or subclass and override this method in a subclass.
        """
        opt_model = self.model_wrapped if is_sagemaker_mp_enabled() else self.model
        opt_model.train()

        if self.optimizer is None:
            optimizer_cls, optimizer_kwargs = Trainer.get_optimizer_cls_and_kwargs(
                self.args)
            decay_parameters = self.get_decay_parameter_names(opt_model)
            decay_parameters = [p for p in decay_parameters if not {x for x in NO_DECAY_PARAMETERS if x in p}]

            optimizer_grouped_parameters = [
                {
                    "params": [
                        p for n, p in opt_model.named_parameters() if (n in decay_parameters and p.requires_grad)
                    ],
                    "weight_decay": self.args.weight_decay,
                    "is_backbone": [
                        1 if ('_model.text_model.' in n or '_model.vision_model.' in n) else 0 for n, p in opt_model.named_parameters() if (n in decay_parameters and p.requires_grad)
                    ]
                },
                {
                    "params": [
                        p for n, p in opt_model.named_parameters() if (n not in decay_parameters and p.requires_grad)
                    ],
                    "weight_decay": 0.0,
                    "is_backbone": [
                        1 if ('_model.text_model.' in n or '_model.vision_model.' in n) else 0 for n, p in opt_model.named_parameters() if (n not in decay_parameters and p.requires_grad)
                    ]
                },
            ]

            if self.args.backbone_learning_rate_multiplier is not None:
                alpha = self.args.backbone_learning_rate_multiplier
                new_optimizer_grouped_parameters = []
                lr = optimizer_kwargs.pop('lr')
                for i, grp in enumerate(optimizer_grouped_parameters):
                    params, back_params = [], []
                    for j, p in enumerate(grp['params']):
                        if grp['is_backbone'][j]:
                            back_params.append(p)
                        else:
                            params.append(p)
                    new_optimizer_grouped_parameters.append(dict(
                        params=params,
                        weight_decay=grp['weight_decay'],
                        lr=lr
                    ))
                    new_optimizer_grouped_parameters.append(dict(
                        params=back_params,
                        weight_decay=grp['weight_decay'],
                        lr=lr * alpha
                    ))
                optimizer_grouped_parameters = new_optimizer_grouped_parameters
            else:
                optimizer_grouped_parameters = [{k: v for k, v in grp.items() if k != 'is_backbone'} for grp in optimizer_grouped_parameters]

            self.optimizer = optimizer_cls(
                optimizer_grouped_parameters, **optimizer_kwargs)
            if optimizer_cls.__name__ == "Adam8bit":
                raise NotImplementedError()

        if is_sagemaker_mp_enabled():
            # self.optimizer = smp.DistributedOptimizer(self.optimizer)
            raise ValueError()

        return self.optimizer

    def parse_model_outputs_for_logs(self, model_outputs):
        for k in MODEL_OUTPUT_KEYS_TO_LOG:
            if k in model_outputs and model_outputs[k] is not None:
                self.train_logs[k].append(model_outputs[k])

    def log(self, logs: Dict[str, float], start_time: Optional[bool]=None) -> None:
        """
        Log `logs` on the various objects watching training.

        Subclass and override this method to inject custom behavior.

        Args:
            logs (`Dict[str, float]`):
                The values to log.
        """
        if self.state.epoch is not None:
            logs["epoch"] = round(self.state.epoch, 2)
        if self.args.include_num_input_tokens_seen:
            logs["num_input_tokens_seen"] = self.state.num_input_tokens_seen

        if not any('eval_' in k for k in logs):
            reduced_train_logs = {
                k: sum(v) / len(v) for k, v in self.train_logs.items() if len(v) > 0}
            logs.update(reduced_train_logs)
        self.train_logs = defaultdict(list)

        output = {**logs, **{"step": self.state.global_step}}
        self.state.log_history.append(output)
        self.control = self.callback_handler.on_log(
            self.args, self.state, self.control, logs)

    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        if self.args.gradient_checkpointing and not self._static_graph_set and hasattr(model, '_set_static_graph'):
            model._set_static_graph()
            self._static_graph_set = True

        loss, outputs = super().compute_loss(model, inputs, return_outputs=True)
        if model.training:
            self.parse_model_outputs_for_logs(outputs)
        if return_outputs:
            return loss, outputs
        else:
            return loss

    def evaluate(
        self,
        eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None,
        ignore_keys: Optional[List[str]] = None,
        metric_key_prefix: str = "eval",
    ) -> Dict[str, float]:
        self.compute_metrics = partial(
            compute_metrics_retrieval,
            query_features=self.all_query_features,
            passage_features=self.all_passage_features,
            eval_labels=self.all_eval_labels,
            batch_size=self.args.eval_similarity_matrix_batch_size
        )
        ret = super().evaluate(
            eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix)
        self.all_query_features = []
        self.all_passage_features = []
        self.all_eval_labels = []
        dist.barrier()
        return ret

    def prediction_step(
        self,
        model: nn.Module,
        inputs: Dict[str, Union[torch.Tensor, Any]],
        prediction_loss_only: bool,
        ignore_keys: Optional[List[str]] = None,
    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
        """
        Return:
            Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: A tuple with the loss,
            logits and labels (each being optional).
        """
        model.eval()
        with torch.inference_mode():
            loss, outputs = self.compute_loss(
                model, inputs, return_outputs=True)
        self.all_query_features.extend(F.normalize(outputs['query_feats'].detach(), p=2, dim=-1).cpu())
        self.all_passage_features.extend(F.normalize(outputs['passage_feats'].detach(), p=2, dim=-1).cpu())
        self.all_eval_labels.extend(inputs['label_ids'])

        # return an empy list for logits and labels to trigger Huggingface's `compute_metrics`
        return (loss, [], [])
    
    def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]:
        if self.args.mbeir_datatask_balance_sampler:
            return MbeirDatataskBalanceSampler(self.train_dataset, self.args)
        else:
            return super()._get_train_sampler()
