from .base import Translator
from .translation import TranslationResult
import warnings

class RSRMTranslator(Translator):
    def loss(self, value):
        losses_list  = ['MAE', 'MSE', 'RMSE', 'R2', 'RobustTanh', 'BinaryCrossEntropy', 'ExpLoss', 'PoissonDeviance']
        if value in losses_list:
            return TranslationResult(backend_kwargs={'base': {'loss': f'{value}'}}, wrapper_kwargs={})
        else: 
            warnings.warn(f"Unknown loss: '{value}'. Valid names are: {losses_list}. Using default, MSE.", UserWarning)

    def num_iterations(self, value):
        return TranslationResult(backend_kwargs={'base': {'epoch': value}}, wrapper_kwargs={})
        
    def operators(self, value):
        operators = {
            '+': 'Add', 
            '-': 'Neg', 
            '*': 'Mul', 
            '/': 'Inv', 
            '**': 'Pow', 
            'exp': 'Exp', 
            'log': 'Log', 
            'sqrt': 'Sqrt', 
            'square': 'N2', 
            'sin': 'Sin', 
            'cos': 'Cos', 
            'atan': 'Atan', 
            'tanh': 'Tanh'
            }
        try:
            tokens = [operators[t] for t in value]
        except:
            tokens = []
            for t in value:
                if t in operators:
                    tokens.append(operators[t])
                else:
                    warnings.warn(f"Unknown operator: '{t}'. Valid names are: {list(operators.keys())}. Skipped.", UserWarning)
        # tokens.append("One")
        return TranslationResult(backend_kwargs={'base': {'tokens': tokens}}, wrapper_kwargs={})

    def max_complexity(self, value):
        return TranslationResult(backend_kwargs={'mcts': {'max_token': value}}, wrapper_kwargs={})

    def max_depth(self, value):
        return TranslationResult(backend_kwargs={'gp': {'max_height': value}, 'mcts': {'max_height': value}}, wrapper_kwargs={})

    def max_params(self, value):
        return TranslationResult(backend_kwargs={'gp': {'max_const': value}, 'mcts': {'max_const': value}}, wrapper_kwargs={})
    
    def force_consistent_sign(self, value):
        return TranslationResult(backend_kwargs={'base': {'force_consistent_sign': value}}, wrapper_kwargs={})
