import os
import json
import joblib
import pandas as pd
import gradio as gr

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
MODELS_DIR = os.path.join(BASE_DIR, "models")

MODEL_CANDIDATES = [
    os.path.join(MODELS_DIR, "best_model.pkl"),
    os.path.join(MODELS_DIR, "best_model (1).pkl"),
]
FEATURE_FILE = os.path.join(MODELS_DIR, "feature_names.txt")
INFO_FILE = os.path.join(MODELS_DIR, "model_info.json")


def find_existing_file(paths):
    for path in paths:
        if os.path.exists(path):
            return path
    return None


def load_artifacts():
    model_path = find_existing_file(MODEL_CANDIDATES)
    if model_path is None:
        raise FileNotFoundError(
            f"Model file not found. Looked for: {', '.join(MODEL_CANDIDATES)}"
        )

    if not os.path.exists(FEATURE_FILE):
        raise FileNotFoundError(f"Feature file not found: {FEATURE_FILE}")

    model = joblib.load(model_path)

    with open(FEATURE_FILE, "r", encoding="utf-8") as f:
        feature_names = [line.strip() for line in f if line.strip()]

    model_info = {}
    if os.path.exists(INFO_FILE):
        with open(INFO_FILE, "r", encoding="utf-8") as f:
            model_info = json.load(f)

    return model, feature_names, model_info, model_path


model, feature_names, model_info, model_path = load_artifacts()


def predict_engine_condition(*values):
    input_df = pd.DataFrame([dict(zip(feature_names, values))])

    try:
        prediction = model.predict(input_df)[0]
        result_text = f"Predicted Engine Condition: {prediction}"

        if hasattr(model, "predict_proba"):
            proba = model.predict_proba(input_df)
            proba_df = pd.DataFrame(
                proba,
                columns=[f"Class {i}" for i in range(proba.shape[1])]
            )
            return result_text, proba_df

        return result_text, pd.DataFrame()

    except Exception as e:
        return f"Prediction failed: {e}", pd.DataFrame()


with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Engine Condition Prediction App")
    gr.Markdown("Enter sensor values to predict the engine condition.")

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Model Details")
            gr.Markdown(f"**Loaded model:** `{os.path.basename(model_path)}`")
            gr.Markdown(f"**Best model name:** {model_info.get('best_model_name', 'Unknown')}")
            gr.Markdown(f"**Test F1:** {model_info.get('test_f1', 'N/A')}")
            gr.Markdown(f"**Number of features:** {len(feature_names)}")

        with gr.Column(scale=2):
            gr.Markdown("### Input Features")

            inputs = []
            for feature in feature_names:
                inputs.append(
                    gr.Number(
                        label=feature,
                        value=0.0,
                        precision=4
                    )
                )

            predict_btn = gr.Button("Predict Engine Condition", variant="primary")

    gr.Markdown("### Prediction Output")
    result_box = gr.Textbox(label="Prediction", interactive=False)
    proba_box = gr.Dataframe(label="Prediction Probabilities")

    predict_btn.click(
        fn=predict_engine_condition,
        inputs=inputs,
        outputs=[result_box, proba_box]
    )

demo.launch()