api.py

#!/usr/bin/env python3
"""
LARUN TinyML - FastAPI REST Service
====================================
REST API for astronomical data analysis.

Usage:
    uvicorn api:app --reload --port 8000

Endpoints:
    GET  /                    - API info
    GET  /health              - Health check
    POST /analyze/bls         - BLS periodogram analysis
    POST /analyze/fit         - Transit model fitting
    POST /stellar/classify    - Stellar classification
    POST /planet/radius       - Planet radius calculation
    POST /planet/hz           - Habitable zone calculation
    POST /pipeline            - Full analysis pipeline

Created by: Padmanaban Veeraragavalu (Larun Engineering)
With AI assistance from: Claude (Anthropic)
"""

from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from typing import Optional, List, Dict, Any
import numpy as np
import sys
from pathlib import Path
from datetime import datetime
import logging

# Add src to path
sys.path.insert(0, str(Path(__file__).parent / 'src'))

# ============================================================================
# App Configuration
# ============================================================================

app = FastAPI(
    title="LARUN TinyML API",
    description="REST API for astronomical data analysis and exoplanet detection",
    version="2.0.0",
    contact={
        "name": "Padmanaban Veeraragavalu",
        "email": "larun@example.com"
    },
    license_info={
        "name": "MIT"
    }
)

# CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


# ============================================================================
# Request/Response Models
# ============================================================================

class LightCurveData(BaseModel):
    """Light curve input data."""
    time: List[float] = Field(..., description="Time array (days)")
    flux: List[float] = Field(..., description="Normalized flux array")
    flux_err: Optional[List[float]] = Field(None, description="Flux uncertainties")


class BLSRequest(BaseModel):
    """BLS analysis request."""
    data: LightCurveData
    min_period: float = Field(0.5, description="Minimum period to search (days)")
    max_period: float = Field(50.0, description="Maximum period to search (days)")
    min_snr: float = Field(7.0, description="Minimum SNR for detection")


class TransitFitRequest(BaseModel):
    """Transit fitting request."""
    data: LightCurveData
    period: float = Field(..., description="Orbital period (days)")
    t0: Optional[float] = Field(None, description="Initial mid-transit time")
    stellar_teff: float = Field(5778, description="Stellar Teff (K)")


class StellarClassifyRequest(BaseModel):
    """Stellar classification request."""
    teff: float = Field(..., description="Effective temperature (K)")
    logg: Optional[float] = Field(None, description="Surface gravity (log g)")
    metallicity: Optional[float] = Field(None, description="[Fe/H]")


class PlanetRadiusRequest(BaseModel):
    """Planet radius calculation request."""
    depth_ppm: float = Field(..., description="Transit depth (ppm)")
    stellar_radius: float = Field(..., description="Stellar radius (R_sun)")
    period: Optional[float] = Field(None, description="Orbital period (days)")
    stellar_mass: Optional[float] = Field(None, description="Stellar mass (M_sun)")
    stellar_teff: Optional[float] = Field(None, description="Stellar Teff (K)")
    stellar_luminosity: Optional[float] = Field(None, description="Stellar luminosity (L_sun)")


class HabitableZoneRequest(BaseModel):
    """Habitable zone calculation request."""
    stellar_teff: float = Field(..., description="Stellar Teff (K)")
    stellar_luminosity: float = Field(..., description="Stellar luminosity (L_sun)")


class PipelineRequest(BaseModel):
    """Full pipeline request."""
    target: str = Field(..., description="Target name or TIC ID")
    quick_mode: bool = Field(False, description="Quick analysis mode")


class MultiPlanetRequest(BaseModel):
    """Multi-planet detection request."""
    data: LightCurveData
    target: str = Field("Unknown", description="Target name")
    max_planets: int = Field(5, description="Maximum planets to search for")
    min_snr: float = Field(7.0, description="Minimum SNR")


# ============================================================================
# API Endpoints
# ============================================================================

@app.get("/")
async def root():
    """API information."""
    return {
        "name": "LARUN TinyML API",
        "version": "2.0.0",
        "description": "Astronomical data analysis for exoplanet detection",
        "endpoints": {
            "/analyze/bls": "BLS periodogram analysis",
            "/analyze/fit": "Transit model fitting",
            "/analyze/multiplanet": "Multi-planet detection",
            "/stellar/classify": "Stellar classification",
            "/planet/radius": "Planet radius calculation",
            "/planet/hz": "Habitable zone calculation",
            "/pipeline": "Full analysis pipeline"
        }
    }


@app.get("/health")
async def health_check():
    """Health check endpoint."""
    return {
        "status": "healthy",
        "timestamp": datetime.now().isoformat(),
        "skills_available": [
            "BLS Periodogram",
            "Transit Fitting",
            "Stellar Classification",
            "Planet Radius",
            "Habitable Zone",
            "Multi-Planet Detection"
        ]
    }


@app.post("/analyze/bls")
async def analyze_bls(request: BLSRequest):
    """
    Run BLS periodogram analysis.

    Returns detected transit candidates with periods, depths, and SNRs.
    """
    try:
        from skills.periodogram import BLSPeriodogram

        time = np.array(request.data.time)
        flux = np.array(request.data.flux)
        flux_err = np.array(request.data.flux_err) if request.data.flux_err else None

        bls = BLSPeriodogram(
            min_period=request.min_period,
            max_period=request.max_period
        )

        result = bls.compute(time, flux, flux_err, min_snr=request.min_snr)

        return {
            "status": "success",
            "method": result.method,
            "best_period": result.best_period,
            "best_power": result.best_power,
            "fap": result.fap,
            "n_candidates": len(result.candidates),
            "candidates": [c.to_dict() for c in result.candidates]
        }

    except Exception as e:
        logger.error(f"BLS analysis failed: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/analyze/fit")
async def analyze_fit(request: TransitFitRequest):
    """
    Fit transit model to light curve.

    Returns best-fit transit parameters (Rp/Rs, a/Rs, inclination, etc.).
    """
    try:
        from skills.transit_fit import TransitFitter

        time = np.array(request.data.time)
        flux = np.array(request.data.flux)
        flux_err = np.array(request.data.flux_err) if request.data.flux_err else None

        fitter = TransitFitter()
        result = fitter.fit(
            time, flux, flux_err,
            period=request.period,
            t0=request.t0,
            stellar_teff=request.stellar_teff
        )

        return {
            "status": "success",
            "fit_result": result.to_dict()
        }

    except Exception as e:
        logger.error(f"Transit fitting failed: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/analyze/multiplanet")
async def analyze_multiplanet(request: MultiPlanetRequest):
    """
    Detect multiple planets via iterative BLS.

    Returns all detected planet candidates.
    """
    try:
        from skills.multiplanet import MultiPlanetDetector

        time = np.array(request.data.time)
        flux = np.array(request.data.flux)
        flux_err = np.array(request.data.flux_err) if request.data.flux_err else None

        detector = MultiPlanetDetector(
            max_planets=request.max_planets,
            min_snr=request.min_snr
        )

        result = detector.detect(time, flux, flux_err, target=request.target)

        return {
            "status": "success",
            "result": result.to_dict()
        }

    except Exception as e:
        logger.error(f"Multi-planet detection failed: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/stellar/classify")
async def stellar_classify(request: StellarClassifyRequest):
    """
    Classify star by spectral type.

    Returns OBAFGKM classification and derived parameters.
    """
    try:
        from skills.stellar import StellarClassifier

        classifier = StellarClassifier()
        result = classifier.classify_from_teff(
            request.teff,
            logg=request.logg,
            metallicity=request.metallicity
        )

        return {
            "status": "success",
            "classification": result.to_dict()
        }

    except Exception as e:
        logger.error(f"Stellar classification failed: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/planet/radius")
async def planet_radius(request: PlanetRadiusRequest):
    """
    Calculate planet radius from transit depth.

    Returns planet radius, classification, and orbital parameters.
    """
    try:
        from skills.planet import PlanetRadiusCalculator

        calc = PlanetRadiusCalculator()
        result = calc.from_depth_ppm(
            request.depth_ppm,
            request.stellar_radius,
            period=request.period,
            stellar_mass=request.stellar_mass,
            stellar_teff=request.stellar_teff,
            stellar_luminosity=request.stellar_luminosity
        )

        return {
            "status": "success",
            "planet": result.to_dict()
        }

    except Exception as e:
        logger.error(f"Planet radius calculation failed: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/planet/hz")
async def habitable_zone(request: HabitableZoneRequest):
    """
    Calculate habitable zone boundaries.

    Returns inner and outer HZ edges in AU.
    """
    try:
        from skills.planet import PlanetRadiusCalculator

        calc = PlanetRadiusCalculator()
        hz = calc.calculate_habitable_zone(
            request.stellar_teff,
            request.stellar_luminosity
        )

        return {
            "status": "success",
            "habitable_zone": hz.to_dict()
        }

    except Exception as e:
        logger.error(f"Habitable zone calculation failed: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/pipeline")
async def run_pipeline(request: PipelineRequest, background_tasks: BackgroundTasks):
    """
    Run full analysis pipeline on a target.

    Steps:
    1. Fetch light curve (TESS/Kepler)
    2. Query stellar parameters (Gaia)
    3. Run BLS periodogram
    4. Phase fold
    5. Calculate planet radius
    6. Assess habitability
    7. Generate report

    Note: This endpoint fetches data from external archives and may take time.
    """
    try:
        import lightkurve as lk
        from skills.periodogram import BLSPeriodogram, phase_fold, bin_phase_curve
        from skills.planet import PlanetRadiusCalculator

        target = request.target

        # Step 1: Fetch light curve
        search = lk.search_lightcurve(target, mission=['TESS', 'Kepler'])
        if len(search) == 0:
            raise HTTPException(status_code=404, detail=f"No light curves found for {target}")

        lc = search[0].download()
        lc = lc.remove_nans().normalize().remove_outliers(sigma=3)
        time = lc.time.value
        flux = lc.flux.value

        # Step 2: Get stellar parameters (simplified - use defaults)
        stellar_params = {'teff': 5778, 'radius': 1.0, 'mass': 1.0, 'luminosity': 1.0}

        try:
            from skills.gaia import GaiaClient
            client = GaiaClient()
            source = client.query_by_name(target)
            if source and source.teff_gspphot:
                stellar_params = client.get_stellar_params(source)
        except:
            pass

        # Step 3: BLS analysis
        bls = BLSPeriodogram(
            min_period=0.5,
            max_period=20,
            n_periods=5000 if request.quick_mode else 10000
        )
        bls_result = bls.compute(time, flux, min_snr=6.0)

        # Step 4: Phase fold
        period = bls_result.candidates[0].period if bls_result.candidates else bls_result.best_period
        t0 = bls_result.candidates[0].t0 if bls_result.candidates else 0
        phase, flux_folded = phase_fold(time, flux, period, t0)
        bin_phase, bin_flux, bin_err = bin_phase_curve(phase, flux_folded)

        # Calculate depth
        depth = 1.0 - np.nanmin(bin_flux)

        # Step 5: Planet radius
        calc = PlanetRadiusCalculator()
        planet = calc.from_transit_depth(
            depth=depth,
            stellar_radius=stellar_params.get('radius', 1.0),
            period=period,
            stellar_mass=stellar_params.get('mass', 1.0),
            stellar_teff=stellar_params.get('teff', 5778),
            stellar_luminosity=stellar_params.get('luminosity', 1.0)
        )

        # Step 6: Habitable zone
        hz = calc.calculate_habitable_zone(
            stellar_params.get('teff', 5778),
            stellar_params.get('luminosity', 1.0)
        )

        return {
            "status": "success",
            "target": target,
            "mission": search[0].mission[0],
            "n_points": len(time),
            "stellar_params": stellar_params,
            "bls_result": {
                "best_period": bls_result.best_period,
                "n_candidates": len(bls_result.candidates),
                "candidates": [c.to_dict() for c in bls_result.candidates]
            },
            "transit": {
                "period": period,
                "depth_ppm": depth * 1e6
            },
            "planet": planet.to_dict(),
            "habitable_zone": hz.to_dict()
        }

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Pipeline failed: {e}")
        raise HTTPException(status_code=500, detail=str(e))


# ============================================================================
# Run Server
# ============================================================================

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

# ============================================================================
# TinyML Endpoints (for Cloud Platform)
# ============================================================================

from fastapi import File, UploadFile, Form
from astropy.io import fits
import tempfile
import os

@app.post("/api/tinyml/analyze")
async def analyze_tinyml(
    file: UploadFile = File(...),
    model_id: str = Form(...),
    user_id: str = Form(...)
):
    """
    Analyze FITS file with TinyML model.
    
    Returns classification results.
    """
    try:
        # Save uploaded file temporarily
        with tempfile.NamedTemporaryFile(delete=False, suffix='.fits') as tmp_file:
            content = await file.read()
            tmp_file.write(content)
            tmp_path = tmp_file.name
        
        try:
            # Read FITS file
            with fits.open(tmp_path) as hdul:
                # Extract light curve data
                # This is a simplified version - adjust based on your FITS format
                if len(hdul) > 1:
                    data = hdul[1].data
                    if 'TIME' in data.columns.names and 'FLUX' in data.columns.names:
                        time = data['TIME']
                        flux = data['FLUX']
                    else:
                        # Try other common column names
                        time = data[data.columns.names[0]]
                        flux = data[data.columns.names[1]]
                else:
                    raise HTTPException(400, "FITS file must have at least 2 HDUs")
            
            # Clean data (remove NaNs)
            mask = ~(np.isnan(time) | np.isnan(flux))
            time = time[mask]
            flux = flux[mask]
            
            # Normalize flux
            flux = (flux - np.median(flux)) / np.std(flux)
            
            # Simple mock analysis for now (replace with actual TinyML model)
            # TODO: Load and run actual TinyML models
            import random
            classification = random.choice([
                'planetary_transit', 'noise', 'stellar_signal', 
                'eclipsing_binary', 'variable_star'
            ])
            confidence = random.uniform(0.6, 0.95)
            
            result = {
                "classification": classification,
                "confidence": confidence,
                "probabilities": {
                    "planetary_transit": confidence if classification == "planetary_transit" else random.uniform(0.1, 0.3),
                    "noise": confidence if classification == "noise" else random.uniform(0.1, 0.3),
                    "stellar_signal": confidence if classification == "stellar_signal" else random.uniform(0.1, 0.3),
                    "eclipsing_binary": confidence if classification == "eclipsing_binary" else random.uniform(0.1, 0.3),
                },
                "inference_time_ms": random.uniform(5, 15),
                "model_id": model_id,
                "data_points": len(time)
            }
            
            return result
            
        finally:
            # Clean up temp file
            os.unlink(tmp_path)
            
    except Exception as e:
        logger.error(f"TinyML analysis failed: {e}")
        raise HTTPException(500, f"Analysis failed: {str(e)}")


@app.get("/api/tinyml/models")
async def get_tinyml_models():
    """List available TinyML models."""
    return {
        "models": [
            {
                "id": "EXOPLANET-001",
                "name": "Exoplanet Transit Detector",
                "accuracy": 0.98,
                "size_kb": 43,
                "status": "active"
            },
            {
                "id": "VSTAR-001",
                "name": "Variable Star Classifier",
                "accuracy": 0.998,
                "size_kb": 27,
                "status": "active"
            },
            {
                "id": "FLARE-001",
                "name": "Stellar Flare Detector",
                "accuracy": 0.967,
                "size_kb": 5,
                "status": "active"
            },
            {
                "id": "MICROLENS-001",
                "name": "Microlensing Detector",
                "accuracy": 0.994,
                "size_kb": 5,
                "status": "active"
            },
            {
                "id": "SUPERNOVA-001",
                "name": "Supernova Detector",
                "accuracy": 1.0,
                "size_kb": 3,
                "status": "active"
            },
            {
                "id": "SPECTYPE-001",
                "name": "Spectral Type Classifier",
                "accuracy": 0.95,
                "size_kb": 5,
                "status": "active"
            },
            {
                "id": "ASTERO-001",
                "name": "Asteroseismology Analyzer",
                "accuracy": 0.998,
                "size_kb": 21,
                "status": "active"
            },
            {
                "id": "GALAXY-001",
                "name": "Galaxy Morphology Classifier",
                "accuracy": 0.999,
                "size_kb": 4,
                "status": "active"
            }
        ]
    }


# ============================================================================
# Mount LARUN v2 API routes (/api/v2/*)
# ============================================================================
try:
    from larun.api.routes import classify as v2_classify
    from larun.api.routes import pipeline as v2_pipeline
    from larun.api.routes import discover as v2_discover
    from larun.api.routes import catalog as v2_catalog
    from larun.api.routes import leaderboard as v2_leaderboard

    app.include_router(v2_classify.router,    prefix="/api/v2", tags=["v2-Classification"])
    app.include_router(v2_pipeline.router,    prefix="/api/v2", tags=["v2-Pipelines"])
    app.include_router(v2_discover.router,    prefix="/api/v2", tags=["v2-Discovery"])
    app.include_router(v2_catalog.router,     prefix="/api/v2", tags=["v2-Catalog"])
    app.include_router(v2_leaderboard.router, prefix="/api/v2", tags=["v2-Leaderboard"])

    @app.get("/api/v2/health")
    async def v2_health():
        return {"status": "ok", "version": "2.0.0", "platform": "larun.space"}

    @app.get("/api/v2/models")
    async def v2_list_models():
        import json
        from pathlib import Path
        registry_path = Path(__file__).parent / "models" / "MODEL_REGISTRY.json"
        if registry_path.exists():
            with open(registry_path) as f:
                return json.load(f)
        return {"models": []}

    logger.info("LARUN v2 API routes mounted at /api/v2/*")
except Exception as _v2_err:
    logger.warning(f"v2 routes not loaded (non-fatal): {_v2_err}")