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", "")

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
# ============================

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 or None,
)
meta_logit = joblib.load(meta_path)

EXPECTED_FEATURES = getattr(meta_logit, "n_features_in_", 15)
if EXPECTED_FEATURES != 15:
    raise RuntimeError(f"Expected 15 features, got {EXPECTED_FEATURES}")

print("✅ Meta-logit loaded")

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

def get_embedding(text: str) -> np.ndarray:
    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))

    # Core features
    phi = float(np.clip(total / (total + 1.0), 0.0, 1.0))
    omega = float(np.clip(dominant_idx / len(power), 0.0, 1.0))

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

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

    hamiltonian_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": hamiltonian_energy,
        "dominant_frequency": float(dominant_idx),
    }


def closest_phi_node(feats: Dict[str, float]) -> int:
    if feats["coherence"] > 0.85 and feats["phi"] > 0.6:
        return 4
    if feats["hamiltonian_energy"] > 0.5:
        return 2
    if feats["omega"] < 0.2:
        return 0
    if feats["coherence"] < 0.4:
        return 5
    if feats["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)
    phi_one_hot = [1.0 if i == phi_idx else 0.0 for i in range(8)]

    # 🔥 EXACTAMENTE 15 FEATURES
    vector = [
        feats["phi"],
        feats["omega"],
        feats["coherence"],
        feats["S_RRF"],
        feats["C_RRF"],
        feats["hamiltonian_energy"],
        feats["dominant_frequency"],
        cos_pa,
        *phi_one_hot  # 8
    ]

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


def log_interaction(prompt, answer, score):
    try:
        with open("rrf_logs.jsonl", "a") as f:
            f.write(json.dumps({
                "t": time.time(),
                "prompt": prompt,
                "answer": answer,
                "score": score
            }) + "\n")
    except:
        pass

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

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

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

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


class EvaluateResponse(BaseModel):
    p_good: float
    scores: Dict[str, float]
    features: Dict[str, float]
    phi_node: str


class ImproveRequest(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", response_model=EvaluateResponse)
def evaluate(req: EvaluateRequest):
    if len(req.prompt) > MAX_PROMPT_CHARS or len(req.answer) > MAX_ANSWER_CHARS:
        raise HTTPException(413, "Payload too large")

    x, feats, phi_node = rrf_vector(req.prompt, req.answer)

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

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

    log_interaction(req.prompt, req.answer, p_good)

    return EvaluateResponse(
        p_good=p_good,
        scores=scores,
        features=feats,
        phi_node=phi_node
    )


@app.post("/quality", response_model=EvaluateResponse)
def quality(req: EvaluateRequest):
    return evaluate(req)


@app.post("/improve")
def improve(req: ImproveRequest):
    x, _, _ = rrf_vector(req.prompt, req.answer)
    p = meta_logit.predict_proba(x.reshape(1, -1))[0][1]

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

    return {
        "original_score": float(p),
        "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,
        convert_to_numpy=True,
        normalize_embeddings=True
    )

    q_emb = embs[0]
    d_embs = embs[1:]

    scores = d_embs @ q_emb
    ranked_idx = np.argsort(-scores)

    results = [
        {
            "id": int(i),
            "score": float(scores[i]),
            "rank": r + 1
        }
        for r, i in enumerate(ranked_idx)
    ]

    return {
        "model_id": ENCODER_MODEL_ID,
        "results": results
    }
]

# ============================
# STARTUP: MODELS
# ============================

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

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

EXPECTED_FEATURES = getattr(meta_logit, "n_features_in_", 15)
if EXPECTED_FEATURES != 15:
    raise RuntimeError(f"Meta-logit expects {EXPECTED_FEATURES} features, expected 15.")

print("✅ Meta-logit loaded (15D)", flush=True)

# ============================
# META-STATE FEATURE EXTRACTION
# ============================

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


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.0, 1.0))
    omega = float(np.clip(dominant_idx / len(power), 0.0, 1.0))

    S_RRF = float(1.0 - np.std(power) / (np.mean(power) + 1e-12))
    S_RRF = float(np.clip(S_RRF, 0.0, 1.0))

    coherence = float(0.5 * (1.0 - np.std(power) / (np.mean(power) + 1e-12)) + 0.5 * C_RRF)

    hamiltonian_energy = float(np.dot(emb, emb))
    dominant_frequency = float(dominant_idx)

    return {
        "phi": phi,
        "omega": omega,
        "coherence": coherence,
        "S_RRF": S_RRF,
        "C_RRF": C_RRF,
        "hamiltonian_energy": hamiltonian_energy,
        "dominant_frequency": dominant_frequency,
    }


def closest_phi_node(feats: Dict[str, float]) -> int:
    # Deterministic ontology mapping
    if feats["coherence"] > 0.85 and feats["phi"] > 0.6:
        return 4  # Φ4_resonance
    if feats["hamiltonian_energy"] > 50:
        return 2  # Φ2_gauge_dirac
    if feats["omega"] < 0.2:
        return 0  # Φ0_seed
    if feats["coherence"] < 0.4:
        return 5  # Φ5_memory_symbiosis
    if feats["phi"] < 0.3:
        return 6  # Φ6_alignment
    return 7  # Φ7_meta_agi


def rrf_state_to_vector(prompt: str, answer: str) -> np.ndarray:
    emb = get_embedding(prompt + "\n" + answer)
    feats = spectral_features(emb)

    phi_idx = closest_phi_node(feats)
    phi_one_hot = [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"],
        *phi_one_hot,
    ]

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

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

app = FastAPI(
    title="Savant RRF Φ12.0 API",
    version="2.0.0",
    description="Meta-state RRF quality evaluation + rerank",
)

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

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


class EvaluateResponse(BaseModel):
    p_good: float
    scores: Dict[str, float]
    features: Dict[str, float]
    phi_node: str


class RerankRequest(BaseModel):
    query: str
    documents: List[str]
    alpha: float = 0.2


class RerankDocument(BaseModel):
    id: int
    score: float
    rank: int


class RerankResponse(BaseModel):
    model_id: str
    results: List[RerankDocument]

# ============================
# MANIFEST / HEALTH
# ============================
@app.get("/")
def root():
    return {
        "status": "ok",
        "model": "RRFSavantMetaLogicV2",
        "version": "Φ12.0",
        "docs": "/docs",
        "endpoints": ["/manifest", "/health", "/evaluate", "/quality", "/v1/rerank"],
    }

@app.get("/manifest")
def manifest():
    return {
        "model": "RRFSavantMetaLogicV2",
        "version": "Φ12.0",
        "encoder": ENCODER_MODEL_ID,
        "features": 15,
        "phi_nodes": PHI_NODES,
    }

@app.get("/health")
def health():
    return {"status": "ok"}

# ============================
# /EVALUATE
# ============================

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

    x, feats, phi_node = rrf_state_to_vector(req.prompt, req.answer)
    proba = meta_logit.predict_proba(x.reshape(1, -1))[0]
    p_good = float(proba[1])

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

    return EvaluateResponse(
        p_good=p_good,
        scores=scores,
        features=feats,
        phi_node=phi_node,
    )


@app.post("/quality", response_model=EvaluateResponse)
def quality_alias(req: EvaluateRequest):
    return evaluate(req)

# ============================
# /v1/rerank (BATCHED)
# ============================

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

    texts = [req.query] + req.documents
    for d in req.documents:
        if len(d) > MAX_DOC_CHARS:
            raise HTTPException(413, "Document too large")

    embs = encoder.encode(
        texts,
        convert_to_numpy=True,
        normalize_embeddings=True,
    )

    q_emb = embs[0]
    d_embs = embs[1:]

    scores = d_embs @ q_emb
    ranked_idx = np.argsort(-scores)

    results = [
        RerankDocument(
            id=int(i),
            score=float(scores[i]),
            rank=r + 1,
        )
        for r, i in enumerate(ranked_idx)
    ]

    return RerankResponse(
        model_id=ENCODER_MODEL_ID,
        results=results,
    )