# =============================================================================
# GRU — Fix artefacto de bipartición en λ_t scan
# Problema: grafo diagonal tiene estructura bipartita → P(σ=impar)≈0
# Fix: usar solo pasos pares para el ajuste
# También: comparar con grafo NO bipartito (conexiones dentro de capa)
# =============================================================================

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

SEED      = 42
N_RADIAL  = 120
N_LAYERS  = 120
N_WALKS   = 5000
SIGMA_MAX = 120

def build_graph_time_diagonal(n_radial, n_layers, lam_t=1.0, seed=SEED):
    rng = np.random.default_rng(seed)
    G = nx.Graph()
    for t in range(n_layers):
        for r in range(n_radial + 1):
            G.add_node(f"{t}_{r}")
    for t in range(n_layers):
        for r in range(n_radial):
            G.add_edge(f"{t}_{r}", f"{t}_{r+1}")
        if t < n_layers - 1:
            for r in range(n_radial + 1):
                if rng.random() < lam_t:
                    G.add_edge(f"{t}_{r}", f"{t+1}_{r}")
                if r > 0 and rng.random() < lam_t/2:
                    G.add_edge(f"{t}_{r}", f"{t+1}_{r-1}")
                if r < n_radial and rng.random() < lam_t/2:
                    G.add_edge(f"{t}_{r}", f"{t+1}_{r+1}")
    return G

def build_graph_no_bipartito(n_radial, n_layers, lam_t=1.0, seed=SEED):
    """
    Mismo grafo pero con self-loops en cada nodo.
    Rompe la bipartición: el caminante puede quedarse en su nodo.
    Equivale a 'lazy random walk' con p_stay=0.5
    """
    rng = np.random.default_rng(seed)
    G = nx.Graph()
    for t in range(n_layers):
        for r in range(n_radial + 1):
            G.add_node(f"{t}_{r}")
            G.add_edge(f"{t}_{r}", f"{t}_{r}")  # self-loop
    for t in range(n_layers):
        for r in range(n_radial):
            G.add_edge(f"{t}_{r}", f"{t}_{r+1}")
        if t < n_layers - 1:
            for r in range(n_radial + 1):
                if rng.random() < lam_t:
                    G.add_edge(f"{t}_{r}", f"{t+1}_{r}")
                if r > 0 and rng.random() < lam_t/2:
                    G.add_edge(f"{t}_{r}", f"{t+1}_{r-1}")
                if r < n_radial and rng.random() < lam_t/2:
                    G.add_edge(f"{t}_{r}", f"{t+1}_{r+1}")
    return G

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 = {}
    for n in G.nodes():
        nb = list(G.neighbors(n))
        # incluir self-loops como vecinos repetidos
        adj[n] = nb if nb else [n]
    P = np.zeros(sigma_max)
    for _ in range(n_walks):
        cur = origin
        for step in range(sigma_max):
            nb = adj[cur]
            cur = nb[rng.integers(len(nb))]
            if cur == origin: P[step] += 1
    return P / n_walks

def fit_ds_todos(P, sigma_arr, i_lo=6, i_hi=100):
    """Fit estándar — todos los pasos"""
    s, p = sigma_arr[i_lo:i_hi], P[i_lo:i_hi]
    mask = p > 1e-10
    if mask.sum() < 5: return None, None
    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]))
        return 2*popt[1], 2*np.sqrt(np.diag(pcov))[1]
    except: return None, None

def fit_ds_pares(P, sigma_arr, i_lo=6, i_hi=100):
    """
    FIX bipartición: usar solo pasos PARES.
    El caminante en grafo bipartito solo regresa en pasos pares.
    """
    mask_par = (sigma_arr.astype(int) % 2 == 0)
    s = sigma_arr[mask_par]
    p = P[mask_par]
    # Recortar ventana
    idx_lo = np.searchsorted(s, sigma_arr[i_lo])
    idx_hi = np.searchsorted(s, sigma_arr[min(i_hi, len(sigma_arr)-1)])
    s = s[idx_lo:idx_hi]
    p = p[idx_lo:idx_hi]
    mask = p > 1e-10
    if mask.sum() < 4: return None, None
    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]))
        return 2*popt[1], 2*np.sqrt(np.diag(pcov))[1]
    except: return None, None

# ── EXPERIMENTO PRINCIPAL ──
sigma_arr = np.arange(1, SIGMA_MAX+1, dtype=float)

print("="*72)
print("Fix Bipartición — λ_t scan con tres métodos de ajuste")
print("="*72)
print(f"N_RADIAL={N_RADIAL}, N_LAYERS={N_LAYERS}, N_WALKS={N_WALKS}")
print()
print(f"{'λ_t':>5} | {'d_s todos':>10} | {'d_s pares':>10} | {'d_s lazy':>10} | Interpretación")
print("-"*72)

for lam_t in [0.0, 0.10, 0.20, 0.30, 0.40, 0.50, 0.60, 0.70, 1.00]:
    # Grafo original
    if lam_t == 0.0:
        G = nx.path_graph(N_RADIAL+1)
        mapping = {i: f"0_{i}" for i in range(N_RADIAL+1)}
        G = nx.relabel_nodes(G, mapping)
    else:
        G = build_graph_time_diagonal(N_RADIAL, N_LAYERS, lam_t)

    P = heat_kernel(G, "0_0")
    ds_todos, _ = fit_ds_todos(P, sigma_arr)
    ds_pares, _ = fit_ds_pares(P, sigma_arr)

    # Grafo lazy (rompe bipartición)
    if lam_t == 0.0:
        G_lazy = nx.path_graph(N_RADIAL+1)
        G_lazy = nx.relabel_nodes(G_lazy,
                    {i: f"0_{i}" for i in range(N_RADIAL+1)})
    else:
        G_lazy = build_graph_no_bipartito(N_RADIAL, N_LAYERS, lam_t)

    P_lazy = heat_kernel(G_lazy, "0_0")
    ds_lazy, _ = fit_ds_todos(P_lazy, sigma_arr)

    def fmt(v): return f"{v:.4f}" if v else "N/A"

    if 0.20 <= lam_t <= 0.40:
        interp = "← zona analizar"
    elif lam_t == 0.0:
        interp = "GRU puro"
    elif 0.50 <= lam_t <= 0.70:
        interp = "Giasemidis ✓"
    elif lam_t == 1.00:
        interp = "sobreconectado"
    else:
        interp = "sub-2"

    print(f"  {lam_t:>3.2f} | {fmt(ds_todos):>10} | {fmt(ds_pares):>10} | "
          f"{fmt(ds_lazy):>10} | {interp}")

print()
print("="*72)
print("INTERPRETACIÓN:")
print("  d_s todos:  fit estándar (afectado por bipartición)")
print("  d_s pares:  fit solo pasos pares (corrige bipartición)")
print("  d_s lazy:   grafo con self-loops (rompe bipartición)")
print()
print("  Si d_s_pares ≈ d_s_lazy para 0.20-0.40 → artefacto puro")
print("  Si d_s_pares sigue siendo inestable → fenómeno físico real")
print("="*72)
