app/main.py

import os
from typing import List, Dict, Optional
import json

import numpy as np
import pandas as pd
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from huggingface_hub import InferenceClient


# =========================
# Configuración
# =========================

HF_TOKEN = os.getenv("HF_TOKEN")
MODEL_ID = os.getenv("CHRONOS_MODEL_ID", "amazon/chronos-t5-large")

app = FastAPI(
    title="Chronos-2 Forecasting API (HF Inference)",
    description=(
        "API de pronósticos usando Chronos-2 via Hugging Face Inference API. "
        "Compatible con Excel Add-in."
    ),
    version="1.0.0",
)

# Configurar CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # En producción, especificar dominios permitidos
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Cliente de HF Inference
if not HF_TOKEN:
    print("⚠️  WARNING: HF_TOKEN no configurado. La API puede no funcionar correctamente.")
    print("   Configura HF_TOKEN en las variables de entorno del Space.")
    client = None
else:
    client = InferenceClient(token=HF_TOKEN)


# =========================
# Modelos Pydantic
# =========================

class UnivariateSeries(BaseModel):
    values: List[float]


class ForecastUnivariateRequest(BaseModel):
    series: UnivariateSeries
    prediction_length: int = Field(7, description="Número de pasos a predecir")
    quantile_levels: Optional[List[float]] = Field(
        default=[0.1, 0.5, 0.9],
        description="Cuantiles para intervalos de confianza"
    )
    freq: str = Field("D", description="Frecuencia temporal (D, W, M, etc.)")


class ForecastUnivariateResponse(BaseModel):
    timestamps: List[str]
    median: List[float]
    quantiles: Dict[str, List[float]]


class AnomalyDetectionRequest(BaseModel):
    context: UnivariateSeries
    recent_observed: List[float]
    prediction_length: int = 7
    quantile_low: float = 0.05
    quantile_high: float = 0.95


class AnomalyPoint(BaseModel):
    index: int
    value: float
    predicted_median: float
    lower: float
    upper: float
    is_anomaly: bool


class AnomalyDetectionResponse(BaseModel):
    anomalies: List[AnomalyPoint]


class BacktestRequest(BaseModel):
    series: UnivariateSeries
    prediction_length: int = 7
    test_length: int = 28


class BacktestMetrics(BaseModel):
    mae: float
    mape: float
    rmse: float


class BacktestResponse(BaseModel):
    metrics: BacktestMetrics
    forecast_median: List[float]
    forecast_timestamps: List[str]
    actuals: List[float]


# =========================
# Función auxiliar para llamar a HF Inference
# =========================

def call_chronos_inference(series: List[float], prediction_length: int) -> Dict:
    """
    Llama a la API de Hugging Face Inference para Chronos.
    Retorna un diccionario con las predicciones.
    """
    if client is None:
        raise HTTPException(
            status_code=503,
            detail="HF_TOKEN no configurado. Contacta al administrador del servicio."
        )
    
    try:
        # Intentar usando el endpoint específico de time series
        import requests
        
        url = f"https://api-inference.huggingface.co/models/{MODEL_ID}"
        headers = {"Authorization": f"Bearer {HF_TOKEN}"}
        
        payload = {
            "inputs": series,
            "parameters": {
                "prediction_length": prediction_length,
                "num_samples": 100  # Para obtener cuantiles
            }
        }
        
        response = requests.post(url, headers=headers, json=payload, timeout=60)
        
        if response.status_code == 503:
            raise HTTPException(
                status_code=503,
                detail="El modelo está cargando. Por favor, intenta de nuevo en 30-60 segundos."
            )
        elif response.status_code != 200:
            raise HTTPException(
                status_code=response.status_code,
                detail=f"Error de la API de HuggingFace: {response.text}"
            )
        
        result = response.json()
        return result
        
    except requests.exceptions.Timeout:
        raise HTTPException(
            status_code=504,
            detail="Timeout al comunicarse con HuggingFace API. El modelo puede estar cargando."
        )
    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Error inesperado: {str(e)}"
        )


def process_chronos_output(raw_output: Dict, prediction_length: int) -> Dict:
    """
    Procesa la salida de Chronos para extraer mediana y cuantiles.
    """
    # La API de Chronos puede devolver diferentes formatos
    # Intentamos adaptarnos a ellos
    
    if isinstance(raw_output, list):
        # Si es una lista de valores, asumimos que es la predicción media
        median = raw_output[:prediction_length]
        return {
            "median": median,
            "quantiles": {
                "0.1": median,  # Sin cuantiles, usar median
                "0.5": median,
                "0.9": median
            }
        }
    
    # Si tiene estructura más compleja, intentar extraer
    if "forecast" in raw_output:
        forecast = raw_output["forecast"]
        if "median" in forecast:
            median = forecast["median"][:prediction_length]
        else:
            median = forecast.get("mean", [0] * prediction_length)[:prediction_length]
        
        quantiles = forecast.get("quantiles", {})
        return {
            "median": median,
            "quantiles": quantiles
        }
    
    # Formato por defecto
    return {
        "median": [0] * prediction_length,
        "quantiles": {
            "0.1": [0] * prediction_length,
            "0.5": [0] * prediction_length,
            "0.9": [0] * prediction_length
        }
    }


# =========================
# Endpoints
# =========================

@app.get("/")
def root():
    """Información básica de la API"""
    return {
        "name": "Chronos-2 Forecasting API",
        "version": "1.0.0",
        "model": MODEL_ID,
        "status": "running",
        "docs": "/docs",
        "health": "/health"
    }


@app.get("/health")
def health():
    """Health check del servicio"""
    return {
        "status": "ok" if HF_TOKEN else "warning",
        "model_id": MODEL_ID,
        "hf_token_configured": HF_TOKEN is not None,
        "message": "Ready" if HF_TOKEN else "HF_TOKEN not configured"
    }


@app.post("/forecast_univariate", response_model=ForecastUnivariateResponse)
def forecast_univariate(req: ForecastUnivariateRequest):
    """
    Pronóstico para una serie temporal univariada.
    
    Compatible con el Excel Add-in.
    """
    values = req.series.values
    n = len(values)
    
    if n == 0:
        raise HTTPException(status_code=400, detail="La serie no puede estar vacía.")
    
    if n < 3:
        raise HTTPException(
            status_code=400,
            detail="La serie debe tener al menos 3 puntos históricos."
        )
    
    # Llamar a la API de HuggingFace
    raw_output = call_chronos_inference(values, req.prediction_length)
    
    # Procesar la salida
    processed = process_chronos_output(raw_output, req.prediction_length)
    
    # Generar timestamps
    timestamps = [f"t+{i+1}" for i in range(req.prediction_length)]
    
    return ForecastUnivariateResponse(
        timestamps=timestamps,
        median=processed["median"],
        quantiles=processed["quantiles"]
    )


@app.post("/detect_anomalies", response_model=AnomalyDetectionResponse)
def detect_anomalies(req: AnomalyDetectionRequest):
    """
    Detecta anomalías comparando valores observados con predicciones.
    """
    n_hist = len(req.context.values)
    
    if n_hist == 0:
        raise HTTPException(status_code=400, detail="El contexto no puede estar vacío.")
    
    if len(req.recent_observed) != req.prediction_length:
        raise HTTPException(
            status_code=400,
            detail="recent_observed debe tener la misma longitud que prediction_length."
        )
    
    # Hacer predicción
    raw_output = call_chronos_inference(req.context.values, req.prediction_length)
    processed = process_chronos_output(raw_output, req.prediction_length)
    
    # Comparar con valores observados
    anomalies: List[AnomalyPoint] = []
    
    median = processed["median"]
    # Intentar obtener cuantiles o usar aproximaciones
    q_low = processed["quantiles"].get(str(req.quantile_low), median)
    q_high = processed["quantiles"].get(str(req.quantile_high), median)
    
    for i, obs in enumerate(req.recent_observed):
        if i < len(median):
            lower = q_low[i] if i < len(q_low) else median[i] * 0.8
            upper = q_high[i] if i < len(q_high) else median[i] * 1.2
            predicted = median[i]
            is_anom = (obs < lower) or (obs > upper)
            
            anomalies.append(
                AnomalyPoint(
                    index=i,
                    value=obs,
                    predicted_median=predicted,
                    lower=lower,
                    upper=upper,
                    is_anomaly=is_anom,
                )
            )
    
    return AnomalyDetectionResponse(anomalies=anomalies)


@app.post("/backtest_simple", response_model=BacktestResponse)
def backtest_simple(req: BacktestRequest):
    """
    Backtesting simple: divide la serie en train/test y evalúa métricas.
    """
    values = np.array(req.series.values, dtype=float)
    n = len(values)
    
    if n <= req.test_length:
        raise HTTPException(
            status_code=400,
            detail="La serie debe ser más larga que test_length."
        )
    
    # Dividir en train/test
    train = values[: n - req.test_length].tolist()
    test = values[n - req.test_length :].tolist()
    
    # Hacer predicción
    raw_output = call_chronos_inference(train, req.test_length)
    processed = process_chronos_output(raw_output, req.test_length)
    
    forecast = np.array(processed["median"], dtype=float)
    test_arr = np.array(test, dtype=float)
    
    # Calcular métricas
    mae = float(np.mean(np.abs(test_arr - forecast)))
    rmse = float(np.sqrt(np.mean((test_arr - forecast) ** 2)))
    
    eps = 1e-8
    mape = float(np.mean(np.abs((test_arr - forecast) / (test_arr + eps)))) * 100.0
    
    timestamps = [f"test_t{i+1}" for i in range(req.test_length)]
    
    metrics = BacktestMetrics(mae=mae, mape=mape, rmse=rmse)
    
    return BacktestResponse(
        metrics=metrics,
        forecast_median=forecast.tolist(),
        forecast_timestamps=timestamps,
        actuals=test,
    )


# =========================
# Endpoints simplificados para testing
# =========================

@app.post("/simple_forecast")
def simple_forecast(series: List[float], prediction_length: int = 7):
    """
    Endpoint simplificado para testing rápido.
    """
    if not series:
        raise HTTPException(status_code=400, detail="Serie vacía")
    
    raw_output = call_chronos_inference(series, prediction_length)
    processed = process_chronos_output(raw_output, prediction_length)
    
    return {
        "input_series": series,
        "prediction_length": prediction_length,
        "forecast": processed["median"],
        "model": MODEL_ID
    }


if __name__ == "__main__":
    import uvicorn
    port = int(os.getenv("PORT", 7860))
    uvicorn.run(app, host="0.0.0.0", port=port)