# ======================================================
# Savant RRF Φ12.5 — CLEAN & STABLE MAIN
# ======================================================

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

import numpy as np
from numpy.linalg import norm
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

from sentence_transformers import SentenceTransformer
from huggingface_hub import hf_hub_download
import joblib

# ============================
# CONFIG
# ============================

HF_TOKEN = os.environ.get("HF_TOKEN", None)

ENCODER_MODEL_ID = "antonypamo/RRFSAVANTMADE"
META_LOGIT_REPO = "antonypamo/RRFSavantMetaLogicV2"
META_LOGIT_FILENAME = "logreg_rrf_savant.joblib"

MAX_PROMPT_CHARS = 8000
MAX_ANSWER_CHARS = 12000
MAX_DOCS = 50
MAX_DOC_CHARS = 6000

PHI_NODES = [
    "Φ0_seed","Φ1_geometric","Φ2_gauge_dirac","Φ3_log_gravity",
    "Φ4_resonance","Φ5_memory_symbiosis","Φ6_alignment","Φ7_meta_agi"
]

# ============================
# LOAD MODELS (ONCE)
# ============================

print("🔄 Loading encoder...", flush=True)
encoder = SentenceTransformer(ENCODER_MODEL_ID)
print("✅ Encoder loaded")

print("🔄 Loading meta-logit...", flush=True)
meta_path = hf_hub_download(
    repo_id=META_LOGIT_REPO,
    filename=META_LOGIT_FILENAME,
    token=HF_TOKEN
)
meta_logit = joblib.load(meta_path)

if getattr(meta_logit, "n_features_in_", 15) != 15:
    raise RuntimeError("Meta-logit must have 15 features")

print("✅ Meta-logit ready")

# ============================
# CORE FUNCTIONS
# ============================

def get_embedding(text: str):
    return encoder.encode([text], convert_to_numpy=True, normalize_embeddings=True)[0]


def cosine_sim(a, b):
    return float(np.dot(a, b) / (norm(a) * norm(b) + 1e-12))


def spectral_features(emb: np.ndarray) -> Dict[str, float]:
    fft = np.fft.rfft(emb)
    power = np.abs(fft) ** 2

    total = power.sum() + 1e-12
    dominant_idx = int(np.argmax(power))

    phi = float(np.clip(total / (total + 1.0), 0, 1))
    omega = float(np.clip(dominant_idx / len(power), 0, 1))

    C_RRF = float(1 - np.std(power) / (np.mean(power) + 1e-12))
    C_RRF = float(np.clip(C_RRF, 0, 1))

    S_RRF = C_RRF
    coherence = float(0.5 * (S_RRF + C_RRF))

    energy = float(np.dot(emb, emb) / len(emb))

    return {
        "phi": phi,
        "omega": omega,
        "coherence": coherence,
        "S_RRF": S_RRF,
        "C_RRF": C_RRF,
        "hamiltonian_energy": energy,
        "dominant_frequency": float(dominant_idx),
    }


def closest_phi_node(f):
    if f["coherence"] > 0.85 and f["phi"] > 0.6:
        return 4
    if f["hamiltonian_energy"] > 0.5:
        return 2
    if f["omega"] < 0.2:
        return 0
    if f["coherence"] < 0.4:
        return 5
    if f["phi"] < 0.3:
        return 6
    return 7


def rrf_vector(prompt: str, answer: str):
    emb_p = get_embedding(prompt)
    emb_a = get_embedding(answer)
    emb = get_embedding(prompt + "\n" + answer)

    feats = spectral_features(emb)
    cos_pa = cosine_sim(emb_p, emb_a)

    phi_idx = closest_phi_node(feats)
    onehot = [1.0 if i == phi_idx else 0.0 for i in range(8)]

    vector = [
        feats["phi"],
        feats["omega"],
        feats["coherence"],
        feats["S_RRF"],
        feats["C_RRF"],
        feats["hamiltonian_energy"],
        feats["dominant_frequency"],
        cos_pa,
        *onehot
    ]

    return np.array(vector, dtype=float), feats, PHI_NODES[phi_idx], cos_pa


def compute_scores(prompt, answer):
    x, feats, phi_node, cosine = rrf_vector(prompt, answer)

    p_good = float(meta_logit.predict_proba(x.reshape(1, -1))[0][1])

    resonance_score = 0.5 * p_good + 0.5 * abs(cosine)

    scores = {
        "SRRF": p_good,
        "CRRF": p_good * feats["coherence"] * feats["phi"],
        "E_phi": 0.5 * (p_good + feats["phi"]),
        "resonance_score": resonance_score,
        "cosine": cosine
    }

    return p_good, scores, feats, phi_node


# ============================
# FASTAPI
# ============================

app = FastAPI(title="Savant RRF Φ12.5 API", version="2.2.0")

# ============================
# SCHEMAS
# ============================

class EvaluateRequest(BaseModel):
    prompt: str
    answer: str


class RerankRequest(BaseModel):
    query: str
    documents: List[str]


# ============================
# ROUTES
# ============================

@app.get("/")
def root():
    return {
        "status": "ok",
        "version": "Φ12.5",
        "endpoints": ["/evaluate","/quality","/improve","/v1/rerank"]
    }


@app.post("/evaluate")
def evaluate(req: EvaluateRequest):
    if len(req.prompt) > MAX_PROMPT_CHARS or len(req.answer) > MAX_ANSWER_CHARS:
        raise HTTPException(413, "Payload too large")

    p_good, scores, feats, phi_node = compute_scores(req.prompt, req.answer)

    return {
        "p_good": p_good,
        "scores": scores,
        "features": feats,
        "phi_node": phi_node
    }


@app.post("/quality")
def quality(req: EvaluateRequest):
    return evaluate(req)


@app.post("/improve")
def improve(req: EvaluateRequest):
    p_good, _, _, _ = compute_scores(req.prompt, req.answer)

    if p_good < 0.6:
        improved = req.answer + " (refined for clarity, coherence and reasoning depth)"
    else:
        improved = req.answer

    return {
        "original_score": p_good,
        "improved_answer": improved
    }


@app.post("/v1/rerank")
def rerank(req: RerankRequest):
    if len(req.documents) > MAX_DOCS:
        raise HTTPException(413, "Too many documents")

    texts = [req.query] + req.documents
    embs = encoder.encode(texts, normalize_embeddings=True)

    q = embs[0]
    docs = embs[1:]

    scores = docs @ q
    idx = np.argsort(-scores)

    return {
        "model_id": ENCODER_MODEL_ID,
        "results": [
            {"id": int(i), "score": float(scores[i]), "rank": r+1}
            for r, i in enumerate(idx)
        ]
    }