import torch
import logging
from transformers import AdamW, get_linear_schedule_with_warmup
from .utils import freeze_bert_parameters, freeze_bert_parameters_KCL
from .__init__ import backbones_map

class ModelManager:

    def __init__(self, args, data, logger_name = 'Detection'):
        
        self.logger = logging.getLogger(logger_name)

    def set_optimizer(self, model, num_train_examples, train_batch_size, num_train_epochs, lr, warmup_proportion):
        num_train_optimization_steps = int(num_train_examples / train_batch_size) * num_train_epochs

        param_optimizer = list(model.named_parameters())
        no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
        optimizer_grouped_parameters = [
            {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},
            {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
        ]
        
        optimizer = AdamW(optimizer_grouped_parameters, lr = lr, correct_bias=False)
        num_warmup_steps= int(num_train_examples * num_train_epochs * warmup_proportion / train_batch_size)
        
        scheduler = get_linear_schedule_with_warmup(optimizer,
                                                    num_warmup_steps=num_warmup_steps,
                                                    num_training_steps=num_train_optimization_steps)
        
        return optimizer, scheduler
    
    def set_model(self, args, pattern):
        backbone = backbones_map[args.backbone]
        args.device = self.device = torch.device('cuda:%d' % int(args.gpu_id) if torch.cuda.is_available() else 'cpu')

        if pattern == 'bert' or pattern == 'llama':
            if hasattr(backbone, 'from_pretrained'):
                model = backbone.from_pretrained('bert-base-uncased', args = args)  
            else:
                model = backbone(args)
                
            if args.freeze_backbone_parameters:
                self.logger.info('Freeze all parameters but the last layer for efficiency')
                if args.method == 'KCL':
                    model = freeze_bert_parameters_KCL(model)
                else: 
                    model = freeze_bert_parameters(model)
        model.to(self.device)
        
        return model