# ======================================================
# Savant RRF Φ12.5 — FULL APP (PRODUCTION ALIGNED)
# ======================================================

from __future__ import annotations

import os, json, math, time
from typing import Optional, Dict, Any, List

import numpy as np
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

from sentence_transformers import SentenceTransformer
from huggingface_hub import hf_hub_download
import joblib


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

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

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

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


# ======================================================
# LOAD MODELS
# ======================================================

embedder = SentenceTransformer(ENCODER_MODEL_ID)

meta_logit_path = hf_hub_download(
    repo_id=META_LOGIT_REPO,
    filename=META_LOGIT_FILENAME,
    token=HF_TOKEN
)

meta_logit = joblib.load(meta_logit_path)


# ======================================================
# CORE FEATURES
# ======================================================

def coherence(vec):
    if len(vec) < 4:
        return 0.0, 0.0

    spectrum = np.fft.rfft(vec)
    power = np.abs(spectrum) ** 2

    total = power.sum() + 1e-9
    C = power.max() / total

    freqs = np.fft.rfftfreq(len(vec))
    f_mean = (freqs * power).sum() / total
    S = 1 - min(1, f_mean / (freqs.max() + 1e-9))

    return S, C


def extract_features(text):
    vec = embedder.encode([text])[0]

    energy = float(np.dot(vec, vec))
    dom_freq = float(np.argmax(np.abs(np.fft.rfft(vec))))

    phi = 1 - math.exp(-energy)
    omega = math.tanh(dom_freq)

    S, C = coherence(vec)
    coh = 0.5 * (S + C)

    # dummy phi node (simplified)
    phi_node = int(abs(hash(text)) % 8)
    onehot = np.zeros(8)
    onehot[phi_node] = 1

    base = np.array([phi, omega, coh, S, C, energy, dom_freq])
    return np.concatenate([base, onehot]), phi_node


def compute_scores(prompt, answer):
    if not prompt or not answer:
        raise ValueError("Empty input")

    if len(prompt) > MAX_PROMPT_CHARS or len(answer) > MAX_ANSWER_CHARS:
        raise HTTPException(status_code=413, detail="Too large")

    qa = f"{prompt}\n{answer}"

    features, phi_node = extract_features(qa)
    p_good = float(meta_logit.predict_proba([features])[0][1])

    e_p = embedder.encode([prompt], normalize_embeddings=True)[0]
    e_a = embedder.encode([answer], normalize_embeddings=True)[0]
    cosine = float(np.dot(e_p, e_a))

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

    return {
        "p_good": p_good,
        "cosine": cosine,
        "resonance_score": resonance_score,
        "phi_node": f"Φ{phi_node}",
    }


# ======================================================
# API MODELS
# ======================================================

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


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


# ======================================================
# APP
# ======================================================

app = FastAPI(title="Savant RRF Φ12.5")


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


# ======================================================
# QUALITY
# ======================================================

@app.post("/quality")
def quality(req: EvalRequest):
    return compute_scores(req.prompt, req.answer)


@app.post("/evaluate")
def evaluate(req: EvalRequest):
    return compute_scores(req.prompt, req.answer)


# ======================================================
# IMPROVEMENT LOOP
# ======================================================

@app.post("/evaluate_and_improve")
def evaluate_and_improve(req: EvalRequest):
    base = compute_scores(req.prompt, req.answer)

    if base["p_good"] > 0.75:
        return {
            "final_answer": req.answer,
            "improved": False,
            "scores": base
        }

    improved = req.answer + "\n\n(Refined for clarity and reasoning)"

    new_score = compute_scores(req.prompt, improved)

    if new_score["p_good"] > base["p_good"]:
        return {
            "final_answer": improved,
            "improved": True,
            "scores": new_score
        }

    return {
        "final_answer": req.answer,
        "improved": False,
        "scores": base
    }


# ======================================================
# SELF LEARNING (BASIC)
# ======================================================

@app.post("/self_learn")
def self_learn(req: EvalRequest):
    result = compute_scores(req.prompt, req.answer)

    if result["resonance_score"] > 0.8:
        with open("rrf_memory.jsonl", "a") as f:
            f.write(json.dumps({
                "prompt": req.prompt,
                "answer": req.answer,
                "score": result["resonance_score"]
            }) + "\n")

        return {"stored": True, "score": result["resonance_score"]}

    return {"stored": False, "score": result["resonance_score"]}


# ======================================================
# RERANK
# ======================================================

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

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

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

    scores = docs @ q

    ranked = sorted(
        [{"id": i, "score": float(s)} for i, s in enumerate(scores)],
        key=lambda x: x["score"],
        reverse=True
    )

    return {"results": ranked}