import networkx as nx
import pandas as pd
import numpy as np
import plotly.graph_objects as go
from lightgbm import LGBMRegressor, LGBMClassifier
from sklearn.inspection import permutation_importance
from sklearn.model_selection import train_test_split
from typing import Tuple, List, Dict


class InsightEngine:
    def __init__(self, target_col: str, task: str = "regression"):
        self.target_col = target_col
        self.task = task
        self.importance_df_ = None
        self.redundancy_clusters_ = {}
        self.leakage_candidates_ = []
        self.useless_features_ = []
        self.corr_matrix_ = None

    def fit(self, df: pd.DataFrame):
        X = df.drop(columns=[self.target_col], errors="ignore")
        y = df[self.target_col]

        # 1. Permutation Importance (Juiz Imparcial)
        X_train, X_val, y_train, y_val = train_test_split(
            X, y, test_size=0.2, random_state=42
        )

        model = (
            LGBMRegressor(n_estimators=100, verbose=-1)
            if self.task == "regression"
            else LGBMClassifier(n_estimators=100, verbose=-1)
        )

        # Simplificação para rodar rápido: Trata tudo como número
        X_train_num = X_train.select_dtypes(include=np.number).fillna(0)
        X_val_num = X_val.select_dtypes(include=np.number).fillna(0)

        model.fit(X_train_num, y_train)
        result = permutation_importance(
            model, X_val_num, y_val, n_repeats=5, random_state=42, n_jobs=-1
        )

        self.importance_df_ = pd.DataFrame(
            {"feature": X_train_num.columns, "importance": result.importances_mean}
        ).sort_values(by="importance", ascending=False)

        # 2. Diagnóstico
        self.leakage_candidates_ = self.importance_df_[
            self.importance_df_["importance"] > 0.8
        ].feature.tolist()
        self.useless_features_ = self.importance_df_[
            self.importance_df_["importance"] <= 0
        ].feature.tolist()

        # 3. Correlação
        self.corr_matrix_ = (
            X.select_dtypes(include=np.number).corr(method="spearman").abs()
        )

        return self

    def plot_correlation_network(self, threshold: float = 0.75) -> go.Figure:
        if self.corr_matrix_ is None:
            return go.Figure()

        G = nx.Graph()
        for col in self.corr_matrix_.columns:
            G.add_node(col)

        for i in range(len(self.corr_matrix_.columns)):
            for j in range(i):
                val = self.corr_matrix_.iloc[i, j]
                if val > threshold:
                    G.add_edge(
                        self.corr_matrix_.columns[i],
                        self.corr_matrix_.columns[j],
                        weight=val,
                    )

        # Comunidades
        communities = nx.community.greedy_modularity_communities(G)
        node_groups = {}
        self.redundancy_clusters_ = {}

        for idx, comm in enumerate(communities):
            members = list(comm)
            if len(members) > 1:
                self.redundancy_clusters_[f"Cluster_{idx}"] = members
            for node in members:
                node_groups[node] = idx

        # Plotly Logic
        pos = nx.spring_layout(G, k=0.3, seed=42)
        edge_x, edge_y = [], []
        for edge in G.edges():
            x0, y0 = pos[edge[0]]
            x1, y1 = pos[edge[1]]
            edge_x.extend([x0, x1, None])
            edge_y.extend([y0, y1, None])

        node_x, node_y, node_text, node_color = [], [], [], []
        for node in G.nodes():
            x, y = pos[node]
            node_x.append(x)
            node_y.append(y)
            node_text.append(f"{node} (Group {node_groups.get(node,0)})")
            node_color.append(node_groups.get(node, 0))

        return go.Figure(
            data=[
                go.Scatter(
                    x=edge_x, y=edge_y, line=dict(width=1, color="#888"), mode="lines"
                ),
                go.Scatter(
                    x=node_x,
                    y=node_y,
                    mode="markers",
                    text=node_text,
                    marker=dict(color=node_color, size=15, colorscale="Turbo"),
                ),
            ],
            layout=go.Layout(title="Redundancy Network", showlegend=False),
        )

    def generate_curation_report(self):
        actions = []
        print("# Relatório de Curadoria ScoutML")

        if self.leakage_candidates_:
            print(f"🔴 CRÍTICO (Leakage): {self.leakage_candidates_}")
            actions.extend([("drop", f) for f in self.leakage_candidates_])

        if self.redundancy_clusters_:
            print(f"🟡 AVISO (Redundância): {self.redundancy_clusters_}")

        if self.useless_features_:
            print(
                f"⚪ INFO (Zombie Features): {len(self.useless_features_)} features inúteis."
            )

        return self._generate_code(actions)

    def _generate_code(self, actions):
        drops = [x[1] for x in actions if x[0] == "drop"]
        return f"cols_to_drop = {drops}"
