import copy

from anonymous_demo.functional.config.config_manager import ConfigManager
from anonymous_demo.core.tad.classic.__bert__.models import TADBERT

_tad_config_template = {
    "model": TADBERT,
    "optimizer": "adamw",
    "learning_rate": 0.00002,
    "patience": 99999,
    "pretrained_bert": "microsoft/mdeberta-v3-base",
    "cache_dataset": True,
    "warmup_step": -1,
    "show_metric": False,
    "max_seq_len": 80,
    "dropout": 0,
    "l2reg": 0.000001,
    "num_epoch": 10,
    "batch_size": 16,
    "initializer": "xavier_uniform_",
    "seed": 52,
    "polarities_dim": 3,
    "log_step": 10,
    "evaluate_begin": 0,
    "cross_validate_fold": -1,
    "use_amp": False,
    # split train and test datasets into 5 folds and repeat 3 training
}

_tad_config_base = {
    "model": TADBERT,
    "optimizer": "adamw",
    "learning_rate": 0.00002,
    "pretrained_bert": "microsoft/deberta-v3-base",
    "cache_dataset": True,
    "warmup_step": -1,
    "show_metric": False,
    "max_seq_len": 80,
    "patience": 99999,
    "dropout": 0,
    "l2reg": 0.000001,
    "num_epoch": 10,
    "batch_size": 16,
    "initializer": "xavier_uniform_",
    "seed": 52,
    "polarities_dim": 3,
    "log_step": 10,
    "evaluate_begin": 0,
    "cross_validate_fold": -1
    # split train and test datasets into 5 folds and repeat 3 training
}

_tad_config_english = {
    "model": TADBERT,
    "optimizer": "adamw",
    "learning_rate": 0.00002,
    "patience": 99999,
    "pretrained_bert": "microsoft/deberta-v3-base",
    "cache_dataset": True,
    "warmup_step": -1,
    "show_metric": False,
    "max_seq_len": 80,
    "dropout": 0,
    "l2reg": 0.000001,
    "num_epoch": 10,
    "batch_size": 16,
    "initializer": "xavier_uniform_",
    "seed": 52,
    "polarities_dim": 3,
    "log_step": 10,
    "evaluate_begin": 0,
    "cross_validate_fold": -1
    # split train and test datasets into 5 folds and repeat 3 training
}

_tad_config_multilingual = {
    "model": TADBERT,
    "optimizer": "adamw",
    "learning_rate": 0.00002,
    "patience": 99999,
    "pretrained_bert": "microsoft/mdeberta-v3-base",
    "cache_dataset": True,
    "warmup_step": -1,
    "show_metric": False,
    "max_seq_len": 80,
    "dropout": 0,
    "l2reg": 0.000001,
    "num_epoch": 10,
    "batch_size": 16,
    "initializer": "xavier_uniform_",
    "seed": 52,
    "polarities_dim": 3,
    "log_step": 10,
    "evaluate_begin": 0,
    "cross_validate_fold": -1
    # split train and test datasets into 5 folds and repeat 3 training
}

_tad_config_chinese = {
    "model": TADBERT,
    "optimizer": "adamw",
    "learning_rate": 0.00002,
    "patience": 99999,
    "cache_dataset": True,
    "warmup_step": -1,
    "show_metric": False,
    "pretrained_bert": "bert-base-chinese",
    "max_seq_len": 80,
    "dropout": 0,
    "l2reg": 0.000001,
    "num_epoch": 10,
    "batch_size": 16,
    "initializer": "xavier_uniform_",
    "seed": 52,
    "polarities_dim": 3,
    "log_step": 10,
    "evaluate_begin": 0,
    "cross_validate_fold": -1
    # split train and test datasets into 5 folds and repeat 3 training
}


class TADConfigManager(ConfigManager):
    def __init__(self, args, **kwargs):
        """
        Available Params:  {'model': BERT,
                            'optimizer': "adamw",
                            'learning_rate': 0.00002,
                            'pretrained_bert': "roberta-base",
                            'cache_dataset': True,
                            'warmup_step': -1,
                            'show_metric': False,
                            'max_seq_len': 80,
                            'patience': 99999,
                            'dropout': 0,
                            'l2reg': 0.000001,
                            'num_epoch': 10,
                            'batch_size': 16,
                            'initializer': 'xavier_uniform_',
                            'seed': {52, 25}
                            'embed_dim': 768,
                            'hidden_dim': 768,
                            'polarities_dim': 3,
                            'log_step': 10,
                            'evaluate_begin': 0,
                            'cross_validate_fold': -1 # split train and test datasets into 5 folds and repeat 3 training
                            }
        :param args:
        :param kwargs:
        """
        super().__init__(args, **kwargs)

    @staticmethod
    def set_tad_config(configType: str, newitem: dict):
        if isinstance(newitem, dict):
            if configType == "template":
                _tad_config_template.update(newitem)
            elif configType == "base":
                _tad_config_base.update(newitem)
            elif configType == "english":
                _tad_config_english.update(newitem)
            elif configType == "chinese":
                _tad_config_chinese.update(newitem)
            elif configType == "multilingual":
                _tad_config_multilingual.update(newitem)
            elif configType == "glove":
                _tad_config_glove.update(newitem)
            else:
                raise ValueError(
                    "Wrong value of config type supplied, please use one from following type: template, base, english, chinese, multilingual, glove"
                )
        else:
            raise TypeError(
                "Wrong type of new config item supplied, please use dict e.g.{'NewConfig': NewValue}"
            )

    @staticmethod
    def set_tad_config_template(newitem):
        TADConfigManager.set_tad_config("template", newitem)

    @staticmethod
    def set_tad_config_base(newitem):
        TADConfigManager.set_tad_config("base", newitem)

    @staticmethod
    def set_tad_config_english(newitem):
        TADConfigManager.set_tad_config("english", newitem)

    @staticmethod
    def set_tad_config_chinese(newitem):
        TADConfigManager.set_tad_config("chinese", newitem)

    @staticmethod
    def set_tad_config_multilingual(newitem):
        TADConfigManager.set_tad_config("multilingual", newitem)

    @staticmethod
    def set_tad_config_glove(newitem):
        TADConfigManager.set_tad_config("glove", newitem)

    @staticmethod
    def get_tad_config_template() -> ConfigManager:
        _tad_config_template.update(_tad_config_template)
        return TADConfigManager(copy.deepcopy(_tad_config_template))

    @staticmethod
    def get_tad_config_base() -> ConfigManager:
        _tad_config_template.update(_tad_config_base)
        return TADConfigManager(copy.deepcopy(_tad_config_template))

    @staticmethod
    def get_tad_config_english() -> ConfigManager:
        _tad_config_template.update(_tad_config_english)
        return TADConfigManager(copy.deepcopy(_tad_config_template))

    @staticmethod
    def get_tad_config_chinese() -> ConfigManager:
        _tad_config_template.update(_tad_config_chinese)
        return TADConfigManager(copy.deepcopy(_tad_config_template))

    @staticmethod
    def get_tad_config_multilingual() -> ConfigManager:
        _tad_config_template.update(_tad_config_multilingual)
        return TADConfigManager(copy.deepcopy(_tad_config_template))

    @staticmethod
    def get_tad_config_glove() -> ConfigManager:
        _tad_config_template.update(_tad_config_glove)
        return TADConfigManager(copy.deepcopy(_tad_config_template))