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 tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Input
from tensorflow.keras.optimizers import Adam
import time

# 1. Helper function: Prepare data
def prepare_data(input_data, n_samples, outliers_fraction=0.01):
    n_outliers = max(int(outliers_fraction * n_samples), 1)
    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_samples, noise=0.05, random_state=0)[0] - np.array([0.5, 0.25])),
        "Noise": 14.0 * (np.random.RandomState(42).rand(n_samples, 2) - 0.5),
    }
    X = DATA_MAPPING[input_data]
    rng = np.random.RandomState(42)
    outliers = rng.uniform(low=-6, high=6, size=(n_outliers, 2))
    X = np.concatenate([X, outliers], axis=0)
    return X

# 2. Autoencoder Anomaly Detection
def build_autoencoder(input_dim):
    model = Sequential([
        Dense(64, activation='relu', input_dim=input_dim),
        Dense(32, activation='relu'),
        Dense(64, activation='relu'),
        Dense(input_dim, activation='sigmoid'),
    ])
    model.compile(optimizer=Adam(learning_rate=0.001), loss='mse')
    return model

def autoencoder_anomaly_detection(X, outliers_fraction=0.01, epochs=50):
    model = build_autoencoder(X.shape[1])
    model.fit(X, X, epochs=epochs, batch_size=32, verbose=0)
    reconstruction = model.predict(X)
    reconstruction_error = np.mean((X - reconstruction) ** 2, axis=1)
    threshold = np.percentile(reconstruction_error, 100 * (1 - outliers_fraction))
    y_pred = (reconstruction_error > threshold).astype(int)
    return y_pred

# 3. Function to train models and generate plots
def train_models(input_data, outliers_fraction, n_samples, clf_name):
    X = prepare_data(input_data, n_samples, outliers_fraction)
    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, random_state=42)
        ),
        "Isolation Forest": IsolationForest(contamination=outliers_fraction, random_state=42),
        "Local Outlier Factor": LocalOutlierFactor(n_neighbors=35, contamination=outliers_fraction),
        "Autoencoders": autoencoder_anomaly_detection(X, outliers_fraction)
    }
    if clf_name == "Autoencoders":
        y_pred = NAME_CLF_MAPPING[clf_name]
    else:
        clf = NAME_CLF_MAPPING[clf_name]
        if clf_name == "Local Outlier Factor":
            y_pred = clf.fit_predict(X)
        else:
            clf.fit(X)
            y_pred = clf.predict(X)
    plt.figure(figsize=(5, 5))
    colors = np.array(["#377eb8", "#ff7f00"])
    plt.scatter(X[:, 0], X[:, 1], c=colors[(y_pred + 1) // 2], s=20)
    plt.title(clf_name)
    return plt.gcf()

# 4. Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("## Anomaly Detection Comparison App")
    input_data = gr.Radio(choices=["Central Blob", "Two Blobs", "Blob with Noise", "Moons", "Noise"], value="Moons", label="Dataset")
    n_samples = gr.Slider(minimum=50, maximum=2000, step=50, value=500, label="Number of Samples")
    outliers_fraction = gr.Slider(minimum=0.01, maximum=0.5, step=0.01, value=0.05, label="Fraction of Outliers")
    input_models = gr.Radio(
        choices=["Robust covariance", "One-Class SVM", "One-Class SVM (SGD)", "Isolation Forest", "Local Outlier Factor", "Autoencoders"],
        value="Isolation Forest",
        label="Select Model"
    )
    plot = gr.Plot(label="Model Results")
    generate_plot = gr.Button("Generate Plot")
    generate_plot.click(fn=train_models, inputs=[input_data, outliers_fraction, n_samples, input_models], outputs=plot)

demo.launch()