import gradio as gr
import pandas as pd

from sklearn.datasets import (
    load_iris,
    load_breast_cancer,
    fetch_california_housing,
    load_diabetes,
)

from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

# Classification Models
from sklearn.linear_model import LogisticRegression
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.neighbors import KNeighborsClassifier

# Regression Models
from sklearn.linear_model import LinearRegression
from sklearn.tree import DecisionTreeRegressor
from sklearn.ensemble import RandomForestRegressor
from sklearn.svm import SVR
from sklearn.svm import SVC

# Metrics
from sklearn.metrics import (
    accuracy_score,
    f1_score,
    mean_squared_error,
    r2_score,
)

# =====================================================
# TITANIC DATASET
# =====================================================

def load_titanic():
    url = "https://raw.githubusercontent.com/datasciencedojo/datasets/master/titanic.csv"
    df = pd.read_csv(url)

    df = df[["Pclass", "Sex", "Age", "Fare", "Survived"]]

    df["Age"] = df["Age"].fillna(df["Age"].mean())
    df["Sex"] = df["Sex"].map({"male": 0, "female": 1})

    X = df.drop("Survived", axis=1)
    y = df["Survived"]

    return X, y


# =====================================================
# BOSTON DATASET
# =====================================================

def load_boston():
    url = "https://raw.githubusercontent.com/selva86/datasets/master/BostonHousing.csv"
    df = pd.read_csv(url)

    X = df.drop("medv", axis=1)
    y = df["medv"]

    return X, y


# =====================================================
# MAIN FUNCTION
# =====================================================


def save_report(results_df, best_model, task_type, dataset_name):
    file_path = "model_report.txt"

    with open(file_path, "w", encoding="utf-8") as f:
        f.write("AI Model Comparison Report\n")
        f.write("=" * 40 + "\n\n")

        f.write(f"Task Type: {task_type}\n")
        f.write(f"Dataset: {dataset_name}\n\n")

        f.write("Results:\n")
        f.write(results_df.to_string(index=False))
        f.write("\n\n")

        f.write(f"Best Model: {best_model}\n")

    return file_path
    
def run_models(task_type, dataset_name):

    # =========================
    # CLASSIFICATION DATASETS
    # =========================

    if task_type == "Classification":

        if dataset_name == "Iris":
            data = load_iris()
            X = pd.DataFrame(data.data, columns=data.feature_names)
            y = data.target

        elif dataset_name == "Breast Cancer":
            data = load_breast_cancer()
            X = pd.DataFrame(data.data, columns=data.feature_names)
            y = data.target

        elif dataset_name == "Titanic":
            X, y = load_titanic()

        models = {
            "Logistic Regression": LogisticRegression(max_iter=1000),
            "SVM": SVC(),
            "Decision Tree": DecisionTreeClassifier(),
            "Random Forest": RandomForestClassifier(),
        }

    # =========================
    # REGRESSION DATASETS
    # =========================

    else:

        if dataset_name == "California Housing":
            data = fetch_california_housing()
            X = pd.DataFrame(data.data, columns=data.feature_names)
            y = data.target

        elif dataset_name == "Diabetes":
            data = load_diabetes()
            X = pd.DataFrame(data.data, columns=data.feature_names)
            y = data.target

        elif dataset_name == "Boston Housing":
            X, y = load_boston()

        models = {
            "Linear Regression": LinearRegression(),
            "SVR": SVR(),
            "Decision Tree": DecisionTreeRegressor(),
            "Random Forest": RandomForestRegressor(),
        }

    # =========================
    # SPLIT + SCALE
    # =========================

    X_train, X_test, y_train, y_test = train_test_split(
        X,
        y,
        test_size=0.2,
        random_state=42,
    )

    scaler = StandardScaler()

    X_train = scaler.fit_transform(X_train)
    X_test = scaler.transform(X_test)

    # =========================
    # TRAIN MODELS
    # =========================

    results = []

    for name, model in models.items():

        model.fit(X_train, y_train)

        predictions = model.predict(X_test)

        # Classification Metrics
        if task_type == "Classification":

            accuracy = accuracy_score(y_test, predictions)
            f1 = f1_score(y_test, predictions, average="weighted")

            results.append([name, accuracy, f1])

        # Regression Metrics
        else:

            mse = mean_squared_error(y_test, predictions)
            r2 = r2_score(y_test, predictions)

            results.append([name, mse, r2])

    # =========================
    # RESULTS TABLE
    # =========================

    if task_type == "Classification":

        results_df = pd.DataFrame(
            results,
            columns=["Model", "Accuracy", "F1 Score"]
        )

        best_model = results_df.loc[
            results_df["Accuracy"].idxmax(),
            "Model"
        ]

    else:

        results_df = pd.DataFrame(
            results,
            columns=["Model", "MSE", "R2 Score"]
        )

        best_model = results_df.loc[
            results_df["MSE"].idxmin(),
            "Model"
        ]

    report_file = save_report(results_df, best_model, task_type, dataset_name)
    return results_df, f"🏆 Best Model: {best_model}", report_file


# =====================================================
# UPDATE DATASET OPTIONS
# =====================================================

def update_datasets(task_type):

    if task_type == "Classification":

        return gr.Dropdown(
            choices=[
                "Iris",
                "Breast Cancer",
                "Titanic"
            ],
            value="Iris"
        )

    else:

        return gr.Dropdown(
            choices=[
                "California Housing",
                "Diabetes",
                "Boston Housing"
            ],
            value="California Housing"
        )


# =====================================================
# GRADIO UI
# =====================================================

with gr.Blocks() as demo:

    gr.Markdown("# AI Model Comparison App")

    task_type = gr.Radio(
        choices=["Classification", "Regression"],
        value="Classification",
        label="Select Task Type"
    )

    dataset_name = gr.Dropdown(
        choices=[
            "Iris",
            "Breast Cancer",
            "Titanic"
        ],
        value="Iris",
        label="Select Dataset"
    )

    task_type.change(
        fn=update_datasets,
        inputs=task_type,
        outputs=dataset_name
    )

    run_button = gr.Button("Run Models")

    output_table = gr.Dataframe()

    output_text = gr.Textbox()

    output_file = gr.File(label="Download Report")

    run_button.click(
    fn=run_models,
    inputs=[task_type, dataset_name],
    outputs=[output_table, output_text, output_file]
    )

demo.launch()