from .base import SymbolicRegressionEngine
from .engine_utils import *
import numpy as np
import pandas as pd
import os
import time
from glob import glob
import uuid
from typing import Sequence

def replace_numeric_parameters(expr: str) -> str:
    import re
    """
    Replace numeric constants in expr with placeholders u1, u2, ...
    preserving their sign in the formula
    """
    # Regex to match signed numbers (int, float, sci notation) not in variable names
    num_pattern = re.compile(
        r"""
        (?<![A-Za-z0-9_])   # not part of a variable name before
        ([+-]?)             # optional sign (group 1)
        (                   # numeric literal (group 2)
            (?:\d+\.\d*|\.\d+|\d+)   # int or decimal
            (?:[eE][+-]?\d+)?        # optional scientific notation
        )
        (?![A-Za-z0-9_])    # not part of a variable name after
        """,
        re.VERBOSE
    )

    counter = 0

    def repl(match):
        nonlocal counter
        start = match.start()
        # Skip if immediately preceded by '**' (exponent)
        if start >= 2 and expr[start-2:start] == '**':
            return match.group(0)  # leave unchanged
        sign = match.group(1) or ''
        counter += 1
        return f"{sign}u{counter}"

    new_expr = num_pattern.sub(repl, expr)
    return new_expr

def cleanup(temp_files, max_retries=3, delay=1):
    for file_path in temp_files:
        try:
            if os.path.isfile(file_path):
                os.remove(file_path)
                print(f"Successfully deleted {file_path}")
        except Exception as e:
            print(f"Failed to delete {file_path} (attempt 1/{max_retries}): {e}")
            
            for retry in range(2, max_retries + 1):
                time.sleep(delay) 
                try:
                    os.remove(file_path)
                    print(f"Successfully deleted {file_path} on retry {retry}")
                    break
                except Exception as e:
                    print(f"Failed to delete {file_path} (attempt {retry}/{max_retries}): {e}")
        except Exception as e:
            print(f"Error deleting {file_path}: {e}")

def cleanup(temp_files):
    for file_path in temp_files:
        try:
            if os.path.isfile(file_path):
                os.remove(file_path)
        except Exception as e:
            print(f"Error deleting {file_path}: {e}")

def run_df(df: pd.DataFrame, **params):
    from .RSRM.main import main, DATA_DIR, OUTPUT_DIR
    uid = uuid.uuid4()
    id_str = f"dataset_{uid}"

    train_path = os.path.join(DATA_DIR, f"{id_str}_train.csv")
    test_path = os.path.join(DATA_DIR, f"{id_str}_test.csv")

    df.to_csv(train_path, index=False)
    df.to_csv(test_path, index=False)

    try:
        results_pf_df = main(id_str, **params)
    finally:
        res_files = glob(os.path.join(OUTPUT_DIR, f'{id_str}*'))
        files_to_cleanup = [train_path, test_path] + res_files
        cleanup(files_to_cleanup)

    return results_pf_df


class RSRMEngine(SymbolicRegressionEngine):
    def __init__(self, backend_kwargs, wrapper_kwargs):
        names_in_extra_args = ['base', 'gp', 'mcts', 'msdb']
        extra_args = {'extra_args': {k: backend_kwargs.pop(k)
                       for k in names_in_extra_args if k in backend_kwargs}}
        backend_kwargs.update(extra_args)
        super().__init__(backend_kwargs, wrapper_kwargs={})

    def fit(self, X: np.ndarray, y: np.ndarray) -> None:
        self._variable_names = [f"X{i+1}" for i in range(X.shape[1])]
        df = pd.DataFrame(X, columns=self._variable_names)
        df['target'] = y
        self._pareto_front = run_df(df, **self.backend_kwargs)
    
    def fit_multi(self, Xs, ys):
        self._variable_names = [f"X{i+1}" for i in range(Xs[0].shape[1])]
        cat_name = 'category'
        df = concat_datasets(Xs, ys, self._variable_names, cat_name=cat_name)
        params = self.backend_kwargs.copy()
        params.update({'split': cat_name})
        pareto_front = run_df(df, **params)
        pareto_front['equation'] = pareto_front['equation'].apply(replace_numeric_parameters)
        self._pareto_front = pareto_front

    def get_expressions(self) -> pd.Series:
        pf = self._pareto_front
        return pf['equation']
    
    @property
    def engine_variable_names(self) -> Sequence[str]:
        return self._variable_names

    @classmethod
    def _default_backend_kwargs(cls):
        from .RSRM.main import SCRIPT_DIR
        DEFAULT_JSON_PATH = os.path.join(SCRIPT_DIR, 'config/ecology.json')
        return {
                'threshold': 1e-10,
                'num_test': 1, 
                'json_path': DEFAULT_JSON_PATH,
                'output': '', 
                'fit': [], 
                'split': [], 
                }
