# =============================================================================
# GRU v1.8.7 — Protocolo Octant-Blind para CDT Formal
# Post-procesamiento aplicable a triangulaciones CDT reales
# Autor: Alfredo Flores Cornejo | dr.alfredo.fc@gmail.com
# DOI: 10.5281/zenodo.20472915
#
# PROPÓSITO:
#   Aplicar el protocolo GRU como capa de post-procesamiento sobre
#   un grafo CDT existente. No requiere modificar:
#   - La acción de Regge
#   - El sampler Monte Carlo
#   - La foliación causal
#   Solo necesita: la lista de adyacencia del grafo CDT termalizado.
#
# USO:
#   1. Obtener lista de adyacencia del ensamble CDT (archivo .txt o .csv)
#   2. Ejecutar este script sobre esa lista
#   3. El script extrae shells radiales, aplica protocolo octant-blind
#      y mide d_s via heat kernel
#
# PREDICCIÓN GRU:
#   Si el protocolo octant-blind (colapso de shells) da d_s ∈ [0.95, 1.05]
#   con separación >5σ respecto al grafo completo → GRU verificado en CDT.
#
# FORMATO DE ENTRADA:
#   Archivo de texto con una arista por línea: "nodo_a nodo_b"
#   Ejemplo:
#     0 1
#     0 5
#     1 2
#     ...
# =============================================================================

import sys, io
if hasattr(sys.stdout, 'buffer'):
    sys.stdout = io.TextIOWrapper(
        sys.stdout.buffer, encoding='utf-8', line_buffering=True)

import numpy as np
import networkx as nx
from scipy.optimize import curve_fit
import time
import argparse

SEED      = 42
N_WALKS   = 4000
SIGMA_MAX = 120
WINDOW    = (6, 100)


# =============================================================================
# CARGA DEL GRAFO CDT
# =============================================================================

def load_cdt_graph(filepath):
    """
    Carga lista de adyacencia desde archivo de texto.
    Formato: una arista por línea 'nodo_a nodo_b'
    """
    G = nx.Graph()
    with open(filepath, 'r') as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith('#'): continue
            parts = line.split()
            if len(parts) >= 2:
                G.add_edge(parts[0], parts[1])
    print(f"  Grafo CDT cargado: {G.number_of_nodes()} nodos, "
          f"{G.number_of_edges()} aristas")
    return G


def build_demo_cdt_graph(n_radial=120, n_layers=10, seed=SEED):
    """
    Genera un grafo CDT sintético de demostración.
    Úsalo si no tienes un grafo CDT real disponible.
    """
    rng = np.random.default_rng(seed)
    G = nx.Graph()
    def node(t, r): return f"{t}_{r}"

    for t in range(n_layers):
        for r in range(n_radial):
            G.add_node(node(t, r))
        # Conexiones radiales S¹
        for r in range(n_radial):
            G.add_edge(node(t, r), node(t, (r+1) % n_radial))
        # Conexiones temporales (CDT causal)
        if t < n_layers - 1:
            for r in range(n_radial):
                if rng.random() < 0.6:
                    G.add_edge(node(t, r), node(t+1, r))
                if rng.random() < 0.2:
                    G.add_edge(node(t, r), node(t+1, (r+1) % n_radial))

    print(f"  Grafo CDT demo: {G.number_of_nodes()} nodos, "
          f"{G.number_of_edges()} aristas")
    return G


# =============================================================================
# PROTOCOLO OCTANT-BLIND
# =============================================================================

def extract_radial_shells(G, origin):
    """
    BFS desde el origen para extraer shells radiales.
    Cada shell k contiene los nodos a distancia BFS k del origen.
    """
    lengths = nx.single_source_shortest_path_length(G, origin)
    max_k = max(lengths.values())

    shells = {k: [] for k in range(max_k + 1)}
    for node, dist in lengths.items():
        shells[dist].append(node)

    print(f"  Shells extraídas: {len(shells)} niveles radiales")
    print(f"  Tamaño máximo de shell: {max(len(s) for s in shells.values())} nodos")
    return shells


def apply_octant_blind(G, shells, lam_r=1.0, lam_t=0.0, seed=SEED):
    """
    Protocolo octant-blind: colapsa cada shell a un nodo representante.
    El grafo resultante es la cadena radial 1D pura.

    lam_r = probabilidad de conexión entre shells adyacentes
    lam_t = probabilidad de conexión temporal entre capas
            (solo relevante si el grafo tiene estructura temporal explícita)
    """
    rng = np.random.default_rng(seed)
    G_blind = nx.Graph()

    max_k = max(shells.keys())

    # Un nodo por shell
    for k in range(max_k + 1):
        G_blind.add_node(k)

    # Conexiones radiales entre shells adyacentes
    for k in range(1, max_k + 1):
        if rng.random() < lam_r:
            G_blind.add_edge(k, k-1)

    # Garantizar conectividad mínima
    for k in range(1, max_k + 1):
        if not G_blind.has_edge(k, k-1):
            G_blind.add_edge(k, k-1)

    print(f"  Grafo octant-blind: {G_blind.number_of_nodes()} nodos "
          f"(1 por shell), {G_blind.number_of_edges()} aristas")
    return G_blind


def apply_full_s1(G, shells, lam_theta=1.0, lam_t=0.6, seed=SEED):
    """
    Grafo completo con estructura S¹ por shell.
    Para comparación con el caso octant-blind.

    lam_theta = activación angular (S¹ dentro de cada shell)
    lam_t = activación temporal (entre shells adyacentes)
    """
    rng = np.random.default_rng(seed)
    G_full = nx.Graph()
    max_k = max(shells.keys())

    for k, nodes in shells.items():
        for nd in nodes:
            G_full.add_node(nd)
        # Cierre S¹ dentro del shell
        n = len(nodes)
        if n > 1 and lam_theta > 0:
            for j in range(n):
                if rng.random() < lam_theta:
                    G_full.add_edge(nodes[j], nodes[(j+1) % n])

    # Conexiones radiales entre shells
    for k in range(1, max_k + 1):
        curr = shells[k]; prev = shells[k-1]
        for j, nd in enumerate(curr):
            pi = int(j * len(prev) / len(curr)) % len(prev)
            G_full.add_edge(nd, prev[pi])

    print(f"  Grafo S¹ completo: {G_full.number_of_nodes()} nodos, "
          f"{G_full.number_of_edges()} aristas")
    return G_full


# =============================================================================
# HEAT KERNEL Y AJUSTE
# =============================================================================

def heat_kernel(G, origin, n_walks=N_WALKS, sigma_max=SIGMA_MAX, seed=SEED):
    if origin not in G:
        origin = list(G.nodes())[0]
    rng = np.random.default_rng(seed)
    adj = {n: list(G.neighbors(n)) for n in G.nodes()}
    P   = np.zeros(sigma_max)
    for _ in range(n_walks):
        cur = origin
        for step in range(sigma_max):
            nb = adj[cur]
            if not nb: break
            cur = nb[rng.integers(len(nb))]
            if cur == origin: P[step] += 1
    return P / n_walks


def fit_ds(sigma_arr, P, i_lo=6, i_hi=100):
    s, p = sigma_arr[i_lo:i_hi], P[i_lo:i_hi]
    mask = p > 1e-10
    if mask.sum() < 5: return None, None, False
    try:
        popt, pcov = curve_fit(
            lambda s, A, a: A * s**(-a),
            s[mask], p[mask],
            p0=[p[mask][0], 0.5],
            maxfev=10000,
            bounds=(0, [np.inf, 3.0])
        )
        return 2.0*popt[1], 2.0*np.sqrt(np.diag(pcov))[1], True
    except: return None, None, False


# =============================================================================
# EXPERIMENTO PRINCIPAL
# =============================================================================

def run_gru_on_cdt(G, origin=None):
    """
    Aplica el protocolo GRU completo sobre un grafo CDT.
    Compara: grafo completo vs octant-blind.
    """
    sigma_arr = np.arange(1, SIGMA_MAX+1, dtype=float)

    # Seleccionar origen (nodo central o especificado)
    if origin is None:
        # Intentar nodo con mayor grado como origen bulk
        origin = max(G.degree(), key=lambda x: x[1])[0]
    print(f"  Nodo origen: {origin}")

    # Extraer shells
    if not nx.is_connected(G):
        print("  ⚠ Grafo desconectado — usando componente más grande")
        G = G.subgraph(max(nx.connected_components(G), key=len)).copy()

    shells = extract_radial_shells(G, origin)

    print("\n" + "="*70)
    print("GRU PROTOCOLO OCTANT-BLIND — Resultados".center(70))
    print("="*70)
    print(f"  {'Configuración':<30} | {'d_s ± err':>14} | {'Criterio':>10}")
    print("  " + "-"*58)

    results = {}

    # 1. Grafo completo CDT
    t0 = time.time()
    P_full = heat_kernel(G, origin)
    ds_full, err_full, ok = fit_ds(sigma_arr, P_full, *WINDOW)
    elapsed = time.time()-t0
    if ok:
        print(f"  {'CDT completo':<30} | {ds_full:.4f}±{err_full:.3f}   | "
              f"{'referencia'} ({elapsed:.1f}s)")
    results['full'] = (ds_full, err_full)

    # 2. Octant-blind (colapso de shells)
    G_blind = apply_octant_blind(G, shells)
    origin_blind = 0  # shell 0 = origen
    t0 = time.time()
    P_blind = heat_kernel(G_blind, origin_blind, sigma_max=min(SIGMA_MAX, G_blind.number_of_nodes()-1))
    sm = min(SIGMA_MAX, G_blind.number_of_nodes()-1)
    sa = np.arange(1, sm+1, dtype=float)
    i_hi = min(WINDOW[1], sm-2)
    ds_blind, err_blind, ok = fit_ds(sa, P_blind, WINDOW[0], i_hi)
    elapsed = time.time()-t0
    if ok:
        tag = "✓ GRU" if abs(ds_blind-1.0) < 0.1 else "⚠ revisar"
        print(f"  {'Octant-blind (GRU puro)':<30} | {ds_blind:.4f}±{err_blind:.3f}   | "
              f"{tag} ({elapsed:.1f}s)")
    results['blind'] = (ds_blind, err_blind)

    # Separación estadística
    if results['full'][0] and results['blind'][0]:
        df = results['full'][0]; ef = results['full'][1]
        db = results['blind'][0]; eb = results['blind'][1]
        sep = abs(df - db) / np.sqrt(ef**2 + eb**2)
        print(f"\n  Separación CDT completo vs GRU puro: {sep:.1f}σ")
        if sep >= 5:
            print(f"  ✓ DIFERENCIA CATEGÓRICA (≥5σ) — GRU verificado en CDT")
        else:
            print(f"  ⚠ Separación insuficiente — más estadística requerida")

    print("="*70)
    print(f"\n  DOI: 10.5281/zenodo.20472915")
    return results


# =============================================================================
# MAIN
# =============================================================================

if __name__ == "__main__":
    print("\n" + "="*70)
    print("GRU v1.8.7 — Protocolo Octant-Blind para CDT".center(70))
    print("="*70)

    import sys
    if len(sys.argv) > 1:
        # Modo real: cargar grafo CDT desde archivo
        filepath = sys.argv[1]
        print(f"\n  Cargando grafo CDT desde: {filepath}")
        G = load_cdt_graph(filepath)
        origin = sys.argv[2] if len(sys.argv) > 2 else None
    else:
        # Modo demo: grafo CDT sintético
        print("\n  Modo DEMO — grafo CDT sintético (N_RADIAL=120, N_LAYERS=10)")
        print("  Para usar con CDT real: python GRU_CDT_postprocessing.py adjacency.txt")
        G = build_demo_cdt_graph()
        origin = "5_60"  # nodo central aproximado

    run_gru_on_cdt(G, origin)
