from .base import SymbolicRegressionEngine
from .engine_utils import *
import numpy as np
import pandas as pd
import shutil
import tempfile
from typing import Sequence

class PySREngine(SymbolicRegressionEngine):
    def __init__(self, backend_kwargs, wrapper_kwargs):
        self._max_params = 8    # by default 
        if isinstance(wrapper_kwargs, dict) and 'max_params' in wrapper_kwargs:
            self._max_params = wrapper_kwargs['max_params']
        super().__init__(backend_kwargs, wrapper_kwargs)

        self._variable_names = None
        self._pareto_front = None
        self._is_classification = None

    @classmethod
    def _default_backend_kwargs(cls):
        return {
            'unary_operators': ["exp", "log"],
            'model_selection': "best",
            'bumper': True
        }

    def fit(self, X: np.ndarray, y: np.ndarray) -> None:
        from pysr import PySRRegressor
        temp_output_dir = tempfile.mkdtemp(prefix="pysr_")
        self._variable_names = [f"x{i}" for i in range(X.shape[1])]
        # Remap y to −1, +1 if needed. PySR uses -1,1 in classification.
        is_binary = np.unique(y).tolist() in ([0, 1], [0], [1])
        self._is_classification = is_binary
        if is_binary:
            y = 2 * y - 1
        try:
            model = PySRRegressor(
                **self.backend_kwargs,
                tempdir=temp_output_dir,
                output_directory=temp_output_dir,
            )
            model.fit(X, y)  
            hall_of_fame = model.equations_
            self._pareto_front = hall_of_fame.copy()
            for name in ("lambda_format", "julia_expression"):
                if name in self._pareto_front:
                    self._pareto_front = self._pareto_front.drop(columns=name)
        finally:
            # Clean up all files generated in the temp directory
            try:
                shutil.rmtree(temp_output_dir)
            except Exception as e:
                print(f"Error cleaning up PySR temp directory: {e}")
    
    def fit_multi(self, Xs: list, ys: list):
        from pysr import PySRRegressor, TemplateExpressionSpec
        import re
        temp_output_dir = tempfile.mkdtemp(prefix="pysr_")
        # Remap y to −1,+1 if needed. PySR uses -1,1 in classification.
        is_binary = np.unique(ys[0]).tolist() in ([0, 1], [0], [1])
        self._is_classification = is_binary
        if is_binary:
            ys = [2*y - 1 for y in ys]
        num_variables = Xs[0].shape[1]
        x_variable_names = [f"x{i}" for i in range(num_variables)]
        self._variable_names = x_variable_names
        cat_name = "category"
        all_variable_names = x_variable_names + [cat_name]
        param_names = [f"u{i}" for i in range(self._max_params)]
        df = concat_datasets(Xs, ys, x_variable_names, cat_name=cat_name)
        X = df.drop(columns='y').to_numpy()
        y = df['y'].to_numpy()
        num_cats = df[cat_name].nunique()

        # building the template expression
        vars_string = ", ".join(x_variable_names)
        params_string =  ", ".join([f"{p}[{cat_name}]" for p in param_names])
        template = TemplateExpressionSpec(
            expressions=["f"],
            variable_names=all_variable_names,
            parameters={p: num_cats for p in param_names},
            combine=f"f({vars_string}, {params_string})"
        )

        try:
            model = PySRRegressor(
                **self.backend_kwargs,
                expression_spec=template,
                tempdir=temp_output_dir,
                output_directory=temp_output_dir,
            )
            model.fit(X, y)  
            # f = #4 + (#3 * (sin(#1) * -1.4412825)); p1 = [-0.69382644, -1.3876529, -0.34691325]; p2 = [0.099999994, 1.5, -0.5]
            hall_of_fame = model.equations_
            def transform_equation(eq: str) -> str:
                """
                Turn an equation like
                    f = #4 + (#3 * (sin(#1) * -1.44)); ...
                into
                    u2 + (u1*(sin(x1)*-1.44))
                given 2 x-vars: x1, x2.
                """
                eq = eq.split(";")[0].strip()
                eq = eq.split("=")[1].strip()
                n_vars = len(x_variable_names)

                def _repl(m):
                    idx = int(m.group(1))
                    if 1 <= idx <= n_vars:
                        # map #1 → x_variable_names[0], etc.
                        return x_variable_names[idx - 1]
                    # beyond n_vars → u{idx - n_vars}
                    return f"u{idx - n_vars}"

                transformed = re.sub(r"#(\d+)", _repl, eq)
                return transformed
            hall_of_fame['pysr_equation'] = hall_of_fame['equation']
            hall_of_fame['equation'] = hall_of_fame['pysr_equation'].apply(transform_equation)
            self._pareto_front = hall_of_fame.copy()
            for name in ("lambda_format", "julia_expression"):
                if name in self._pareto_front:
                    self._pareto_front = self._pareto_front.drop(columns=name)
        finally:
            # Clean up all files generated in the temp directory
            try:
                shutil.rmtree(temp_output_dir)
            except Exception as e:
                print(f"Error cleaning up PySR temp directory: {e}")

    def get_expressions(self) -> pd.Series:
            return self._pareto_front['equation']

    @property
    def engine_variable_names(self) -> Sequence[str]:
        return self._variable_names
