app.py

import io
import requests
import pandas as pd
import numpy as np
import gradio as gr
import json
import pickle
import os
from typing import Dict, List, Any

os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

print("🚀 Starting Eco Finder API...")

# Configuration
try:
    import tensorflow as tf
    print(f"✅ TensorFlow version: {tf.__version__}")
    from tensorflow.keras.models import load_model
    TENSORFLOW_AVAILABLE = True
except ImportError as e:
    print(f"❌ TensorFlow not available: {e}")
    TENSORFLOW_AVAILABLE = False

# Load resources
def load_resources():
    try:
        with open("feature_stats.json", "r") as f:
            feature_stats = json.load(f)
        print("✅ Feature stats loaded")
            
        with open("scaler.pkl", "rb") as f:
            scaler = pickle.load(f)
        print("✅ Scaler loaded")
            
        with open("label_encoder.pkl", "rb") as f:
            label_encoder = pickle.load(f)
        print("✅ Label encoder loaded")
        
        model = None
        if TENSORFLOW_AVAILABLE:
            model = load_model("modulo_tabular.h5")
            print("✅ Model loaded")
        
        return model, scaler, label_encoder, feature_stats
        
    except Exception as e:
        print(f"❌ Error loading resources: {str(e)}")
        feature_stats = {
            "feature_columns": [
                "koi_period", "koi_duration", "koi_depth", "koi_prad", 
                "koi_srad", "koi_teq", "koi_steff", "koi_slogg", 
                "koi_smet", "koi_kepmag", "koi_model_snr", "koi_num_transits"
            ],
            "train_medians": {
                "koi_period": 10.0, "koi_duration": 5.0, "koi_depth": 1000.0,
                "koi_prad": 2.0, "koi_srad": 1.0, "koi_teq": 1000.0,
                "koi_steff": 6000.0, "koi_slogg": 4.5, "koi_smet": 0.0,
                "koi_kepmag": 12.0, "koi_model_snr": 10.0, "koi_num_transits": 3.0
            }
        }
        return None, None, None, feature_stats

# Load resources
model, scaler, label_encoder, feature_stats = load_resources()
feature_columns = feature_stats.get("feature_columns", [])
train_medians = feature_stats.get("train_medians", {})

BASE = "https://exoplanetarchive.ipac.caltech.edu/cgi-bin/nstedAPI/nph-nstedAPI"

# ==================== FUNCTIONS FOR GRADIO ====================

def predict_single(features: Dict) -> Dict:
    """Function to predict a single object - USED BY GRADIO"""
    try:
        if model is None or scaler is None or label_encoder is None:
            return {"error": "Model not available"}
        
        # Create feature array
        input_features = []
        for feature in feature_columns:
            value = features.get(feature, train_medians.get(feature, 0))
            input_features.append(float(value))
        
        # Predict
        input_array = np.array([input_features])
        X_input = scaler.transform(input_array)
        
        if TENSORFLOW_AVAILABLE:
            probs = model.predict(X_input, verbose=0)[0]
        else:
            probs = np.random.dirichlet(np.ones(3), size=1)[0]
        
        pred_idx = np.argmax(probs)
        pred_label = label_encoder.inverse_transform([pred_idx])[0]
        
        return {
            "prediction": pred_label,
            "probabilities": {
                "CONFIRMED": float(probs[0]),
                "CANDIDATE": float(probs[1]),
                "FALSE_POSITIVE": float(probs[2])
            },
            "input_features": dict(zip(feature_columns, input_features))
        }
        
    except Exception as e:
        return {"error": str(e)}

def predict_from_dict(
    koi_period: float, koi_duration: float, koi_depth: float, 
    koi_prad: float, koi_srad: float, koi_teq: float,
    koi_steff: float, koi_slogg: float, koi_smet: float,
    koi_kepmag: float, koi_model_snr: float, koi_num_transits: float
) -> Dict:
    """Wrapper that takes individual parameters and converts them to dict"""
    features = {
        "koi_period": koi_period,
        "koi_duration": koi_duration,
        "koi_depth": koi_depth,
        "koi_prad": koi_prad,
        "koi_srad": koi_srad,
        "koi_teq": koi_teq,
        "koi_steff": koi_steff,
        "koi_slogg": koi_slogg,
        "koi_smet": koi_smet,
        "koi_kepmag": koi_kepmag,
        "koi_model_snr": koi_model_snr,
        "koi_num_transits": koi_num_transits
    }
    return predict_single(features)

def predict_koi_realtime():
    """Function for real-time KOI"""
    try:
        if model is None or scaler is None or label_encoder is None:
            return "❌ Model not available"
        
        # Query exoplanet API
        where = ("(tfopwg_disp like 'PC' or tfopwg_disp like 'APC') "
                 "and (pl_orbper is not null or tce_period is not null)")
        
        params = {"table": "koi", "where": where, "format": "csv"}
        resp = requests.get(BASE, params=params, timeout=60)
        resp.raise_for_status()
        koi_df = pd.read_csv(io.StringIO(resp.text))
        
        if koi_df.empty:
            return "❌ No KOI objects found"
        
        # Take sample
        koi_sample = koi_df.sample(min(3, len(koi_df)), random_state=7)
        koi_sample.columns = [c.strip().lower() for c in koi_sample.columns]
        
        # Synonym mapping
        candidates_map = {
            "koi_period": ["pl_orbper", "tce_period", "orbper", "period"],
            "koi_duration": ["pl_trandurh", "tce_duration", "tran_dur", "trandur", "duration", "dur"],
            "koi_depth": ["pl_trandep", "tce_depth", "depth", "trandep"],
            "koi_prad": ["pl_rade", "prad", "rade", "planet_radius"],
            "koi_srad": ["st_rad", "srad", "stellar_radius", "star_radius"],
            "koi_teq": ["pl_eqt", "teq", "equilibrium_temp"],
            "koi_steff": ["st_teff", "teff", "stellar_teff", "effective_temp"],
            "koi_slogg": ["st_logg", "logg", "slogg"],
            "koi_smet": ["st_met", "feh", "metallicity", "smet"],
            "koi_kepmag": ["st_tmag", "tmag", "kepmag", "koi_kepmag"],
            "koi_model_snr": ["tce_model_snr", "model_snr", "snr"],
            "koi_num_transits": ["tce_num_transits", "num_transits", "ntransits", "tran_count"]
        }
        
        def first_present(candidates, cols_set):
            for name in candidates:
                if name in cols_set:
                    return name
            for name in candidates:
                found = [c for c in cols_set if name in c]
                if found:
                    return found[0]
            return None
        
        cols_set = set(koi_sample.columns)
        results = []
        
        for idx, row in koi_sample.iterrows():
            # Prepare features
            features = {}
            for feat in feature_columns:
                src = first_present(candidates_map.get(feat, []), cols_set)
                if src and src in row and pd.notna(row[src]):
                    features[feat] = float(row[src])
                else:
                    features[feat] = train_medians.get(feat, 0)
            
            # Predict
            result = predict_single(features)
            
            if "error" not in result:
                results.append({
                    "KOI": row.get('koi', f"KOI-{idx}"),
                    "Disposition": row.get('tfopwg_disp', 'Unknown'),
                    "Prediction": result['prediction'],
                    "P(Confirmed)": f"{result['probabilities']['CONFIRMED']:.3f}",
                    "P(Candidate)": f"{result['probabilities']['CANDIDATE']:.3f}",
                    "P(False Positive)": f"{result['probabilities']['FALSE_POSITIVE']:.3f}"
                })
        
        if not results:
            return "❌ Could not generate predictions"
        
        result_df = pd.DataFrame(results)
        return f"**KOI Predictions:**\n\n{result_df.to_markdown(index=False)}"
        
    except Exception as e:
        return f"❌ Error: {str(e)}"

def predict_manual(
    period, duration, depth, prad, srad, teq, steff, slogg, smet, kepmag, snr, num_transits
):
    """Function for manual prediction in Gradio"""
    try:
        result = predict_from_dict(
            period, duration, depth, prad, srad, teq, steff, slogg, smet, kepmag, snr, num_transits
        )
        
        if "error" in result:
            return f"❌ {result['error']}"
        
        output = f"**Prediction:** {result['prediction']}\n\n**Probabilities:**\n"
        for clase, prob in result['probabilities'].items():
            output += f"- {clase}: {prob:.3f}\n"
            
        return output
        
    except Exception as e:
        return f"❌ Error: {str(e)}"

# ==================== GRADIO INTERFACE ====================

with gr.Blocks(theme=gr.themes.Soft(), title="Eco Finder API") as demo:
    gr.Markdown("# 🌌 Eco Finder API")
    gr.Markdown("Exoplanet classifier")
    
    with gr.Tab("🎯 API Prediction"):
        gr.Markdown("### Endpoint for frontend consumption")
        gr.Markdown("""
        **URL:** `https://jarpalucas-echo-finder-api.hf.space/api/predict`
        
        **Method:** POST
        **Content-Type:** application/json
        
        **Usage example with curl:**
        ```bash
        curl -X POST "https://jarpalucas-echo-finder-api.hf.space/api/predict" \\
          -H "Content-Type: application/json" \\
          -d '{
            "koi_period": 10.0,
            "koi_duration": 5.0,
            "koi_depth": 1000.0,
            "koi_prad": 2.0,
            "koi_srad": 1.0,
            "koi_teq": 1000.0,
            "koi_steff": 6000.0,
            "koi_slogg": 4.5,
            "koi_smet": 0.0,
            "koi_kepmag": 12.0,
            "koi_model_snr": 10.0,
            "koi_num_transits": 3.0
          }'
        ```
        """)
        
        # Inputs to test the API locally
        with gr.Row():
            with gr.Column():
                period = gr.Number(label="koi_period", value=10.0)
                duration = gr.Number(label="koi_duration", value=5.0)
                depth = gr.Number(label="koi_depth", value=1000.0)
                prad = gr.Number(label="koi_prad", value=2.0)
            with gr.Column():
                srad = gr.Number(label="koi_srad", value=1.0)
                teq = gr.Number(label="koi_teq", value=1000.0)
                steff = gr.Number(label="koi_steff", value=6000.0)
                slogg = gr.Number(label="koi_slogg", value=4.5)
            with gr.Column():
                smet = gr.Number(label="koi_smet", value=0.0)
                kepmag = gr.Number(label="koi_kepmag", value=12.0)
                snr = gr.Number(label="koi_model_snr", value=10.0)
                num_transits = gr.Number(label="koi_num_transits", value=3.0)
        
        api_btn = gr.Button("🚀 Test Prediction")
        api_output = gr.JSON()
        
        api_btn.click(
            fn=predict_from_dict,
            inputs=[period, duration, depth, prad, srad, teq, steff, slogg, smet, kepmag, snr, num_transits],
            outputs=api_output
        )
    
    with gr.Tab("🔭 Real-time KOI"):
        gr.Markdown("Real-time KOI object predictions")
        koi_btn = gr.Button("🔍 Analyze KOI")
        koi_output = gr.Markdown()
        koi_btn.click(predict_koi_realtime, outputs=koi_output)
    
    with gr.Tab("📊 Manual Interface"):
        gr.Markdown("Manual interface for predictions")
        manual_btn = gr.Button("🎯 Predict")
        manual_output = gr.Markdown()
        manual_btn.click(
            fn=predict_manual,
            inputs=[period, duration, depth, prad, srad, teq, steff, slogg, smet, kepmag, snr, num_transits],
            outputs=manual_output
        )

print("🎉 Application started successfully!")
print("🌐 Interface available at: /")
print("🔗 API endpoint available at: /api/predict")

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)