import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
from sklearn.metrics import (
    accuracy_score,
    precision_score,
    recall_score,
    f1_score,
    confusion_matrix,
    classification_report
)
import gradio as gr

# -------- Data --------
iris = load_iris()
X_full = pd.DataFrame(iris.data, columns=iris.feature_names)
y_full = pd.Series(iris.target, name="target")
class_names = iris.target_names

FEATURE_CHOICES = list(X_full.columns)

# -------- Plot helpers --------
def plot_confusion_matrix(cm, class_names):
    fig, ax = plt.subplots(figsize=(5, 4), dpi=120)
    im = ax.imshow(cm, interpolation="nearest")
    ax.figure.colorbar(im, ax=ax)
    ax.set(
        xticks=np.arange(cm.shape[1]),
        yticks=np.arange(cm.shape[0]),
        xticklabels=class_names,
        yticklabels=class_names,
        ylabel="True label",
        xlabel="Predicted label",
        title="Confusion Matrix",
    )

    # Show counts on cells
    thresh = cm.max() / 2.0
    for i in range(cm.shape[0]):
        for j in range(cm.shape[1]):
            ax.text(
                j, i, format(cm[i, j], "d"),
                ha="center", va="center",
                color="white" if cm[i, j] > thresh else "black"
            )
    fig.tight_layout()
    return fig

def plot_feature_importances(model, feature_names):
    importances = model.feature_importances_
    order = np.argsort(importances)[::-1]
    fig, ax = plt.subplots(figsize=(6, 4), dpi=120)
    ax.bar(range(len(importances)), importances[order])
    ax.set_xticks(range(len(importances)))
    ax.set_xticklabels([feature_names[i] for i in order], rotation=30, ha="right")
    ax.set_ylabel("Importance")
    ax.set_title("Feature Importances")
    fig.tight_layout()
    return fig

def plot_decision_regions_2d(model, X, y, feature_x_name, feature_y_name, class_names):
    # X: dataframe with exactly two columns (selected features)
    # Create a mesh
    x_min, x_max = X.iloc[:, 0].min() - 0.5, X.iloc[:, 0].max() + 0.5
    y_min, y_max = X.iloc[:, 1].min() - 0.5, X.iloc[:, 1].max() + 0.5
    xx, yy = np.meshgrid(
        np.linspace(x_min, x_max, 300),
        np.linspace(y_min, y_max, 300)
    )
    grid = np.c_[xx.ravel(), yy.ravel()]
    Z = model.predict(grid).reshape(xx.shape)

    fig, ax = plt.subplots(figsize=(6, 5), dpi=120)
    ax.contourf(xx, yy, Z, alpha=0.2)

    # Scatter original points
    for idx, cname in enumerate(class_names):
        mask = (y == idx)
        ax.scatter(
            X.loc[mask, feature_x_name],
            X.loc[mask, feature_y_name],
            label=cname, s=24
        )

    ax.set_xlabel(feature_x_name)
    ax.set_ylabel(feature_y_name)
    ax.set_title("Decision Regions (2 features)")
    ax.legend(loc="upper right", fontsize=8)
    fig.tight_layout()
    return fig

# -------- Core training + evaluation --------
def run_decision_tree(
    test_size,
    random_state,
    criterion,
    splitter,
    unlimited_depth,
    max_depth,
    min_samples_split,
    min_samples_leaf,
    max_features,
    class_weight_mode,
    feature_x_name,
    feature_y_name
):
    # Map UI values to sklearn-friendly options
    md = None if unlimited_depth else int(max_depth)
    if max_features == "None":
        mf = None
    elif max_features == "auto/sqrt":
        mf = "sqrt"
    elif max_features == "log2":
        mf = "log2"
    else:
        mf = None

    cw = None if class_weight_mode == "None" else "balanced"

    # Train / Test split
    X_train, X_test, y_train, y_test = train_test_split(
        X_full, y_full,
        test_size=float(test_size),
        random_state=int(random_state),
        stratify=y_full
    )

    # Model
    clf = DecisionTreeClassifier(
        criterion=criterion,
        splitter=splitter,
        max_depth=md,
        min_samples_split=int(min_samples_split),
        min_samples_leaf=int(min_samples_leaf),
        max_features=mf,
        class_weight=cw,
        random_state=int(random_state)
    )
    clf.fit(X_train, y_train)

    # Predictions & metrics
    y_pred = clf.predict(X_test)
    acc = accuracy_score(y_test, y_pred)
    prec = precision_score(y_test, y_pred, average="macro", zero_division=0)
    rec = recall_score(y_test, y_pred, average="macro", zero_division=0)
    f1 = f1_score(y_test, y_pred, average="macro", zero_division=0)
    report = classification_report(y_test, y_pred, target_names=class_names, zero_division=0)

    # Confusion matrix
    cm = confusion_matrix(y_test, y_pred)
    cm_fig = plot_confusion_matrix(cm, class_names)

    # Feature importances
    fi_fig = plot_feature_importances(clf, X_full.columns)

    # Decision boundary for 2 chosen features
    # Refit a new tree **on those two features only** to plot clear regions (same hyperparams)
    if feature_x_name == feature_y_name:
        # If same feature accidentally chosen, pick a safe default distinct pair
        feature_x_name, feature_y_name = FEATURE_CHOICES[0], FEATURE_CHOICES[1]

    two_feat_cols = [feature_x_name, feature_y_name]
    X2_train = X_train[two_feat_cols]
    X2_test = X_test[two_feat_cols]

    clf2 = DecisionTreeClassifier(
        criterion=criterion,
        splitter=splitter,
        max_depth=md,
        min_samples_split=int(min_samples_split),
        min_samples_leaf=int(min_samples_leaf),
        max_features=None,  # force 2D features for plotting
        class_weight=cw,
        random_state=int(random_state)
    )
    clf2.fit(X2_train, y_train)
    boundary_fig = plot_decision_regions_2d(
        clf2,
        pd.concat([X2_train, X2_test], axis=0),
        pd.concat([y_train, y_test], axis=0).values,
        feature_x_name, feature_y_name, class_names
    )

    # Make a small metrics dataframe for display
    metrics_df = pd.DataFrame(
        [{"accuracy": acc, "precision_macro": prec, "recall_macro": rec, "f1_macro": f1}]
    )

    # Classification report as preformatted text
    report_text = f"```\n{report}\n```"

    return (
        metrics_df,
        report_text,
        cm_fig,
        fi_fig,
        boundary_fig
    )

# -------- Gradio UI --------
with gr.Blocks(title="Iris Decision Tree Explorer") as demo:
    gr.Markdown(
        """
        # 🌸 Iris Decision Tree Explorer
        Train a Decision Tree on the classic Iris dataset. Tweak hyperparameters on the left, then inspect metrics and plots on the right.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Hyperparameters")

            criterion = gr.Dropdown(
                label="criterion",
                choices=["gini", "entropy", "log_loss"],
                value="gini"
            )
            splitter = gr.Dropdown(
                label="splitter",
                choices=["best", "random"],
                value="best"
            )

            unlimited_depth = gr.Checkbox(
                label="Unlimited depth (ignore max_depth)",
                value=True
            )
            max_depth = gr.Slider(
                label="max_depth (used only if Unlimited depth = False)",
                minimum=1, maximum=30, step=1, value=5
            )

            min_samples_split = gr.Slider(
                label="min_samples_split",
                minimum=2, maximum=20, step=1, value=2
            )
            min_samples_leaf = gr.Slider(
                label="min_samples_leaf",
                minimum=1, maximum=20, step=1, value=1
            )
            max_features = gr.Dropdown(
                label="max_features",
                choices=["None", "auto/sqrt", "log2"],
                value="None"
            )
            class_weight_mode = gr.Dropdown(
                label="class_weight",
                choices=["None", "balanced"],
                value="None"
            )

            gr.Markdown("### Data & Reproducibility")
            test_size = gr.Slider(
                label="test_size",
                minimum=0.1, maximum=0.5, step=0.05, value=0.2
            )
            random_state = gr.Number(
                label="random_state",
                value=42, precision=0
            )

            gr.Markdown("### Decision Region (2D) Features")
            feature_x_name = gr.Dropdown(
                label="X-axis feature",
                choices=FEATURE_CHOICES,
                value=FEATURE_CHOICES[0]
            )
            feature_y_name = gr.Dropdown(
                label="Y-axis feature",
                choices=FEATURE_CHOICES,
                value=FEATURE_CHOICES[1]
            )

            run_btn = gr.Button("Train & Evaluate", variant="primary")

        with gr.Column(scale=2):
            gr.Markdown("### Results")
            metrics_df = gr.Dataframe(
                label="Metrics (test set)",
                interactive=False,
                headers=["accuracy", "precision_macro", "recall_macro", "f1_macro"]
            )
            report_md = gr.Markdown(label="Classification Report")

            with gr.Row():
                cm_plot = gr.Plot(label="Confusion Matrix")
                fi_plot = gr.Plot(label="Feature Importances")

            boundary_plot = gr.Plot(label="Decision Regions (2 features)")

    run_btn.click(
        fn=run_decision_tree,
        inputs=[
            test_size, random_state,
            criterion, splitter, unlimited_depth, max_depth,
            min_samples_split, min_samples_leaf, max_features, class_weight_mode,
            feature_x_name, feature_y_name
        ],
        outputs=[metrics_df, report_md, cm_plot, fi_plot, boundary_plot]
    )

if __name__ == "__main__":
    # You can set share=True if you want a public link when running locally
    demo.launch()