import numpy as np
import matplotlib.pyplot as plt
from sklearn import svm
from sklearn.covariance import EllipticEnvelope
from sklearn.ensemble import IsolationForest
from sklearn.neighbors import LocalOutlierFactor
from sklearn.linear_model import SGDOneClassSVM
from sklearn.kernel_approximation import Nystroem
from sklearn.pipeline import make_pipeline
from sklearn.datasets import make_blobs, make_moons
import gradio as gr
import time

# Function to train models and generate plots
def train_models(input_data, outliers_fraction, n_samples, clf_name):
    # Prepare data
    n_outliers = int(outliers_fraction * n_samples)
    n_inliers = n_samples - n_outliers
    blobs_params = dict(random_state=0, n_samples=n_inliers, n_features=2)
    
    DATA_MAPPING = {
        "Central Blob": make_blobs(centers=[[0, 0], [0, 0]], cluster_std=0.5, **blobs_params)[0],
        "Two Blobs": make_blobs(centers=[[2, 2], [-2, -2]], cluster_std=[0.5, 0.5], **blobs_params)[0],
        "Blob with Noise": make_blobs(centers=[[2, 2], [-2, -2]], cluster_std=[1.5, 0.3], **blobs_params)[0],
        "Moons": 4.0 * (make_moons(n_samples=n_inliers, noise=0.05, random_state=0)[0] - np.array([0.5, 0.25])),
        "Noise": 14.0 * (np.random.RandomState(42).rand(n_inliers, 2) - 0.5),
    }
    
    NAME_CLF_MAPPING = {
        "Robust covariance": EllipticEnvelope(contamination=outliers_fraction),
        "One-Class SVM": svm.OneClassSVM(nu=outliers_fraction, kernel="rbf", gamma=0.1),
        "One-Class SVM (SGD)": make_pipeline(
            Nystroem(gamma=0.1, random_state=42, n_components=150),
            SGDOneClassSVM(
                nu=outliers_fraction,
                shuffle=True,
                fit_intercept=True,
                random_state=42,
                tol=1e-6,
            ),
        ),
        "Isolation Forest": IsolationForest(contamination=outliers_fraction, random_state=42),
        "Local Outlier Factor": LocalOutlierFactor(n_neighbors=35, contamination=outliers_fraction),
    }
    
    X = DATA_MAPPING[input_data]
    rng = np.random.RandomState(42)
    X_outliers = rng.uniform(low=-6, high=6, size=(n_outliers, 2))
    X = np.concatenate([X, X_outliers], axis=0)

    xx, yy = np.meshgrid(np.linspace(-7, 7, 150), np.linspace(-7, 7, 150))
    clf = NAME_CLF_MAPPING[clf_name]

    t0 = time.time()
    if clf_name == "Local Outlier Factor":
        y_pred = clf.fit_predict(X)
        # The decision_function is inverse of the LocalOutlierFactor._score_samples
        Z = clf._decision_function(np.c_[xx.ravel(), yy.ravel()])
    else:
        clf.fit(X)
        y_pred = clf.predict(X)
        Z = clf.decision_function(np.c_[xx.ravel(), yy.ravel()])
    t1 = time.time()

    # Plot
    plt.figure(figsize=(6, 6))
    Z = Z.reshape(xx.shape)
    plt.contourf(xx, yy, Z, levels=np.linspace(Z.min(), 0, 7), cmap=plt.cm.Blues_r)
    a = plt.contour(xx, yy, Z, levels=[0], linewidths=2, colors="red")
    plt.contourf(xx, yy, Z, levels=[0, Z.max()], colors="palevioletred")
    
    s = 20
    b1 = plt.scatter(X[:-n_outliers, 0], X[:-n_outliers, 1], c="white", s=s, edgecolors="k")
    b2 = plt.scatter(X[-n_outliers:, 0], X[-n_outliers:, 1], c="black", s=s, edgecolors="k")
    plt.axis("tight")
    plt.xlim((-7, 7))
    plt.ylim((-7, 7))
    plt.xticks(())
    plt.yticks(())
    plt.title(f"{clf_name} ({t1 - t0:.2f}s)")
    return plt.gcf()

# Gradio Interface
description = "Compare how different anomaly detection algorithms perform on various datasets."
title = "🕵️‍♀️ Compare Anomaly Detection Algorithms 🕵️‍♂️"

with gr.Blocks() as demo:
    gr.Markdown(f"## {title}")
    gr.Markdown(description)
    
    with gr.Row():
        with gr.Column(scale=1):
            # Inputs
            input_data = gr.Radio(
                choices=["Central Blob", "Two Blobs", "Blob with Noise", "Moons", "Noise"],
                value="Moons",
                label="Dataset"
            )
            n_samples = gr.Slider(minimum=100, maximum=500, step=25, value=300, label="Number of Samples")
            outliers_fraction = gr.Slider(minimum=0.1, maximum=0.9, step=0.1, value=0.2, label="Fraction of Outliers")
        with gr.Column(scale=3):
            # Models and their plots
            input_models = ["Robust covariance", "One-Class SVM", "One-Class SVM (SGD)", "Isolation Forest", "Local Outlier Factor"]
            plots = []
            with gr.Row():
                for model_name in input_models:
                    plot = gr.Plot(label=model_name)
                    plots.append((model_name, plot))
    
    # Update function
    def update(input_data, outliers_fraction, n_samples):
        results = []
        for clf_name, plot in plots:
            fig = train_models(input_data, outliers_fraction, n_samples, clf_name)
            results.append(fig)
            plt.close(fig)
        return results

    # Set change triggers
    inputs = [input_data, outliers_fraction, n_samples]
    demo_outputs = [plot for _, plot in plots]
    input_data.change(fn=update, inputs=inputs, outputs=demo_outputs)
    n_samples.change(fn=update, inputs=inputs, outputs=demo_outputs)
    outliers_fraction.change(fn=update, inputs=inputs, outputs=demo_outputs)
    
    # Initial update to display plots on load
    demo.load(fn=update, inputs=inputs, outputs=demo_outputs)

demo.launch(debug=True)