import gradio as gr
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from joblib import dump, load
import os
from huggingface_hub import HfApi
from huggingface_hub import login

login(token="botcoin")
huggingface-cli login

# Função para treinar e salvar o modelo
def train_model():
    # Carregar e dividir o dataset
    data = load_iris()
    X = data.data
    y = data.target
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # Treinando o modelo
    model = LogisticRegression(max_iter=200)
    model.fit(X_train, y_train)

    # Criar diretório para salvar o modelo
    diretorio = "/mnt/data"
    os.makedirs(diretorio, exist_ok=True)

    # Salvar o modelo em /mnt/data/
    model_filename = os.path.join(diretorio, "model.pkl")
    dump(model, model_filename)
    
    # Criar repositório no Hugging Face se não existir
    api = HfApi()
    repo_id = "coan/botClaiton"  # Substitua pelo seu repositório

    # Verificar se o repositório já existe
    try:
        api.repo_info(repo_id)
    except:
        # Criar o repositório caso não exista
        api.create_repo(repo_id, repo_type="model")

    # Upload do modelo para o repositório Hugging Face
    api.upload_file(
        path_or_fileobj=model_filename,
        path_in_repo="model.pkl",  # Nome do arquivo no repositório
        repo_id=repo_id,
        repo_type="model"
    )
    
    return f"Modelo treinado e salvo em: {model_filename}"

# Função para carregar o modelo treinado e fazer previsões
def predict(input_data):
    # Carregar o modelo treinado
    model = load("/mnt/data/model.pkl")  # Caminho do modelo salvo

    # Garantir que a entrada seja um vetor (as características)
    input_data = [float(i) for i in input_data.split()]
    
    # Realizar a previsão
    prediction = model.predict([input_data])  # input_data deve ser um vetor com as características
    return f"Previsão: {prediction[0]}"

# Interface para treinamento do modelo
train_iface = gr.Interface(
    fn=train_model, 
    inputs=[],  # Nenhuma entrada necessária para treinar
    outputs=["text"]  # Exibe uma mensagem de sucesso
)

# Interface para previsões
predict_iface = gr.Interface(
    fn=predict,                  # Função para chamar ao gerar previsão
    inputs=gr.Textbox(label="Input (4 features, espaço separado)"),  # Caixa de entrada de texto
    outputs=gr.Textbox(label="Output")     # Caixa de saída com a previsão
)

# Iniciar a interface de treinamento
train_iface.launch()

# Iniciar a interface de previsões
predict_iface.launch()