# pkgs/mariner/src/mariner/estimator.py
import numpy as np
from sklearn.base import BaseEstimator
from sklearn.utils.validation import check_X_y, check_array, check_is_fitted
import matplotlib.pyplot as plt
import seaborn as sns
import warnings

from hiutils.core import (
    UniLasso, OLS,
    make_interactions, cv,
    has_unit_leverage
)

from importlib.metadata import version as pkg_version, PackageNotFoundError
    

class MARINER(BaseEstimator):
    """
    Two-stage regression with marginal selection and UniLasso on both main effects and interactions via CV.

    Parameters
    ----------
    lmda_path : array-like of shape (n_lmdas,)
        Path of the lambda values for UniLasso. If None, it will be generated using generate_lambda_path.  

    n_folds : int
        Number of folds for cross-validation.

    plot_cv_curve : bool  
        If true, plot the CV R2 against -log(λ).

    cv1se : bool 
        If true, use the one-standard error rule at CV. 

    verbose : bool
        If true, print progress messages.

    interaction_candidates : array-like of int, optional
        Restrict interactions to those involving any of these variable indices (0-based indices).

    interaction_pairs : array-like of shape (n_pairs, 2), optional
        Restrict interactions to exactly these pairs of variables (0-based indices).

    vars_names : list of str, optional
        Names of the main effects variables/features.

    zero_cutoff : float, optional
        All pvalues below this threshold are considered zero. Defaults to 1e-20.

    interactions_threshold : float, optional
        The p-value threshold used to select interactions. If None, it will be determined using the largest log-gap rule.

    save_plots : str, optional
        Directory path to save the plots as PNG files. If None, plots will not be saved.
        like "/path/to/dir". Defaults to None.

    Attributes
    ----------
    n_features\_ : int
        Number of features in the dataset.

    regressor\_ : UniLasso object
        Fitted Unilasso model after CV.

    triplet_regressors\_ : dict of dict of OLS
        Dictionary of OLS models for each triplet (X_j, X_k, X_j*X_k).

    selected_pairs\_ : array-like of shape (n_pairs, 2)
        Selected interaction pairs (j, k).

    main_effects_active_set\_ : array-like
        Indices of active main effects.

    interactios_active_set\_ : array-like
        Indices of active interactions.

    n_folds\_ : int
        Number of folds for cross-validation.

    main_effects_names\_ : list of str
        Names of the main effects.

    interactions_names\_ : list of str
        Names of the interaction terms.

    cv_errors\_ : array-like of shape (n_folds, n_lmdas)
        Cross-validation errors.

    lmda_path\_ : array-like
        Lmda path for cv  

    plot_cv_curve\_ : bool
        If true, plot the CV R2 against -log(λ).

    cv1se\_ : bool 
        If true, use the one-standard error rule at CV. 

    verbose\_ : bool
        If true, print progress messages.

    interaction_candidates\_ : array-like of int
        Restrict interactions to those involving any of these variable indices (0-based indices).
    
    interaction_pairs\_ : array-like of shape (n_pairs, 2)
        Restrict interactions to exactly these pairs of variables (0-based indices).

    vars_names\_ : list of str, optional
        Names of the main effects variables/features.

    zero_cutoff\_ : float
        All pvalues below this threshold are considered zero. Defaults to 1e-20.

    interactions_threshold\_ : float
        The p-value threshold used to select interactions. If None, it will be determined using the largest log-gap rule.
    
    save_plots\_ : str
        Directory path to save the plots as PNG files. If None, plots will not be saved.
        like "/path/to/dir". Defaults to None.
    """

    def __init__(self, lmda_path=None, n_folds=10, plot_cv_curve=False, cv1se=False, verbose=False, interaction_candidates=None, interaction_pairs=None, vars_names=None, zero_cutoff=1e-20, interactions_threshold=None, save_plots=None):

        """
        Two-stage regression with marginal selection and UniLasso on both main effects and interactions via CV.

        Parameters
        ----------
        lmda_path : array-like of shape (n_lmdas,)
            Path of the lambda values for UniLasso. If None, it will be generated using generate_lambda_path.  

        n_folds : int
            Number of folds for cross-validation.

        plot_cv_curve : bool  
            If true, plot the CV R2 against -log(λ).

        cv1se : bool 
            If true, use the one-standard error rule at CV. 

        verbose : bool
            If true, print progress messages.

        interaction_candidates : array-like of int, optional
            Restrict interactions to those involving any of these variable indices (0-based indices).

        interaction_pairs : array-like of shape (n_pairs, 2), optional
            Restrict interactions to exactly these pairs of variables (0-based indices).

        vars_names : list of str, optional
            Names of the main effects variables/features.

        zero_cutoff : float, optional
            All pvalues below this threshold are considered zero. Defaults to 1e-20.

        interactions_threshold : float, optional
            The p-value threshold used to select interactions. If None, it will be determined using the largest log-gap rule.

        save_plots : str, optional
            Directory path to save the plots as PNG files. If None, plots will not be saved.
            like "/path/to/dir". Defaults to None.

        Attributes
        ----------
        n_features\_ : int
            Number of features in the dataset.

        regressor\_ : UniLasso object
            Fitted Unilasso model after CV.

        triplet_regressors\_ : dict of dict of OLS
            Dictionary of OLS models for each triplet (X_j, X_k, X_j*X_k).

        selected_pairs\_ : array-like of shape (n_pairs, 2)
            Selected interaction pairs (j, k).

        main_effects_active_set\_ : array-like
            Indices of active main effects.

        interactios_active_set\_ : array-like
            Indices of active interactions.

        n_folds\_ : int
            Number of folds for cross-validation.

        main_effects_names\_ : list of str
            Names of the main effects.

        interactions_names\_ : list of str
            Names of the interaction terms.

        cv_errors\_ : array-like of shape (n_folds, n_lmdas)
            Cross-validation errors.

        lmda_path\_ : array-like
            Lmda path for cv  

        plot_cv_curve\_ : bool
            If true, plot the CV R2 against -log(λ).

        cv1se\_ : bool 
            If true, use the one-standard error rule at CV. 

        verbose\_ : bool
            If true, print progress messages.

        interaction_candidates\_ : array-like of int
            Restrict interactions to those involving any of these variable indices (0-based indices).
        
        interaction_pairs\_ : array-like of shape (n_pairs, 2)
            Restrict interactions to exactly these pairs of variables (0-based indices).

        vars_names\_ : list of str, optional
            Names of the main effects variables/features.

        zero_cutoff\_ : float
            All pvalues below this threshold are considered zero. Defaults to 1e-20.

        interactions_threshold\_ : float
            The p-value threshold used to select interactions. If None, it will be determined using the largest log-gap rule.
        
        save_plots\_ : str
            Directory path to save the plots as PNG files. If None, plots will not be saved
            like "/path/to/dir". Defaults to None.
        """
        
        if interaction_candidates is not None and interaction_pairs is not None:
            raise ValueError("Specify only one of interaction_candidates or interaction_pairs, not both.")

        self.regressor_ = UniLasso(lmda_path=lmda_path)
        self.triplet_regressors_ = None
        
        self.n_folds_ = n_folds

        self.n_features_ = None
        self.main_effects_active_set_ = None
        self.selected_pairs_ = None

        self.main_effects_names_ = None
        self.interactions_names_ = None
        self.cv_errors = None
        self.lmda_path_ = None
        self.plot_cv_curve_ = plot_cv_curve
        self.cv1se_ = cv1se
        self.verbose_ = verbose
        self.interaction_candidates_ = (
            np.asarray(interaction_candidates, dtype=int)
            if interaction_candidates is not None else None
        )
        self.interaction_pairs_ = (
            np.asarray(interaction_pairs, dtype=int)
            if interaction_pairs is not None else None
        )
        if self.interaction_pairs_ is not None:
            if self.interaction_pairs_.ndim != 2 or self.interaction_pairs_.shape[1] != 2:
                raise ValueError("interaction_pairs must be a 2D array-like with shape (n_pairs, 2).")
        
        if self.interaction_candidates_ is not None:
            if self.interaction_candidates_.ndim != 1:
                raise ValueError("interaction_candidates must be a 1D array-like of variable indices.")
            
        self.vars_names_ = vars_names
        self.zero_cutoff_ = zero_cutoff
        self.interactions_threshold_ = interactions_threshold
        self.save_plots_ = save_plots

    @property
    def version(self):
        try:
            return pkg_version("mariner-stats")
        except PackageNotFoundError:
            from mariner_stats import __version__ as local_version
            return local_version
    
    def fit_triplet_models(self,X, y):
        """
        Regress y on [X_j,X_k,X_j*X_k] with intercept for all j<k

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Design matrix.
        y : array-like of shape (n_samples,)
            Response vector.

        Returns
        -------
        None 
        """
        n, p = X.shape
        assert p>1
        if not (n > 3):
            raise ValueError("Need at least 4 samples to fit triplet OLS models.")
        
        if self.interaction_pairs_ is not None:
            allowed_pairs = [tuple(pair) for pair in self.interaction_pairs_]
        elif self.interaction_candidates_ is not None:
            allowed_pairs = [
                (j, k)
                for j in range(p - 1)
                for k in range(j + 1, p)
                if (j in self.interaction_candidates_ or k in self.interaction_candidates_)
            ]
        else:
            allowed_pairs = [(j, k) for j in range(p - 1) for k in range(j + 1, p)]
    
        total_pairs = len(allowed_pairs)
        if self.verbose_:
            print(f"  Fitting {total_pairs} triplet OLS models...")
        
        self.instable_pairs_ = []
        for count, (j, k) in enumerate(allowed_pairs, start=1):
            F = np.hstack([X[:, j : j + 1], X[:, k : k + 1], X[:, j : j + 1] * X[:, k : k + 1]])                        # (n, 3)

            if np.linalg.matrix_rank(np.hstack((np.ones((F.shape[0], 1)), F))) < 4: 
                self.instable_pairs_.append((j, k))
            elif has_unit_leverage(np.hstack((np.ones((F.shape[0], 1)), F))):
                self.instable_pairs_.append((j, k))
            else:
                self.triplet_regressors_[j][k].fit(F, y)
            if self.verbose_ and count % max(1, total_pairs // 10) == 0:
                print(f"  Progress: {count}/{total_pairs} triplets fitted...")
        if self.verbose_:
            print("  Triplet model fitting complete.")
    
    def scan_interactions(self):
        """
        Implements the screening procedure for the interactions.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Design matrix.
        y : array-like of shape (n_samples,)
            Response vector.
    
        Returns
        -------
        None
        """

        n_features = self.n_features_

        if self.triplet_regressors_ is None:
            raise RuntimeError("triplet_regressors_ not initialized. Call fit_triplet_models first.")

        if self.verbose_:
            print("Scanning interaction using p-values ...")

        if self.interaction_pairs_ is not None:
            allowed_pairs = [tuple(pair) for pair in self.interaction_pairs_]
        elif self.interaction_candidates_ is not None:
            allowed_pairs = [
                (j, k)
                for j in range(n_features - 1)
                for k in range(j + 1, n_features)
                if (j in self.interaction_candidates_ or k in self.interaction_candidates_)
            ]
        else:
            allowed_pairs = [(j, k) for j in range(n_features - 1) for k in range(j + 1, n_features)]

        allowed_pairs = list(set(allowed_pairs)-set(self.instable_pairs_))
        
        scan_coefs = np.zeros((n_features,n_features))
        pvals = np.ones((n_features,n_features))
        for j, k in allowed_pairs:
            pvals[j,k] = self.triplet_regressors_[j][k].p_value_t_test_[-1]
            scan_coefs[j,k] = np.abs(self.triplet_regressors_[j][k].slopes_[-1])
    

        pvals_ea = {(j,k):self.triplet_regressors_[j][k].p_value_t_test_[-1] for (j,k) in allowed_pairs}
        self.pvals_ = {(self.main_effects_names_[j],self.main_effects_names_[k]):pval for (j,k),pval in pvals_ea.items()}
        
        # point that best separates the two clusters 
        pvals_ea_sorted = sorted(pvals_ea.items(),key=lambda x: x[1],reverse=False) # ascending order 
        num_zeros_ea = len(pvals_ea_sorted) - len([pval for _,pval in pvals_ea_sorted if pval>self.zero_cutoff_])

        # threshold by the biggest log-gap heuristic
        if self.interactions_threshold_ is not None:
            selected_pairs_ea = [pair for (pair,pval) in pvals_ea_sorted if pval<=self.interactions_threshold_]
            self.threshold_pval_ = self.interactions_threshold_
        else  :
            tmp = [pval for _,pval in pvals_ea_sorted if pval>self.zero_cutoff_]
            if len(tmp)<=1 : 
                selected_pairs_ea = [pair for (pair,_) in pvals_ea_sorted]
                self.threshold_pval_ = tmp[0] if len(tmp)==1 else 0.0
            else : 
                tmp = np.log(np.array(tmp)) # need to investigate why taking log of pvalues does better than using their raw value
                i_tmp = np.argmax(tmp[1:]-tmp[:-1]) + num_zeros_ea
                selected_pairs_ea = [pair for (pair,_) in pvals_ea_sorted[:i_tmp+1]]
                self.threshold_pval_ = pvals_ea_sorted[i_tmp][1]

        if self.verbose_:
            add_on_tmp = "largest log-gap rule" if self.interactions_threshold_ is None else f" threshold={self.interactions_threshold_:.2e}"
            print(f"  Selected {len(selected_pairs_ea)} interactions ({add_on_tmp}).")

        if self.plot_cv_curve_ : 
            
            mask = np.triu(np.ones_like(pvals, dtype=bool), k=1)
            pv = pvals[mask]
            plt.figure(figsize=(8, 5))
            # plt.hist(pv, bins=int(len(pv) / 20), edgecolor="black")
            plt.hist(-np.log(np.maximum(pv,self.zero_cutoff_)), bins=int(len(pv) / 20), edgecolor="black")

            sel_pv = [pvals[j,k] for j,k in selected_pairs_ea]
            # plt.scatter(sel_pv, np.zeros_like(sel_pv), color="orange", s=50, label="selected pairs", zorder=10)
            plt.scatter(-np.log(np.maximum(sel_pv,self.zero_cutoff_)), np.zeros_like(sel_pv), color="orange", s=50, label="selected pairs", zorder=10)
            title = "Interaction -log(p-values) from triplet models | " 
            title += "largest log-gap rule" if self.interactions_threshold_ is None else f" threshold={self.interactions_threshold_:.2e}"
            plt.title(title)

            plt.xlabel('-log(p-value)')
            plt.ylabel('Count')
            plt.legend()
            plt.tight_layout()
            if self.save_plots_:
                plt.savefig(self.save_plots_+"/interactions_pvalues_histogram.png", dpi=300)
            plt.show()

            mask = np.tril(np.ones_like(pvals, dtype=bool), k=0)
            plt.figure(figsize=(8, 8))
            ax = sns.heatmap(
                pvals,
                mask=mask,
                cmap='viridis',
                xticklabels=self.main_effects_names_,
                yticklabels=self.main_effects_names_,
                vmin=0,
                vmax=1,
                square=True,
                cbar_kws={"shrink": .8}
            )
            ax.set_facecolor('white')
            for j, k in selected_pairs_ea:
                ax.text(k + 0.5,j + 0.5,'★',ha='center', va='center',color='white',fontsize=16,zorder=10)
            plt.tight_layout()
            plt.title(title)
            if self.save_plots_:
                plt.savefig(self.save_plots_+"/interactions_pvalues_heatmap.png", dpi=300)
            plt.show()

        self.selected_pairs_ = np.array(selected_pairs_ea)
        
    def fit(self, X, y, lmda_path=None, tolerance=1e-10):
        """
        Fit the Unilasso regression model with interactions.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Design matrix.
        y : array-like of shape (n_samples,)
            Response vector.
        lmda_path : array-like of shape (n_lmdas,), optional
            Path of the lambda values for UniLasso. If None, it will be generated using generate_lambda_path.  
        tolerance : float, optional
            Tolerance for determining active set / removing zero-std features. Default is 1e-10.
        

        Returns
        -------
        None
        """
        
        X, y = check_X_y(X, y, y_numeric=True)

        self.main_effects_stds = np.std(X, axis=0, ddof=1)
        zero_std_mask = np.abs(self.main_effects_stds) < tolerance
        if np.any(zero_std_mask):
            warnings.warn(
                f"Features at indices {np.where(zero_std_mask)[0]} have zero standard deviation and will be removed.",
                category=UserWarning,
                stacklevel=2,
            )
            X = X[:, ~zero_std_mask]
        self.main_effects_means = np.mean(X, axis=0)
        self.main_effects_stds = self.main_effects_stds[~zero_std_mask]
        X = (X - self.main_effects_means) / self.main_effects_stds


        _, n_features = X.shape
        self.n_features_ = n_features

        if self.interaction_pairs_ is not None:
            max_idx = self.interaction_pairs_.max()
            if max_idx >= n_features:
                raise ValueError("interaction_pairs contain indices out of range for X.")
        if self.interaction_candidates_ is not None:
            max_idx = self.interaction_candidates_.max()
            if max_idx >= n_features:
                raise ValueError("interaction_candidates contain indices out of range for X.")
            
        self.main_effects_names_ = [self.vars_names_[j] if self.vars_names_ is not None else f'X{j}' for j in range(n_features)]
        self.triplet_regressors_  = {j:{k:OLS(vars_names=[self.main_effects_names_[j],self.main_effects_names_[k],f'{self.main_effects_names_[j]}*{self.main_effects_names_[k]}']) for k in range(j+1,n_features)} for j in range(n_features-1)}

        if self.verbose_:
            print(f"Starting MARINER fit with {n_features} features ...")

        if n_features > 1 : 
            if self.verbose_:
                print("Step 1: Fitting triplet models ...")
            self.fit_triplet_models(X, y)
            if self.verbose_:
                print("Step 2: Scanning interactions ...")
            self.scan_interactions()
            if self.verbose_:
                print("Step 3: Constructing interaction matrix ...")
            interactions_X, self.interactions_names_ = make_interactions(X, self.selected_pairs_, vars_names=self.vars_names_)      # (n_samples, n_pairs)
            full_X = np.hstack([X,interactions_X])
        else :
            full_X = X
            self.interactions_names_ = []


        self.regressor_.set_vars_names(self.main_effects_names_+self.interactions_names_)
        if self.verbose_:
            print("Step 4: Cross-validating UniLasso ...")
        save_plots = self.save_plots_+"/cv_curve_unilasso.png" if self.save_plots_ else None
        cv_results = cv(
            base=self.regressor_,
            X=full_X,
            y=y,
            n_folds=self.n_folds_,
            lmda_path=lmda_path,
            plot_cv_curve=self.plot_cv_curve_,
            save_plots=save_plots,
            cv1se=self.cv1se_
        ) # self.regressor_ is fited in-place after CV
        self.lmda_path_ = cv_results['lmda_path']
        self.cv_errors_ = cv_results['cv_errors']
        main_effects_slopes = self.regressor_.slopes_[:1,:n_features]                                                           # (1, n_features)
        interactions_slopes = self.regressor_.slopes_[:1,n_features:]                                                           # (1, n_pairs)
        self.main_effects_active_set_ = np.where(np.abs(main_effects_slopes) > tolerance)[1] 
        self.interactions_active_set_ = np.where(np.abs(interactions_slopes) > tolerance)[1] 

        # change back to original scale
        main_effects_new_slopes = np.empty_like(main_effects_slopes)
        interactions_new_slopes = np.empty_like(interactions_slopes)
        intercept_new = self.regressor_.intercept_.copy()

        for idx, (j, k) in enumerate(self.selected_pairs_):
            interactions_new_slopes[0, idx] = interactions_slopes[0, idx] / (self.main_effects_stds[j] * self.main_effects_stds[k])
            intercept_new += interactions_new_slopes[0, idx] * (self.main_effects_means[j] * self.main_effects_means[k])

        for j in range(n_features):
            main_effects_new_slopes[0, j] = main_effects_slopes[0, j] / self.main_effects_stds[j]
            for idx, (jj, kk) in enumerate(self.selected_pairs_):
                if jj == j:
                    main_effects_new_slopes[0, j] -= interactions_new_slopes[0, idx] * self.main_effects_means[kk] 
                if kk == j:
                    main_effects_new_slopes[0, j] -= interactions_new_slopes[0, idx] * self.main_effects_means[jj]
            intercept_new -= main_effects_slopes[0, j] * self.main_effects_means[j] / self.main_effects_stds[j]
        
        self.main_effects_active_set_ = np.where(np.abs(main_effects_new_slopes) > tolerance)[1]
        self.interactions_active_set_ = np.where(np.abs(interactions_new_slopes) > tolerance)[1]

        self.regressor_.slopes_ = np.hstack([main_effects_new_slopes, interactions_new_slopes])
        self.regressor_.intercept_ = intercept_new

        if self.verbose_:
            print("MARINER fit complete.\n")
    def get_active_variables(self):
        """
        Get the names of the active variables for the fitted model.
        This includes both main effects and interactions.

        Parameters
        ----------
        None

        Returns
        -------
        active_vars : list of str
            List of names of the active variables.
        """
        check_is_fitted(self.regressor_)
        
        active_vars = []
        
        for i in self.main_effects_active_set_:
            active_vars.append(self.main_effects_names_[i])
    
        if self.n_features_>1:
            for i in self.interactions_active_set_:
                active_vars.append(self.interactions_names_[i])

        return active_vars
    
    def get_fitted_function(self, tolerance=1e-10):
        """
        Get the fitted model string representation.

        Parameters
        ----------
        None

        Returns
        -------
        fitted_model_rep : str
            string representation of the fitted model
        """
        check_is_fitted(self.regressor_)
        
        return self.regressor_.get_fitted_function(self.regressor_.lmda_path_[0],tolerance) 
         
    def predict(self, X):
        """
        Predict using the two-stage regression model with interactions.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Design matrix.

        Returns
        -------
        y_pred : array-like of shape (n_samples,)
            Predicted values using the two-stage regression model.
        """
        check_is_fitted(self.regressor_)
        
        X = check_array(X)

        if self.n_features_>1:
            interactions_X, _ = make_interactions(X,self.selected_pairs_)
            full_X = np.hstack([X,interactions_X])
        else : 
            full_X = X

        return self.regressor_.predict(full_X)[:,0]                                                                             # (n_samples, )
    