import gradio as gr
import matplotlib.pyplot as plt
from sklearn import datasets
from sklearn.tree import DecisionTreeClassifier, plot_tree
from io import BytesIO
from PIL import Image

# Load available datasets
dataset_names = ["iris", "wine", "breast_cancer", "digits"]
datasets_dict = {
    "iris": datasets.load_iris(),
    "wine": datasets.load_wine(),
    "breast_cancer": datasets.load_breast_cancer(),
    "digits": datasets.load_digits(),
}

# Define the function to visualize the decision tree
def visualize_decision_tree(dataset_name, max_depth, min_samples_split, min_samples_leaf, max_features, criterion, splitter, max_leaf_nodes, random_state):
    dataset = datasets_dict[dataset_name]
    X, y = dataset.data, dataset.target
    clf = DecisionTreeClassifier(max_depth=max_depth, min_samples_split=min_samples_split, min_samples_leaf=min_samples_leaf, max_features=max_features, criterion=criterion, splitter=splitter, max_leaf_nodes=max_leaf_nodes, random_state=random_state)
    clf.fit(X, y)
    fig, ax = plt.subplots(figsize=(10, 8))
    plot_tree(clf, feature_names=dataset.feature_names, class_names=dataset.target_names, filled=True, ax=ax)
    buf = BytesIO()
    fig.savefig(buf, format='png')
    buf.seek(0)
    image_data = buf.getvalue()
    image = Image.open(BytesIO(image_data))
    return image

# Define the hyperparameters and their ranges
max_depth_range = [None, 2, 3, 4, 5, 6, 7, 8, 9, 10]
min_samples_split_range = [2, 3, 4, 5, 6, 7, 8, 9, 10]
min_samples_leaf_range = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
max_features_range = [None, 'sqrt', 'log2', 0.1, 0.2, 0.3, 0.4, 0.5]
criterion_range = ['gini', 'entropy']
splitter_range = ['best', 'random']
max_leaf_nodes_range = [None, 2, 3, 4, 5, 6, 7, 8, 9, 10]
random_state_range = [None, 42, 100, 200, 300, 400, 500, 600]

# Create the Gradio interface
dataset_dropdown = gr.components.Dropdown(choices=dataset_names, label="Dataset", value="iris")
max_depth_dropdown = gr.components.Dropdown(choices=max_depth_range, label="Max Depth", value=None)
min_samples_split_dropdown = gr.components.Dropdown(choices=min_samples_split_range, label="Min Samples Split", value=2)
min_samples_leaf_dropdown = gr.components.Dropdown(choices=min_samples_leaf_range, label="Min Samples Leaf", value=1)
max_features_dropdown = gr.components.Dropdown(choices=max_features_range, label="Max Features", value=None)
criterion_dropdown = gr.components.Dropdown(choices=criterion_range, label="Criterion", value="gini")
splitter_dropdown = gr.components.Dropdown(choices=splitter_range, label="Splitter", value="best")
max_leaf_nodes_dropdown = gr.components.Dropdown(choices=max_leaf_nodes_range, label="Max Leaf Nodes", value=None)
random_state_dropdown = gr.components.Dropdown(choices=random_state_range, label="Random State", value=None)

iface = gr.Interface(
    fn=visualize_decision_tree,
    inputs=[dataset_dropdown, max_depth_dropdown, min_samples_split_dropdown, min_samples_leaf_dropdown, max_features_dropdown, criterion_dropdown, splitter_dropdown, max_leaf_nodes_dropdown, random_state_dropdown],
    outputs=gr.Image(type="pil"),
    title="Decision Tree Visualization",
    description="Visualize a decision tree classifier on various datasets by adjusting hyperparameters."
)

iface.launch()