import gradio as gr
import joblib
from huggingface_hub import hf_hub_download

# ==================== LOAD MODEL ====================
def load_model_package():
    """Load the complete model package"""
    try:
        # Try to load the model from a local path first
        package = joblib.load("exoplanet_final_model.joblib")
        print("✅ Loaded model from local file.")
        return package
    except FileNotFoundError:
        print("ℹ️ Model file not found locally. Attempting to download from Hugging Face Hub...")
        try:
            # Download from Hugging Face Hub
            model_path = hf_hub_download(
                repo_id="Nugget-cloud/nasa-space-apps-exoplanet",
                filename="exoplanet_final_model.joblib"
            )
            package = joblib.load(model_path)
            print("✅ Model successfully downloaded and loaded from Hugging Face Hub.")
            return package
        except Exception as hub_e:
            print(f"❌ Failed to download or load model from Hugging Face Hub: {hub_e}")
            return None
    except Exception as e:
        print(f"❌ An unexpected error occurred while loading the model: {e}")
        return None

# Load model
package = load_model_package()
if package is None:
    raise SystemExit("Model could not be loaded. Please check your setup.")

model = package['ensemble_model']

# ==================== DEFINE PREDICTION FUNCTION ====================
def classify_exoplanet(
    koi_period, koi_duration, koi_depth, koi_prad, koi_teq, koi_insol, koi_steff
):
    """
    Run the trained model on the input features.
    """
    try:
        # Convert inputs to the expected 2D array format for the model
        features = [[
            float(koi_period),
            float(koi_duration),
            float(koi_depth),
            float(koi_prad),
            float(koi_teq),
            float(koi_insol),
            float(koi_steff)
        ]]
        prediction = model.predict(features)[0]
        label = "🌍 Likely Exoplanet" if prediction == 1 else "❌ Not an Exoplanet Candidate"
        return label
    except Exception as e:
        return f"Error during prediction: {e}"

# ==================== BUILD GRADIO INTERFACE ====================
with gr.Blocks(title="Exoplanet Classification 🌌") as demo:
    gr.Markdown(
        """
        # 🌌 **NASA Exoplanet Classifier**
        Enter the Kepler Object of Interest (KOI) parameters to classify whether it's likely an exoplanet.
        """
    )

    with gr.Row():
        koi_period = gr.Number(label="Orbital Period (days)", value=10.0)
        koi_duration = gr.Number(label="Transit Duration (hrs)", value=5.0)
        koi_depth = gr.Number(label="Transit Depth (ppm)", value=1000.0)
        koi_prad = gr.Number(label="Planet Radius (Earth radii)", value=1.0)
        koi_teq = gr.Number(label="Equilibrium Temperature (K)", value=300.0)
        koi_insol = gr.Number(label="Insolation Flux (Earth flux)", value=1.0)
        koi_steff = gr.Number(label="Stellar Effective Temperature (K)", value=5700.0)

    predict_btn = gr.Button("🚀 Classify")
    output = gr.Textbox(label="Prediction", placeholder="Result will appear here...")

    predict_btn.click(
        classify_exoplanet,
        inputs=[koi_period, koi_duration, koi_depth, koi_prad, koi_teq, koi_insol, koi_steff],
        outputs=output
    )

# ==================== LAUNCH ====================
if __name__ == "__main__":

    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True
    )