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

class PySRTranslator(Translator):
    """
    Translate generic Sucellus arguments into PySR-specific engine args
    and wrapper-only metadata.
    """

    def loss(self, value):
        loss_map = {
            'MAE': 'L1DistLoss()',
            'MSE': 'L2DistLoss()',
            'BinaryCrossEntropy': 'LogitMarginLoss()',
            'ExpLoss': 'ExpLoss()',
        }

        if value in loss_map:
            return TranslationResult(backend_kwargs={'elementwise_loss': loss_map[value]}, wrapper_kwargs={})

        if value == 'PoissonDeviance':
            # Elementwise Poisson deviance (Julia snippet)
            fn = (
                "f(x,y) = begin\n"
                "  c = max(x, 1e-8)\n"
                "  return y * log((y + 1e-8) / c) - (y - c)\n"
                "end"
            )
            return TranslationResult(backend_kwargs={'elementwise_loss': fn}, wrapper_kwargs={})

        warnings.warn(
            f"Unknown loss '{value}'. Known losses are: {list(loss_map.keys()) + ['PoissonDeviance']}. "
            "Defaulting to L2DistLoss().",
            UserWarning
        )
        return TranslationResult(backend_kwargs={'elementwise_loss': 'L2DistLoss()'}, wrapper_kwargs={})

    def num_iterations(self, value):
        return TranslationResult(backend_kwargs={'niterations': value}, wrapper_kwargs={})

    def operators(self, value):
        operator_map = {
            '+': '+',
            '-': '-',
            '*': '*',
            '/': '/',
            '**': '**',
            'exp': 'exp',
            'log': 'log',
            'sqrt': 'sqrt',
            'square': 'square',
            'sin': 'sin',
            'cos': 'cos',
            'atan': 'atan',
            'tanh': 'tanh',
        }

        valid = []
        for op in value:
            if op in operator_map:
                valid.append(operator_map[op])
            else:
                warnings.warn(f"Unknown operator: '{op}'. Skipped. Valid operators are: {list(operator_map.keys())}", UserWarning)

        return TranslationResult(
            backend_kwargs={
                'binary_operators': [op for op in valid if op in ['+', '-', '*', '/', '**']],
                'unary_operators': [op for op in valid if op not in ['+', '-', '*', '/', '**']]
            },
            wrapper_kwargs={}
        )

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

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

    def max_params(self, value):
        # This is a wrapper-only parameter (Sucellus-level), not a PySR argument.
        # Send it to wrapper_meta so engine wrapper can use it.
        return TranslationResult(backend_kwargs={}, wrapper_kwargs={'max_params': value})
