from omegaconf import DictConfig
from ppd.utils.logger import Log

class LRTable:
    def __init__(self,
                 default_lr: float = 1e-4,
                 prefix_table: DictConfig = None,
                 postfix_table: DictConfig = None,
                 else_table: DictConfig = None):
        # 
        self.default_lr = default_lr
        self.prefix_table = prefix_table
        self.postfix_table = postfix_table
        self.else_table = else_table
        self.tables = [self.prefix_table, self.postfix_table, self.else_table]
        self.tags = ["prefix_", "postfix_", "else_"]
        
    def match_table(self, key) -> bool:
        if self.prefix_table is not None:
            for prefix in self.prefix_table:
                if key.startswith(prefix):
                    return 0, prefix
        
        if self.postfix_table is not None:
            for postfix in self.postfix_table:
                if key.endswith(postfix):
                    return 1, postfix
        
        if self.else_table is not None:
            for else_key in self.else_table:
                if else_key in key and self.tables[2][else_key] is not None:
                    return 2, else_key
        
        return 3, None        
        
    def get_lr(self, key) -> float:
        table_idx, match_key = self.match_table(key)
        if table_idx == 3:
            Log.debug(f'{key} is not matched to any table, use default lr: {self.default_lr}')
            return 'default', self.default_lr
        else:
            Log.debug(f'{key} is matched to table {self.tags[table_idx]}: {match_key}')
            return self.tags[table_idx] + match_key, self.tables[table_idx][match_key]