"""
PolyBloom ML Service — Space A v20.0
PRICE MEMORY TIER: Chronos (zero-shot T5 foundation model)

Tier role: reads closes + price-derived indicators only.
No orderflow (CVD, OB, funding). No sentiment (polymarket, fear/greed).

Endpoints:
  GET  /          health check
  POST /predict   returns single Chronos prediction
"""

import os
import math
import time
import logging
import numpy as np

from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse

logging.basicConfig(level=logging.INFO, format="%(levelname)s  %(name)s  %(message)s")
logger = logging.getLogger("space-a")

app = FastAPI(title="PolyBloom Space A — Price Memory", version="20.0.0")
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

# ─────────────────────────────────────────────────────────────────────────────
# Global state
# ─────────────────────────────────────────────────────────────────────────────

_model = None
_chronos_loaded = False

# ─────────────────────────────────────────────────────────────────────────────
# Helpers
# ─────────────────────────────────────────────────────────────────────────────

def _tanh(x: float) -> float:
    return math.tanh(x)

def _clamp(v: float, lo: float, hi: float) -> float:
    return max(lo, min(hi, v))

def _slope(arr: np.ndarray, p: int) -> float:
    if len(arr) < p + 1:
        return 0.0
    b = arr[-p - 1]
    return float((arr[-1] - b) / b) if b != 0 else 0.0

def _ema(arr: np.ndarray, p: int) -> float:
    k = 2.0 / (p + 1)
    e = float(arr[0])
    for v in arr[1:]:
        e = float(v) * k + e * (1.0 - k)
    return e

def _score_to_prediction(score: float, reasoning: str) -> dict:
    direction  = "UP" if score > 0.05 else ("DOWN" if score < -0.05 else "NEUTRAL")
    confidence = _clamp(45.0 + abs(score) * 55.0, 40.0, 95.0)
    return {
        "name":       "Chronos",
        "direction":  direction,
        "confidence": round(confidence),
        "slope_pct":  round(score * 0.5, 6),
        "reasoning":  reasoning,
        "tier":       "price_memory",
    }

# ─────────────────────────────────────────────────────────────────────────────
# Startup
# ─────────────────────────────────────────────────────────────────────────────

@app.on_event("startup")
async def startup_event():
    global _model, _chronos_loaded
    try:
        from chronos import ChronosPipeline
        import torch
        model_name = os.environ.get("CHRONOS_MODEL", "amazon/chronos-t5-tiny")
        logger.info("Chronos: loading %s …", model_name)
        _model = ChronosPipeline.from_pretrained(
            model_name,
            device_map="cpu",
            torch_dtype=torch.float32,
        )
        _chronos_loaded = True
        logger.info("Chronos loaded ✓")
    except Exception as exc:
        logger.warning("Chronos not loaded: %s", exc)

# ─────────────────────────────────────────────────────────────────────────────
# Inference
# ─────────────────────────────────────────────────────────────────────────────

def _predict_real(closes: list, feats: dict) -> dict:
    import torch
    c       = np.array(closes[-64:], dtype=np.float32)
    context = torch.tensor(c).unsqueeze(0)
    with torch.no_grad():
        forecast = _model.predict(context, prediction_length=2)
    median = float(forecast[0].median(dim=0).values.mean().item())
    last   = float(c[-1])
    score  = _tanh((median / last - 1.0) * 20.0) if last != 0 else 0.0
    return _score_to_prediction(score, f"Chronos: median={median:.2f} last={last:.2f}")


def _predict_bootstrap(closes: list, feats: dict) -> dict:
    c    = np.array(closes, dtype=np.float64)
    n    = len(c)
    last = float(c[-1]) if n > 0 else 1.0
    e20  = _ema(c[-max(40, n):], 20)
    e50  = _ema(c[-max(80, n):], min(50, max(n // 2, 5)))
    sl20 = _slope(c, min(20, n - 1))
    rsi  = float(feats.get("rsi") or 50.0)
    score = (
        math.copysign(min(1.0, abs(sl20) * 150.0), sl20) * 0.45 +
        math.copysign(min(1.0, abs((e20 - e50) / last) * 200.0), (e20 - e50)) * 0.35 +
        ((rsi - 50.0) / 50.0) * 0.20
    )
    return _score_to_prediction(
        score,
        f"Chronos bootstrap: sl20={sl20*100:.4f}% ema_diff={((e20-e50)/last)*100:.4f}%",
    )

# ─────────────────────────────────────────────────────────────────────────────
# Endpoints
# ─────────────────────────────────────────────────────────────────────────────

@app.get("/")
async def health():
    return {
        "ok":      True,
        "version": "20.0.0",
        "space":   "A",
        "tier":    "price_memory",
        "models":  ["Chronos"],
        "chronos": _chronos_loaded,
    }


@app.post("/predict")
async def predict(request: Request):
    t0   = time.perf_counter()
    body = await request.json()

    closes = list(body.get("closes") or [])
    if len(closes) < 16:
        return JSONResponse(status_code=422, content={"error": "closes too short — need 16+"})

    # Price-memory tier: only price-derived features allowed
    allowed = {"rsi", "stoch_rsi_k", "stoch_rsi_d", "vwap_dev", "atr_pct_of_spot",
               "atr", "bollinger_squeeze", "vol_regime", "vol_percentile",
               "garman_klass_vol", "market_regime", "htf_trend_15m", "htf_trend_1h",
               "candle_pattern", "nearest_support", "nearest_resistance", "session"}
    feats = {k: v for k, v in body.items() if k != "closes" and k in allowed}

    if _chronos_loaded and _model is not None:
        try:
            pred = _predict_real(closes, feats)
        except Exception as exc:
            logger.warning("Chronos real inference failed: %s", exc)
            pred = _predict_bootstrap(closes, feats)
    else:
        pred = _predict_bootstrap(closes, feats)

    pred["latency_ms"] = round((time.perf_counter() - t0) * 1000)
    return pred