import os
import numpy as np
from numpy.linalg import norm
from scipy.linalg import expm
from sentence_transformers import SentenceTransformer
from huggingface_hub import hf_hub_download
import joblib

from fastapi import FastAPI
from pydantic import BaseModel, Field
from typing import Optional, Dict, Any


# NOTE: HF_TOKEN is expected to be set as an environment variable in a real deployment
# For local testing, you might set it here or pass it directly
HF_TOKEN = os.environ.get("HF_TOKEN", "") # Use environment variable, default to empty
os.environ["HF_TOKEN"] = HF_TOKEN

ENCODER_MODEL_ID = "antonypamo/RRFSAVANTMADE"        # encoder RRF
META_LOGIT_REPO = "antonypamo/RRFSavantMetaLogit"    # repo del meta-logit
META_LOGIT_FILENAME = "logreg_rrf_savant_v2.joblib"  # NUEVO archivo del meta-logit en HF

print("🔄 Cargando encoder RRFSAVANTMADE...")
encoder = SentenceTransformer(ENCODER_MODEL_ID)

print("🔄 Descargando meta-logit v2 desde HF Hub...")
meta_logit_path = hf_hub_download(
    repo_id=META_LOGIT_REPO,
    filename=META_LOGIT_FILENAME,
    token=os.environ.get("HF_TOKEN")
)

print("🔄 Cargando modelo meta-logit v2...")
meta_logit = joblib.load(meta_logit_path)

print("✅ Encoder y meta-logit v2 cargados correctamente.")


# =========================
# Geometría icosaédrica
# (Copied from cell lyVrwdhgIOlq)
# =========================

phi = (1 + np.sqrt(5)) / 2
nodes = np.array([
    [0, 1, phi], [0, -1, phi], [0, 1, -phi], [0, -1, -phi],
    [1, phi, 0], [-1, phi, 0], [1, -phi, 0], [-1, -phi, 0],
    [phi, 0, 1], [phi, 0, -1], [-phi, 0, 1], [-phi, 0, -1]
], dtype=float)
nodes /= norm(nodes, axis=1, keepdims=True)
N = nodes.shape[0]  # 12 nodos

# Pauli
sigma_x = np.array([[0, 1], [1, 0]], dtype=complex)
sigma_y = np.array([[0, -1j], [1j, 0]], dtype=complex)
sigma_z = np.array([[1, 0], [0, -1]], dtype=complex)

def kron_IN(M, N_sites):
    return np.kron(M, np.eye(N_sites, dtype=complex))

def site_op(block_2x2, i, j, N_sites):
    K = np.zeros((N_sites, N_sites), dtype=complex)
    K[i, j] = 1.0
    return np.kron(K, block_2x2)

def geodesic_kernel(nodes, sigma=0.618, alpha_log=0.10):
    diff = nodes[:, None, :] - nodes[None, :, :]
    dist = norm(diff, axis=-1)

    W = np.exp(-(dist**2) / (sigma**2))
    np.fill_diagonal(W, 0.0)

    if alpha_log > 0.0:
        corr = 1.0 + alpha_log * np.log1p(dist**2)
        corr[range(N), range(N)] = 1.0
        W = W / corr

    row_sums = W.sum(axis=1, keepdims=True)
    row_sums[row_sums == 0] = 1.0
    return W / row_sums

def u1_edge_phases(nodes, flux_vector=(0.0, 0.0, 0.0), q=1.0, gauge_scale=1.0):
    A = gauge_scale * np.asarray(flux_vector, dtype=float)
    midpoints = (nodes[:, None, :] + nodes[None, :, :]) / 2.0
    theta = (midpoints @ A).astype(float)
    theta = 0.5 * (theta - theta.T)
    return theta * q

def build_dirac_hamiltonian(
    m=0.25,
    v=1.0,
    sigma=0.618,
    alpha_log=0.10,
    q=1.0,
    flux_vector=(0.0, 0.0, 0.0),
    gauge_scale=0.0
):
    W = geodesic_kernel(nodes, sigma=sigma, alpha_log=alpha_log)

    if gauge_scale != 0.0 and any(flux_vector):
        theta = u1_edge_phases(nodes, flux_vector=flux_vector,
                               q=q, gauge_scale=gauge_scale)
        U = np.exp(1j * theta)
    else:
        U = np.ones((N, N), dtype=complex)

    # Término de masa
    H = np.kron(np.eye(N, dtype=complex), m * sigma_z)

    # Término cinético acoplado
    diff = nodes[:, None, :] - nodes[None, :, :]
    dist = norm(diff, axis=-1) + 1e-12
    d_hat = diff / dist[..., None]

    for i in range(N):
        for j in range(N):
            if i == j or W[i, j] == 0:
                continue
            nvec = d_hat[i, j]
            S = (nvec[0] * sigma_x +
                 nvec[1] * sigma_y +
                 nvec[2] * sigma_z)
            H += v * W[i, j] * U[i, j] * site_op(S, i, j, N)

    # Hermitizar por seguridad numérica
    H = 0.5 * (H + H.conj().T)
    return H

def site_probs(psi):
    N2 = psi.shape[0]
    n = N2 // 2
    psi_mat = psi.reshape(n, 2)
    return np.sum(np.abs(psi_mat)**2, axis=1).real

def chirality(psi):
    S = kron_IN(sigma_z, N)
    return float(np.vdot(psi, S @ psi).real)

def energy_expectation(psi, H):
    return float(np.vdot(psi, H @ psi).real)

def spatial_entropy(p):
    p = np.clip(p, 1e-12, 1.0)
    return float(-np.sum(p * np.log(p)).real)

def evolve_dirac_shell(psi0, H, dt=0.05, steps=200, record_every=20):
    U = expm(-1j * dt * H)
    psi = psi0.copy()

    probs_hist = []
    energy_hist = []
    chir_hist = []
    ent_hist = []

    for t in range(steps + 1):
        if t % record_every == 0:
            p = site_probs(psi)
            probs_hist.append(p)
            energy_hist.append(energy_expectation(psi, H))
            chir_hist.append(chirality(psi))
            ent_hist.append(spatial_entropy(p))

        psi = U @ psi
        psi /= np.sqrt(np.vdot(psi, psi))

    return {
        "probs": np.array(probs_hist, dtype=float),
        "energy": np.array(energy_hist, dtype=float),
        "chirality": np.array(chir_hist, dtype=float),
        "entropy": np.array(ent_hist, dtype=float),
        "dt": dt,
        "record_every": record_every,
    }


# =========================
# Feature extraction and scoring
# (Copied from cell DiknqWJZIZ5q)
# =========================

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

def compute_rrf_features(prompt: str, answer: str) -> dict:
    # Embeddings RRF
    e_p = get_embedding(prompt)
    e_a = get_embedding(answer)

    cosine_pa = float(np.dot(e_p, e_a))
    len_ratio = len(answer) / (len(prompt) + 1.0)

    # Estado inicial ligado al texto (seed reproducible)
    rng = np.random.default_rng(abs(hash(prompt + answer)) % (2**32))
    vec = rng.normal(0, 1, (2*N,)) + 1j * rng.normal(0, 1, (2*N,))
    vec /= np.sqrt(np.vdot(vec, vec))
    psi0 = vec

    # Hamiltoniano Dirac Φ12.0
    H = build_dirac_hamiltonian(
        m=0.25, v=1.0, sigma=0.618,
        alpha_log=0.10, q=1.0,
        flux_vector=(0.0, 0.0, 0.0),
        gauge_scale=0.0
    )

    out = evolve_dirac_shell(psi0, H, dt=0.05, steps=200, record_every=20)

    probs = out["probs"]
    energy = out["energy"]
    chir = out["chirality"]
    entropy = out["entropy"]

    S_initial = float(entropy[0])
    S_final = float(entropy[-1])
    S_delta = S_final - S_initial
    C_final = float(chir[-1])
    E_mean = float(np.mean(energy))
    E_std = float(np.std(energy))

    return {
        "cosine_pa": cosine_pa,
        "len_ratio": len_ratio,
        "dirac_entropy_final": S_final,
        "dirac_entropy_delta": S_delta,
        "dirac_chirality_final": C_final,
        "dirac_energy_mean": E_mean,
        "dirac_energy_std": E_std,
    }

def features_to_vector(feats: dict) -> np.ndarray:
    keys = [
        "cosine_pa",
        "len_ratio",
        "dirac_entropy_final",
        "dirac_entropy_delta",
        "dirac_chirality_final",
        "dirac_energy_mean",
        "dirac_energy_std",
    ]
    return np.array([feats[k] for k in keys], dtype=float)

def compute_scores_srff_crrf_ephi(prompt: str, answer: str):
    feats = compute_rrf_features(prompt, answer)
    x = features_to_vector(feats).reshape(1, -1)

    # meta-logit v2: pipeline (scaler + logistic regression)
    proba = meta_logit.predict_proba(x)[0]
    p_good = float(proba[1])

    SRRF = p_good
    CRRF = p_good * feats["cosine_pa"]

    S_final = feats["dirac_entropy_final"]
    S_max = np.log(N)
    norm_entropy = float(S_final / S_max)

    E_phi = 0.5 * (SRRF + norm_entropy)

    scores = {
        "SRRF": SRRF,
        "CRRF": CRRF,
        "E_phi": E_phi,
        "p_good": p_good,
    }
    return scores, feats


# =========================
# FastAPI App
# (Copied from cell LwlyX4-LIgKK)
# =========================

app = FastAPI(
    title="Savant RRF Φ12.0 API",
    description="Evaluación conceptual resonante para texto generado por LLMs (SRRF / CRRF / E_phi).",
    version="1.0.0",
)

class EvaluateRequest(BaseModel):
    prompt: str = Field(..., description="Pregunta / instrucción original.")
    answer: str = Field(..., description="Respuesta generada por un LLM.")
    model_label: Optional[str] = Field(
        None, description="Etiqueta opcional del modelo que generó la respuesta."
    )

class EvaluateResponse(BaseModel):
    scores: Dict[str, float]
    features: Dict[str, float]
    sim_summary: Dict[str, Any]

@app.post("/evaluate", response_model=EvaluateResponse)
def evaluate_endpoint(req: EvaluateRequest):
    scores, feats = compute_scores_srff_crrf_ephi(req.prompt, req.answer)

    # mini-sim extra para resumen diagnóstico simple
    H = build_dirac_hamiltonian(
        m=0.25, v=1.0, sigma=0.618,
        alpha_log=0.10, q=1.0,
        flux_vector=(0.0, 0.0, 0.0),
        gauge_scale=0.0
    )
    rng = np.random.default_rng(abs(hash(req.prompt + req.answer)) % (2**32))
    vec = rng.normal(0, 1, (2*N,)) + 1j * rng.normal(0, 1, (2*N,))
    vec /= np.sqrt(np.vdot(vec, vec))
    psi0 = vec

    sim = evolve_dirac_shell(psi0, H, dt=0.05, steps=100, record_every=25)

    sim_summary = {
        "entropy_initial": float(sim["entropy"][0]),
        "entropy_final": float(sim["entropy"][-1]),
        "chirality_initial": float(sim["chirality"][0]),
        "chirality_final": float(sim["chirality"][-1]),
        "energy_mean": float(np.mean(sim["energy"])),
        "energy_std": float(np.std(sim["energy"])),
        "N_sites": int(N),
    }

    return EvaluateResponse(
        scores=scores,
        features=feats,
        sim_summary=sim_summary,
    )