import numpy as np
import plotly.graph_objs as go
import plotly.figure_factory as ff
from sklearn.datasets import make_moons, make_circles, make_classification
from sklearn.neighbors import KNeighborsClassifier
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import confusion_matrix, roc_curve, auc

def generate_data(dataset_name, noise, random_state=42):
    if dataset_name == "moons":
        X, y = make_moons(n_samples=300, noise=noise, random_state=random_state)
    elif dataset_name == "circles":
        X, y = make_circles(n_samples=300, noise=noise, factor=0.5, random_state=random_state)
    else:
        X, y = make_classification(
            n_samples=300, n_features=2, n_redundant=0, n_informative=2,
            random_state=random_state, n_clusters_per_class=1, flip_y=noise/10.0
        )
    return X, y

def serve_prediction_plot(dataset_name, noise, n_neighbors, weights, metric, p_value, test_size=0.3):
    # 1. Fetch and Preprocess Data
    X, y = generate_data(dataset_name, noise)
    X = StandardScaler().fit_transform(X)
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=test_size, random_state=42)

    # 2. Fit KNN Model
    clf = KNeighborsClassifier(n_neighbors=n_neighbors, weights=weights, metric=metric, p=p_value)
    clf.fit(X_train, y_train)
    
    train_score = clf.score(X_train, y_train)
    test_score = clf.score(X_test, y_test)

    # 3. Create a Fine Mesh Grid for Contours
    x_min, x_max = X[:, 0].min() - 0.5, X[:, 0].max() + 0.5
    y_min, y_max = X[:, 1].min() - 0.5, X[:, 1].max() + 0.5
    h = 0.05
    xx, yy = np.meshgrid(np.arange(x_min, x_max, h), np.arange(y_min, y_max, h))

    grid_points = np.c_[xx.ravel(), yy.ravel()]
    if hasattr(clf, "predict_proba"):
        Z = clf.predict_proba(grid_points)[:, 1]
    else:
        Z = clf.predict(grid_points)
    Z = Z.reshape(xx.shape)

    # 4. Construct the Main Boundary Plot
    fig_boundary = go.Figure()
    fig_boundary.add_trace(go.Contour(
        x=np.arange(x_min, x_max, h), y=np.arange(y_min, y_max, h), z=Z,
        colorscale='RdBu', opacity=0.35, showscale=False, hoverinfo='skip'
    ))

    for label, color, name in [(0, '#FF4136', 'Train Class 0'), (1, '#0074D9', 'Train Class 1')]:
        mask = (y_train == label)
        fig_boundary.add_trace(go.Scatter(
            x=X_train[mask, 0], y=X_train[mask, 1], mode='markers',
            marker=dict(color=color, size=8, line=dict(width=1, color='black')), name=name
        ))

    for label, color, name in [(0, '#FF4136', 'Test Class 0'), (1, '#0074D9', 'Test Class 1')]:
        mask = (y_test == label)
        fig_boundary.add_trace(go.Scatter(
            x=X_test[mask, 0], y=X_test[mask, 1], mode='markers',
            marker=dict(color=color, size=10, symbol='diamond', line=dict(width=1.5, color='white')), name=name
        ))

    fig_boundary.update_layout(
        title="KNN Decision Space Map",
        xaxis=dict(title="Feature 1", showgrid=False),
        yaxis=dict(title="Feature 2", showgrid=False),
        legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1),
        margin=dict(l=40, r=40, t=80, b=40),
        plot_bgcolor='white', paper_bgcolor='white'
    )

    # 5. Construct the Confusion Matrix Heatmap
    y_pred = clf.predict(X_test)
    cm = confusion_matrix(y_test, y_pred)
    
    x_labels = ['Predicted 0', 'Predicted 1']
    y_labels = ['Actual 1', 'Actual 0']
    cm_text = [[str(y) for y in x] for x in cm]

    fig_cm = ff.create_annotated_heatmap(
        z=cm[::-1], x=x_labels, y=y_labels, annotation_text=cm_text[::-1], colorscale='Blues'
    )
    fig_cm.update_layout(
        title="Confusion Matrix (Test Data)",
        margin=dict(l=60, r=20, t=80, b=40),
        height=280
    )

    # 6. Construct the ROC Curve with AUC Calculation
    # Get predicted probabilities for class 1
    y_probs = clf.predict_proba(X_test)[:, 1] if hasattr(clf, "predict_proba") else y_pred
    fpr, tpr, thresholds = roc_curve(y_test, y_probs)
    roc_auc = auc(fpr, tpr)

    fig_roc = go.Figure()
    
    # Add Baseline Random Guess Line
    fig_roc.add_trace(go.Scatter(
        x=[0, 1], y=[0, 1], mode='lines', 
        line=dict(color='gray', width=1.5, dash='dash'), 
        name='Random Guess (AUC = 0.50)', hoverinfo='skip'
    ))
    
    # Add ROC Curve Line
    fig_roc.add_trace(go.Scatter(
        x=fpr, y=tpr, mode='lines+markers', 
        line=dict(color='#e11d48', width=3), 
        name=f'KNN Model (AUC = {roc_auc:.2f})',
        hovertext=[f"Threshold: {t:.2f}" for t in thresholds],
        hoverinfo='text+x+y'
    ))

    fig_roc.update_layout(
        title=f"ROC Curve (AUC: {roc_auc:.2f})",
        xaxis=dict(title="False Positive Rate (1 - Specificity)", range=[-0.02, 1.02], gridcolor='#f1f5f9'),
        yaxis=dict(title="True Positive Rate (Sensitivity / Recall)", range=[-0.02, 1.02], gridcolor='#f1f5f9'),
        margin=dict(l=50, r=20, t=80, b=50),
        plot_bgcolor='white', paper_bgcolor='white',
        height=280, showlegend=False
    )

    return fig_boundary, fig_cm, fig_roc, train_score, test_score