import numpy as np
import networkx as nx
import shap
import lightgbm as lgb
import seaborn as sns
import matplotlib.pyplot as plt
from xgboost import XGBClassifier
from catboost import CatBoostClassifier
from itertools import combinations
from sklearn.model_selection import train_test_split

class NetworkConstructor:
    def __init__(self, cutoff, core_name="LightGBM"):
        """
        Initialize the NetworkConstructor with the dataset, feature combinations, cutoff value, and core model.
        
        :param df: DataFrame containing the data.
        :param feature_combination_dict: Dictionary with feature combinations for each pair of categories.
        :param cutoff: Threshold to determine whether an edge should be added to the network.
        :param core_name: Name of the core model to use ('LightGBM', 'XGBoost', or 'CatBoost').
        """
        self.cutoff = cutoff
        self.core_name = core_name
        self.core_model = self._select_core_model(core_name)

    def _select_core_model(self, core_name):
        """
        Select the core model based on the given core name.
        
        :param core_name: Name of the core model ('LightGBM', 'XGBoost', or 'CatBoost').
        :return: The core model.
        """
        if core_name == "LightGBM":
            # No need to return a model, as LightGBM's training code is used directly.
            return None
        elif core_name == "XGBoost":
            return XGBClassifier(n_estimators=1000, max_depth=5)
        elif core_name == "CatBoost":
            return CatBoostClassifier(verbose=False, iterations=1500, max_depth=5)
        else:
            raise ValueError("Unsupported core. Choose from LightGBM, XGBoost, or CatBoost.")

    def construct_network(self,df,feature_combination_dict, edge_weight = False, show_matrix = True):
        """
        Construct networks for each pair of categories and return a list of interaction networks.
        
        :return: List of constructed interaction networks.
        """
        self.feature_combination_dict = feature_combination_dict
        category = dict.fromkeys(df['Disease'])
        category_pairs = list(combinations(category, 2))
        Interactions = []

        for i, (Cat_A, Cat_B) in enumerate(category_pairs):
            df_filtered = df[df['Disease'].isin([Cat_A, Cat_B])]
            X = df_filtered.drop("Disease", axis=1)
            X = X[feature_combination_dict[f'{Cat_A} vs {Cat_B}']]
            y = df_filtered['Disease']
            
            # Shuffle and reset indices
            X = X.reset_index(drop=True)
            y = y.reset_index(drop=True)
            shuffle_index = np.random.permutation(X.index)
            X = X.iloc[shuffle_index]
            y = y.iloc[shuffle_index]
            y = y.map({Cat_A: 0, Cat_B: 1})
            
            # Split the data
            X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
            
            # Train the model based on the selected core
            if self.core_name == "LightGBM":
                d_train = lgb.Dataset(X_train, label=y_train)
                d_test = lgb.Dataset(X_test, label=y_test)
                params = {
                    "max_bin": 512,
                    "learning_rate": 0.05,
                    "boosting_type": "gbdt",
                    "objective": "binary",
                    "metric": "binary_logloss",
                    "num_leaves": 10,
                    "verbose": -1,
                    "boost_from_average": True,
                    "early_stopping_rounds": 50,
                    "verbose_eval": 1000,
                    "class_weight": 'balanced'
                }
                model = lgb.train(
                    params,
                    d_train,
                    1000,
                    valid_sets=[d_test],
                )
            else:
                model = self.core_model.fit(X, y)
            
            # Compute SHAP interaction values
            shap_interaction_values = shap.TreeExplainer(model, feature_perturbation="interventional").shap_interaction_values(X)
            interaction_matrix = np.abs(shap_interaction_values).sum(0)
            
            # Remove self-interactions
            np.fill_diagonal(interaction_matrix, 0)
            
            # Sort and select the top interactions
            inds = np.argsort(-interaction_matrix.sum(0))[:len(self.feature_combination_dict[f'{Cat_A} vs {Cat_B}'])]
            sorted_ia_matrix = interaction_matrix[inds, :][:, inds]
            feature_names = X.columns[inds]
            # show matrix
            if show_matrix:
                plt.figure(figsize=(8, 8))
                sns.heatmap(
                    sorted_ia_matrix, 
                    annot=True, 
                    fmt=".2f", 
                    cmap='coolwarm', 
                    cbar_kws={'label': 'Intensity'}, 
                    cbar=False
                )

                plt.xticks(
                    ticks=range(len(X.columns[inds])), 
                    labels=X.columns[inds], 
                    rotation=45, 
                    horizontalalignment="right"
                )
                plt.yticks(
                    ticks=range(len(X.columns[inds])), 
                    labels=X.columns[inds], 
                    rotation=45, 
                    horizontalalignment="right"
                )
                plt.title(f"{Cat_A} vs {Cat_B}")
                plt.tight_layout()
                plt.show()
            # Create a network graph
            G = nx.Graph()
            
            # Add nodes with feature names
            G.add_nodes_from(feature_names)
            
            # Add edges with weights from the interaction matrix
            for row in range(sorted_ia_matrix.shape[0]):
                for col in range(row + 1, sorted_ia_matrix.shape[1]):
                    weight = sorted_ia_matrix[row, col]
                    if edge_weight:
                        G.add_edge(feature_names[row], feature_names[col], weight=weight)
                    else:
                        if weight > self.cutoff:
                            G.add_edge(feature_names[row], feature_names[col])
            
            Interactions.append(G)

        return Interactions

    def compose_all(self, Interactions):
        """
        Combine all networks into one single network.
        
        :param Interactions: List of interaction networks.
        :return: A single merged network graph.
        """
        combined_graph = nx.compose_all(Interactions)
        return combined_graph

    def remove_isolated_nodes(self, Graph_BMI):
        """
        Remove isolated nodes from the graph.
        
        :param Graph_BMI: The network graph from which to remove isolated nodes.
        :return: The graph with isolated nodes removed.
        """
        isolated_nodes = list(nx.isolates(Graph_BMI))
        Graph_BMI.remove_nodes_from(isolated_nodes)
        return Graph_BMI
    
    def save_graph(self, Graph_BMI, path):
        """
        Save .graphml file.
        """
        nx.write_graphml(Graph_BMI, f'{path}/network.graphml')
        print(f"Network file saved to {path}/network.graphml")
        return None
    
    def read_graph(self, path):
        """
        Read .graphml file.
        """
        Graph_BMI = nx.read_graphml(path)
        return Graph_BMI

