import sys, os
import random
import numpy as np
import torch


def START_SEED(seed=73):
    np.random.seed(seed)
    random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True


##========================== WEIGTHS ===========================================

def load_pretrained(model, weight_path, flexible = False):
    if not weight_path:
        print("No weight file to be loaded returning Model with Random weights")
        return model

    def _purge(key): # hardcoded logic
        return key.replace("backbone.", "")

    model_dict = model.state_dict()
    weight_dict = torch.load(weight_path)

    if 'model' in weight_dict.keys():
        pretrain_dict = weight_dict['model']
    else:
        pretrain_dict = weight_dict

    pretrain_dict = { _purge(k) : v for k, v in pretrain_dict.items()}

    if flexible:
        pretrain_dict = {k: v for k, v in pretrain_dict.items() if k in model_dict}
    if not len(pretrain_dict):
        raise Exception(f"No weight names match to be loaded; though file exits ! {weight_path}, Dict: {weight_dict.keys()}")

    print(f"Pretrained layers:{pretrain_dict.keys()}")

    model_dict.update(pretrain_dict)
    model.load_state_dict(model_dict)

    return model



def count_train_param(model):
    train_params_count = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print('The model has {} trainable parameters'.format(train_params_count))
    return train_params_count


def freeze_params(model, exclusion_list = []):
    ## TODO: Exclusion lists
    for param in model.parameters():
        if param not in exclusion_list:
            param.requires_grad = False
    return model

##==============================================================================


class ObjDict(dict):
    """
    reference: https://stackoverflow.com/a/32107024
    """
    def __init__(self, *args, **kwargs):
        super(ObjDict, self).__init__(*args, **kwargs)
        for arg in args:
            if isinstance(arg, dict):
                for k, v in arg.items():
                    self[k] = v
        if kwargs:
            for k, v in kwargs.items():
                self[k] = v

    def __getattr__(self, attr):
        return self.get(attr)

    def __setattr__(self, key, value):
        self.__setitem__(key, value)

    def __setitem__(self, key, value):
        super(ObjDict, self).__setitem__(key, value)
        self.__dict__.update({key: value})

    def __delattr__(self, item):
        self.__delitem__(item)

    def __delitem__(self, key):
        super(ObjDict, self).__delitem__(key)
        del self.__dict__[key]