import os

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.base import clone
from sklearn.cluster import AgglomerativeClustering, KMeans
from sklearn.compose import ColumnTransformer
from sklearn.decomposition import PCA
from sklearn.ensemble import GradientBoostingClassifier, GradientBoostingRegressor, RandomForestClassifier, RandomForestRegressor
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression, Ridge
from sklearn.metrics import (
    accuracy_score,
    confusion_matrix,
    f1_score,
    mean_absolute_error,
    mean_squared_error,
    r2_score,
    silhouette_score,
)
from sklearn.model_selection import GridSearchCV, train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler


os.environ.setdefault("LOKY_MAX_CPU_COUNT", "2")


class PredictiveMLWorkbenchService:
    def run(self, csv_path, workflow, target_column, test_size, cv_folds, max_clusters):
        if not csv_path:
            return "", "", "", None, "Upload a CSV file first."

        try:
            df = pd.read_csv(csv_path)
        except Exception as exc:
            return "", "", "", None, f"Could not read CSV: {type(exc).__name__}: {exc}"

        if df.empty:
            return "", "", "", None, "Dataset is empty."

        try:
            if workflow == "Classification":
                return self._run_classification(df, target_column, test_size, cv_folds)
            if workflow == "Regression":
                return self._run_regression(df, target_column, test_size, cv_folds)
            if workflow == "Clustering":
                return self._run_clustering(df, target_column, max_clusters)
            return self._run_dimensionality_reduction(df, target_column)
        except Exception as exc:
            return "", "", "", None, f"Workflow failed: {type(exc).__name__}: {exc}"

    def _run_classification(self, df, target_column, test_size, cv_folds):
        x, y = self._supervised_split(df, target_column)
        preprocessor = self._build_preprocessor(x)

        candidates = [
            (
                "LogisticRegression",
                LogisticRegression(max_iter=600),
                {"model__C": [0.5, 1.0, 2.0]},
            ),
            (
                "RandomForestClassifier",
                RandomForestClassifier(random_state=42),
                {"model__n_estimators": [120, 220], "model__max_depth": [None, 8]},
            ),
            (
                "GradientBoostingClassifier",
                GradientBoostingClassifier(random_state=42),
                {"model__n_estimators": [80, 140], "model__learning_rate": [0.05, 0.1]},
            ),
        ]

        x_train, x_test, y_train, y_test = train_test_split(
            x,
            y,
            test_size=test_size,
            random_state=42,
            stratify=y if y.nunique() > 1 else None,
        )

        best_name, best_search = self._select_model(
            candidates=candidates,
            preprocessor=preprocessor,
            x_train=x_train,
            y_train=y_train,
            cv_folds=cv_folds,
            scoring="f1_macro",
        )

        preds = best_search.best_estimator_.predict(x_test)
        acc = accuracy_score(y_test, preds)
        macro_f1 = f1_score(y_test, preds, average="macro")
        metrics = "\n".join(
            [
                f"Accuracy: {acc:.4f}",
                f"Macro F1: {macro_f1:.4f}",
                f"CV Best Score: {best_search.best_score_:.4f}",
                f"Train Rows: {len(x_train)}",
                f"Test Rows: {len(x_test)}",
                f"Classes: {y.nunique()}",
                f"Best Params: {best_search.best_params_}",
            ]
        )

        fig = self._plot_confusion_matrix(y_test, preds)
        preview = x.head(8).to_string(index=False)
        status = "Completed end-to-end classification workflow with preprocessing, model selection, and evaluation."
        return best_name, metrics, preview, fig, status

    def _run_regression(self, df, target_column, test_size, cv_folds):
        x, y = self._supervised_split(df, target_column)
        if not pd.api.types.is_numeric_dtype(y):
            raise ValueError("Regression target column must be numeric.")

        preprocessor = self._build_preprocessor(x)
        candidates = [
            (
                "Ridge",
                Ridge(),
                {"model__alpha": [0.5, 1.0, 2.0, 5.0]},
            ),
            (
                "RandomForestRegressor",
                RandomForestRegressor(random_state=42),
                {"model__n_estimators": [120, 220], "model__max_depth": [None, 8]},
            ),
            (
                "GradientBoostingRegressor",
                GradientBoostingRegressor(random_state=42),
                {"model__n_estimators": [80, 140], "model__learning_rate": [0.05, 0.1]},
            ),
        ]

        x_train, x_test, y_train, y_test = train_test_split(
            x,
            y,
            test_size=test_size,
            random_state=42,
        )

        best_name, best_search = self._select_model(
            candidates=candidates,
            preprocessor=preprocessor,
            x_train=x_train,
            y_train=y_train,
            cv_folds=cv_folds,
            scoring="r2",
        )

        preds = best_search.best_estimator_.predict(x_test)
        r2 = r2_score(y_test, preds)
        mae = mean_absolute_error(y_test, preds)
        rmse = float(np.sqrt(mean_squared_error(y_test, preds)))
        metrics = "\n".join(
            [
                f"R2: {r2:.4f}",
                f"MAE: {mae:.4f}",
                f"RMSE: {rmse:.4f}",
                f"CV Best Score: {best_search.best_score_:.4f}",
                f"Train Rows: {len(x_train)}",
                f"Test Rows: {len(x_test)}",
                f"Best Params: {best_search.best_params_}",
            ]
        )

        fig = self._plot_regression_scatter(y_test, preds)
        preview = x.head(8).to_string(index=False)
        status = "Completed end-to-end regression workflow with preprocessing, model selection, and evaluation."
        return best_name, metrics, preview, fig, status

    def _run_clustering(self, df, target_column, max_clusters):
        x = df.copy()
        if target_column and target_column in x.columns:
            x = x.drop(columns=[target_column])

        preprocessor = self._build_preprocessor(x)
        transformed = preprocessor.fit_transform(x)
        transformed = np.asarray(transformed)

        sample = transformed
        if transformed.shape[0] > 1200:
            sample = transformed[:1200]

        best = None
        best_labels = None
        for n_clusters in range(2, max_clusters + 1):
            for name, estimator in [
                ("KMeans", KMeans(n_clusters=n_clusters, random_state=42, n_init=10)),
                ("AgglomerativeClustering", AgglomerativeClustering(n_clusters=n_clusters)),
            ]:
                labels = estimator.fit_predict(sample)
                if len(np.unique(labels)) < 2:
                    continue
                score = silhouette_score(sample, labels)
                if best is None or score > best["score"]:
                    best = {"name": name, "clusters": n_clusters, "score": score}
                    best_labels = labels

        if best is None:
            raise ValueError("Could not produce a valid clustering result.")

        reduced = PCA(n_components=2, random_state=42).fit_transform(sample)
        metrics = "\n".join(
            [
                f"Algorithm: {best['name']}",
                f"Clusters: {best['clusters']}",
                f"Silhouette Score: {best['score']:.4f}",
                f"Rows Used: {sample.shape[0]}",
                f"Features After Preprocessing: {sample.shape[1]}",
            ]
        )
        fig = self._plot_cluster_scatter(reduced, best_labels, title=f"{best['name']} clustering")
        preview = x.head(8).to_string(index=False)
        status = "Completed clustering workflow with preprocessing, model selection across algorithms, and evaluation."
        return best["name"], metrics, preview, fig, status

    def _run_dimensionality_reduction(self, df, target_column):
        x = df.copy()
        labels = None
        if target_column and target_column in x.columns:
            labels = x[target_column].astype(str)
            x = x.drop(columns=[target_column])

        preprocessor = self._build_preprocessor(x)
        transformed = preprocessor.fit_transform(x)
        transformed = np.asarray(transformed)

        n_components = 2 if transformed.shape[1] >= 2 else 1
        pca = PCA(n_components=n_components, random_state=42)
        reduced = pca.fit_transform(transformed)

        explained = pca.explained_variance_ratio_
        metrics = "\n".join(
            [
                f"Method: PCA",
                f"Components: {n_components}",
                f"Explained Variance: {', '.join(f'{v:.4f}' for v in explained)}",
                f"Cumulative Variance: {explained.sum():.4f}",
                f"Rows: {transformed.shape[0]}",
                f"Features After Preprocessing: {transformed.shape[1]}",
            ]
        )

        fig = self._plot_pca_scatter(reduced, labels)
        preview_df = pd.DataFrame(reduced, columns=[f"PC{i+1}" for i in range(n_components)])
        preview = preview_df.head(8).to_string(index=False)
        status = "Completed dimensionality reduction workflow with preprocessing and PCA evaluation."
        return "PCA", metrics, preview, fig, status

    def _build_preprocessor(self, x):
        numeric_cols = x.select_dtypes(include=["number"]).columns.tolist()
        categorical_cols = [col for col in x.columns if col not in numeric_cols]

        numeric_pipeline = Pipeline(
            steps=[
                ("imputer", SimpleImputer(strategy="median")),
                ("scaler", StandardScaler()),
            ]
        )

        categorical_pipeline = Pipeline(
            steps=[
                ("imputer", SimpleImputer(strategy="most_frequent")),
                ("onehot", OneHotEncoder(handle_unknown="ignore", sparse_output=False)),
            ]
        )

        return ColumnTransformer(
            transformers=[
                ("num", numeric_pipeline, numeric_cols),
                ("cat", categorical_pipeline, categorical_cols),
            ],
            remainder="drop",
        )

    def _supervised_split(self, df, target_column):
        if not target_column:
            raise ValueError("Target column is required for supervised workflows.")
        if target_column not in df.columns:
            raise ValueError(f"Target column `{target_column}` was not found.")

        x = df.drop(columns=[target_column])
        y = df[target_column]
        if x.shape[1] == 0:
            raise ValueError("Dataset needs at least one feature column.")
        return x, y

    def _select_model(self, candidates, preprocessor, x_train, y_train, cv_folds, scoring):
        best_name = None
        best_search = None

        for name, estimator, param_grid in candidates:
            pipeline = Pipeline(
                steps=[
                    ("preprocessor", clone(preprocessor)),
                    ("model", estimator),
                ]
            )
            search = GridSearchCV(
                estimator=pipeline,
                param_grid=param_grid,
                cv=cv_folds,
                scoring=scoring,
                n_jobs=1,
            )
            search.fit(x_train, y_train)
            if best_search is None or search.best_score_ > best_search.best_score_:
                best_name = name
                best_search = search

        return best_name, best_search

    def _plot_confusion_matrix(self, y_true, y_pred):
        fig, ax = plt.subplots(figsize=(5.5, 4.5))
        labels = np.unique(np.concatenate([np.asarray(y_true), np.asarray(y_pred)]))
        matrix = confusion_matrix(y_true, y_pred, labels=labels)
        im = ax.imshow(matrix, cmap="Blues")
        ax.set_title("Confusion Matrix")
        ax.set_xlabel("Predicted")
        ax.set_ylabel("True")
        ax.set_xticks(range(len(labels)))
        ax.set_xticklabels(labels, rotation=45, ha="right")
        ax.set_yticks(range(len(labels)))
        ax.set_yticklabels(labels)
        for i in range(matrix.shape[0]):
            for j in range(matrix.shape[1]):
                ax.text(j, i, str(matrix[i, j]), ha="center", va="center", color="black")
        fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
        fig.tight_layout()
        return fig

    def _plot_regression_scatter(self, y_true, y_pred):
        fig, ax = plt.subplots(figsize=(5.5, 4.5))
        ax.scatter(y_true, y_pred, alpha=0.75)
        min_val = min(np.min(y_true), np.min(y_pred))
        max_val = max(np.max(y_true), np.max(y_pred))
        ax.plot([min_val, max_val], [min_val, max_val], linestyle="--", color="red")
        ax.set_title("Actual vs Predicted")
        ax.set_xlabel("Actual")
        ax.set_ylabel("Predicted")
        fig.tight_layout()
        return fig

    def _plot_cluster_scatter(self, reduced, labels, title):
        fig, ax = plt.subplots(figsize=(5.5, 4.5))
        scatter = ax.scatter(reduced[:, 0], reduced[:, 1], c=labels, cmap="tab10", alpha=0.85)
        ax.set_title(title)
        ax.set_xlabel("Component 1")
        ax.set_ylabel("Component 2")
        fig.colorbar(scatter, ax=ax, fraction=0.046, pad=0.04)
        fig.tight_layout()
        return fig

    def _plot_pca_scatter(self, reduced, labels):
        fig, ax = plt.subplots(figsize=(5.5, 4.5))
        if reduced.shape[1] == 1:
            ax.scatter(reduced[:, 0], np.zeros_like(reduced[:, 0]), alpha=0.75)
            ax.set_ylabel("Zero Axis")
        else:
            if labels is not None:
                unique_labels = labels.astype(str)
                for label in sorted(unique_labels.unique())[:8]:
                    mask = unique_labels == label
                    ax.scatter(reduced[mask, 0], reduced[mask, 1], alpha=0.75, label=label)
                ax.legend(loc="best", fontsize=8)
            else:
                ax.scatter(reduced[:, 0], reduced[:, 1], alpha=0.75)
            ax.set_ylabel("PC2")
        ax.set_title("PCA Projection")
        ax.set_xlabel("PC1")
        fig.tight_layout()
        return fig