import mlflow
import mlflow.sklearn
from sklearn.datasets import load_diabetes
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error
from pyngrok import ngrok
import gradio as gr

# MLflow setup
mlflow.set_tracking_uri("./mlruns")  # Local directory for tracking
mlflow.set_experiment("House Price Prediction")

# Training function
def train_and_log_model():
    # Load dataset
    data = load_diabetes()
    X = data.data
    y = data.target

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

    # Train model
    model = LinearRegression()
    model.fit(X_train, y_train)

    # Predict and evaluate
    y_pred = model.predict(X_test)
    mse = mean_squared_error(y_test, y_pred)

    # Log to MLflow
    with mlflow.start_run():
        mlflow.log_param("model", "Linear Regression")
        mlflow.log_metric("mse", mse)
        mlflow.sklearn.log_model(model, "model")

    return mse, "Model training complete and logged to MLflow!"

# Start MLflow UI with Ngrok
def start_mlflow_ui():
    public_url = ngrok.connect(5000)  # Expose the MLflow UI running on port 5000
    mlflow_command = "mlflow ui --host 0.0.0.0 --port 5000"
    return_code = os.system(mlflow_command)
    if return_code != 0:
        return "Error: Unable to start MLflow UI."
    return f"MLflow UI is accessible at {public_url}"

# Gradio Interface Functions
def train_model():
    mse, message = train_and_log_model()
    return f"MSE: {mse}\n{message}"

def get_mlflow_ui_link():
    public_url = start_mlflow_ui()
    return public_url

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## House Price Prediction with MLflow")
    train_btn = gr.Button("Train Model and Log to MLflow")
    mlflow_btn = gr.Button("Start MLflow UI")
    output = gr.Textbox(label="Output")
    
    train_btn.click(train_model, inputs=[], outputs=output)
    mlflow_btn.click(get_mlflow_ui_link, inputs=[], outputs=output)

# Launch Gradio App
if __name__ == "__main__":
    demo.launch()