import itertools
import pandas as pd
import numpy as np
from typing import Sequence, List
import warnings
from .base import SymbolicRegressionEngine

class LinearEngine(SymbolicRegressionEngine):
    def __init__(self, backend_kwargs, wrapper_kwargs):
        """
        loss: one of "MSE" (default), "PoissonDeviance", or "BinaryCrossEntropy"
        params: reserved for future extensions (e.g., regularization)
        """
        loss = backend_kwargs.get("loss", "MSE")
        max_params = backend_kwargs.get("max_params", None)

        valid = {"MSE", "PoissonDeviance", "BinaryCrossEntropy"}
        if loss not in valid:
            raise ValueError(f"Invalid loss for the LinearEngine: '{loss}'. Valid choices: {valid}")
        self.loss = loss
        self.max_params = max_params

    @classmethod
    def _default_backend_kwargs(cls):
        return {}

    def fit(self, X: np.ndarray, y: np.ndarray) -> None:
        import statsmodels.api as sm
        import statsmodels.formula.api as smf
        from statsmodels.tools.sm_exceptions import DomainWarning

        df = pd.DataFrame(X, columns=[f"x{i+1}" for i in range(X.shape[1])])
        df["y"] = y
        features = list(df.columns[:-1])
        num_params = self.max_params or len(features)
        self._feature_names = features
        model_metrics = []

        def evaluate(formula, family):
            model = smf.glm(formula=formula, data=df, family=family).fit()
            aic = model.aic
            loss = -model.llf / len(df)
            return (formula,
                    family.link.__class__.__name__,
                    aic,
                    loss,
                    model.params)

        if self.loss == "PoissonDeviance":
            with warnings.catch_warnings():
                # Filter out Identity link function warnings
                warnings.filterwarnings("ignore", category=DomainWarning)   
                warnings.filterwarnings("ignore", category=RuntimeWarning, message="invalid value encountered in log")
                # Poisson count models, test both identity & log links
                for intercept_flag in [True, False]:
                    N = num_params - int(intercept_flag)
                    for L in range(1, N + 1):
                        for subset in itertools.combinations(features, L):
                            if intercept_flag:
                                formula = "y ~ " + " + ".join(subset)  # with intercept
                            else:
                                formula = "y ~ " + " + ".join(subset) + " - 1"  # no intercept
                            for link in [sm.families.links.Identity(),
                                        sm.families.links.Log()]:
                                fm, lk, aic, loss, params = evaluate(
                                    formula,
                                    sm.families.Poisson(link=link)
                                )
                                model_metrics.append(("count", fm, lk, aic, loss, params))

        elif self.loss == "BinaryCrossEntropy":
            uni = set(np.unique(y))
            if not uni.issubset({0, 1}):
                raise ValueError("y must be 0/1 for binary task")
            for L in range(1, num_params + 1):
                for subset in itertools.combinations(features, L):
                    formula = "y ~ " + " + ".join(subset)
                    fm, lk, aic, loss, params = evaluate(
                        formula,
                        sm.families.Binomial(link=sm.families.links.Logit())
                    )
                    model_metrics.append(("binary", fm, lk, aic, loss, params))

        else:  
            # Ordinary least squares
            for intercept_flag in [True, False]:
                N = num_params - int(intercept_flag)
                for L in range(1, N + 1):
                    for subset in itertools.combinations(features, L):
                        if intercept_flag:
                            formula = "y ~ " + " + ".join(subset)  # with intercept
                        else:
                            formula = "y ~ " + " + ".join(subset) + " - 1"  # no intercept
                        model = smf.ols(formula=formula, data=df).fit()
                        model_metrics.append((
                            "ols",
                            formula,
                            "OLS",
                            model.aic,
                            model.mse_resid,
                            model.params
                        ))

        columns = ["Task", "Formula", "Link/Method", "AIC", "Loss/MSE", "Parameters"]
        self._results = (
            pd.DataFrame(model_metrics, columns=columns)
              .sort_values(by="AIC")
              .reset_index(drop=True)
        )
    
    def fit_multi(self, Xs: List[np.ndarray], ys: List[np.ndarray]) -> None:
        import statsmodels.api as sm
        import statsmodels.formula.api as smf
        from statsmodels.tools.sm_exceptions import DomainWarning

        if len(Xs) != len(ys):
            raise ValueError("Xs and ys must have the same length")

        # Assume all datasets have same feature count
        n_features = Xs[0].shape[1]
        features = [f"x{i+1}" for i in range(n_features)]
        num_params = self.max_params or (n_features + 1)
        self._feature_names = features
        model_metrics = []

        def evaluate_all_datasets(formula, family_or_method, is_glm=True):
            aics = []
            losses = []
            n_params = None

            for X, y in zip(Xs, ys):
                df = pd.DataFrame(X, columns=features)
                df["y"] = y

                if is_glm:
                    model = smf.glm(formula=formula, data=df, family=family_or_method).fit()
                    aics.append(model.aic)
                    losses.append(-model.llf / len(df))
                else:
                    model = smf.ols(formula=formula, data=df).fit()
                    aics.append(model.aic)
                    losses.append(model.mse_resid)

                if n_params is None:
                    n_params = len(model.params)

            # Worst-case metrics
            worst_aic = max(aics)
            worst_loss = max(losses)
            # Placeholder params
            placeholder_params = {f"u{i}": None for i in range(n_params)}

            return worst_aic, worst_loss, placeholder_params

        if self.loss == "PoissonDeviance":
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore", category=DomainWarning)
                warnings.filterwarnings("ignore", category=RuntimeWarning,
                                        message="invalid value encountered in log")
                for intercept_flag in [True, False]:
                    N = min(num_params - int(intercept_flag), n_features)
                    for L in range(1, N + 1):
                        for subset in itertools.combinations(features, L):
                            if intercept_flag:
                                formula = "y ~ " + " + ".join(subset)
                            else:
                                formula = "y ~ " + " + ".join(subset) + " - 1"
                            for link in [sm.families.links.Identity(),
                                        sm.families.links.Log()]:
                                worst_aic, worst_loss, placeholder_params = evaluate_all_datasets(
                                    formula,
                                    sm.families.Poisson(link=link),
                                    is_glm=True
                                )
                                model_metrics.append((
                                    "count",
                                    formula,
                                    link.__class__.__name__,
                                    worst_aic,
                                    worst_loss,
                                    placeholder_params
                                ))

        elif self.loss == "BinaryCrossEntropy":
            # Validate all ys
            for y in ys:
                uni = set(np.unique(y))
                if not uni.issubset({0, 1}):
                    raise ValueError("All y must be 0/1 for binary task")
            for L in range(1, num_params + 1):
                for subset in itertools.combinations(features, L):
                    formula = "y ~ " + " + ".join(subset)
                    worst_aic, worst_loss, placeholder_params = evaluate_all_datasets(
                        formula,
                        sm.families.Binomial(link=sm.families.links.Logit()),
                        is_glm=True
                    )
                    model_metrics.append((
                        "binary",
                        formula,
                        "Logit",
                        worst_aic,
                        worst_loss,
                        placeholder_params
                    ))

        else:  # OLS
            for intercept_flag in [True, False]:
                N = min(num_params - int(intercept_flag), n_features)
                for L in range(1, N + 1):
                    for subset in itertools.combinations(features, L):
                        if intercept_flag:
                            formula = "y ~ " + " + ".join(subset)
                        else:
                            formula = "y ~ " + " + ".join(subset) + " - 1"
                        worst_aic, worst_loss, placeholder_params = evaluate_all_datasets(
                            formula,
                            None,
                            is_glm=False
                        )
                        model_metrics.append((
                            "ols",
                            formula,
                            "OLS",
                            worst_aic,
                            worst_loss,
                            placeholder_params
                        ))
        columns = ["Task", "Formula", "Link/Method", "AIC", "Loss/MSE", "Parameters"]
        self._results = (
            pd.DataFrame(model_metrics, columns=columns)
            .sort_values(by="AIC")
            .reset_index(drop=True)
        )

    def predict(self, X: np.ndarray, expression_id=0) -> np.ndarray:
        # By default uses the best expression (first one)
        df = pd.DataFrame(X, columns=[f"x{i+1}" for i in range(X.shape[1])])
        best = self._results.iloc[expression_id]
        params = best.Parameters

        # Linear prediction
        linpred = params.Intercept + sum(df[col] * params[col] for col in df.columns)

        if self.loss == "BinaryCrossEntropy":
            # logistic
            return 1 / (1 + np.exp(-linpred))

        if self.loss == "PoissonDeviance" and best["Link/Method"] == "log":
            return np.exp(linpred)

        return linpred

    def get_expressions(self) -> pd.Series:
        def formula_to_expression(row) -> str:
            params = row['Parameters']
            formula = row['Formula']
            link = row['Link/Method'].strip().lower()  # e.g., "log", "logit", "identity"

            # Build the linear predictor
            features = [f for f in self._feature_names if f in formula]
            terms = []
            # Handle intercept if present
            if hasattr(params, "index"):  # real params from statsmodels
                if "Intercept" in params.index:
                    terms.append(f"{params['Intercept']:.4f}")
                terms.extend(f"{params[f]:.4f}*{f}" for f in features)
            else:
                # Placeholder params from fit_multi
                terms = []
                if "- 1" not in formula:
                    terms.append("u0")
                terms.extend(f"u{i+1} * {f}" for i, f in enumerate(features))

            linpred = " + ".join(terms) if terms else "0.0000"

            # Wrap in inverse link function
            if "log" == link:
                expr = f"exp({linpred})"
            elif "logit" == link:
                expr = f"sigmoid({linpred})" 
            elif "identity" == link or "ols" in link:
                expr = linpred
            else:
                expr = f"{link}^-1({linpred})"  # generic fallback

            return expr
        res = self._results.apply(formula_to_expression, axis=1)
        return res
        
    @property
    def engine_variable_names(self) -> Sequence[str]:
        return self._feature_names
