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_toi_realtime():
    """Function for real-time TOI"""
    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": "toi", "where": where, "format": "csv"}
        resp = requests.get(BASE, params=params, timeout=60)
        resp.raise_for_status()
        toi_df = pd.read_csv(io.StringIO(resp.text))

        if toi_df.empty:
            return "❌ No TOI objects found"

        # Take sample
        toi_sample = toi_df.sample(min(3, len(toi_df)), random_state=7)
        toi_sample.columns = [c.strip().lower() for c in toi_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(toi_sample.columns)
        results = []

        for idx, row in toi_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(
                    {
                        "TOI": row.get("toi", f"tOI-{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"**TOI 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 TOI"):
        gr.Markdown("Real-time TOI object predictions")
        toi_btn = gr.Button("🔍 Analyze TOI")
        toi_output = gr.Markdown()
        toi_btn.click(predict_toi_realtime, outputs=toi_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)