import gradio as gr
import pandas as pd
import xgboost as xgb
import json

def train_and_predict(training_json, new_data_json, target_col):
    try:
        # 1. Parse Training Data
        train_data = json.loads(training_json)
        df_train = pd.DataFrame(train_data)
        
        if target_col not in df_train.columns:
            return {"error": f"Target column '{target_col}' not found in training data"}

        # Split into X (Features) and y (Target)
        X_train = df_train.drop(columns=[target_col])
        y_train = df_train[target_col]

        # 2. Train XGBoost (runs on server CPU)
        model = xgb.XGBRegressor()
        model.fit(X_train, y_train)

        # 3. Parse New Data for Prediction
        new_data = json.loads(new_data_json)
        df_new = pd.DataFrame(new_data)
        
        # Ensure columns match training data order
        # (Align columns to handle JSON key scrambling)
        existing_cols = [c for c in X_train.columns if c in df_new.columns]
        df_new = df_new[existing_cols]

        # 4. Predict
        predictions = model.predict(df_new)
        
        # 5. Format Output
        df_new['prediction'] = predictions.astype(float)
        return df_new.to_json(orient='records')

    except Exception as e:
        return json.dumps({"error": str(e)})

# We use Interface to guarantee the '/predict' API name exists
demo = gr.Interface(
    fn=train_and_predict,
    inputs=[
        gr.Textbox(label="Training Data (JSON)"),
        gr.Textbox(label="New Data (JSON)"),
        gr.Textbox(label="Target Column", value="reach")
    ],
    outputs="json"
)

demo.launch()