"""
Compute and visualize resistance distance matrix for Slipnet concepts (Figure 3)
Resistance distance considers all paths between nodes, weighted by conductance
"""

import matplotlib.pyplot as plt
import numpy as np
import networkx as nx
from scipy.linalg import pinv

# Define key Slipnet nodes
key_nodes = [
    'a', 'b', 'c',
    'letterCategory',
    'left', 'right',
    'leftmost', 'rightmost',
    'first', 'last',
    'predecessor', 'successor', 'sameness',
    'identity', 'opposite',
]

# Create graph with resistances (link lengths)
G = nx.Graph()
edges = [
    # Letters to category
    ('a', 'letterCategory', 97),
    ('b', 'letterCategory', 97),
    ('c', 'letterCategory', 97),
    # Sequential relationships
    ('a', 'b', 50),
    ('b', 'c', 50),
    # Bond types
    ('predecessor', 'successor', 60),
    ('sameness', 'identity', 50),
    # Opposite relations
    ('left', 'right', 80),
    ('first', 'last', 80),
    ('leftmost', 'rightmost', 90),
    # Slippable connections
    ('left', 'leftmost', 90),
    ('right', 'rightmost', 90),
    ('first', 'leftmost', 100),
    ('last', 'rightmost', 100),
    # Abstract relations
    ('identity', 'opposite', 70),
    ('predecessor', 'identity', 60),
    ('successor', 'identity', 60),
    ('sameness', 'identity', 40),
]

for src, dst, link_len in edges:
    # Resistance = link length, conductance = 1/resistance
    G.add_edge(src, dst, resistance=link_len, conductance=1.0/link_len)

# Only keep nodes that are in our key list and connected
connected_nodes = [n for n in key_nodes if n in G.nodes()]

def compute_resistance_distance(G, nodes):
    """Compute resistance distance matrix using graph Laplacian"""
    # Create mapping from nodes to indices
    node_to_idx = {node: i for i, node in enumerate(nodes)}
    n = len(nodes)

    # Build Laplacian matrix (weighted by conductance)
    L = np.zeros((n, n))
    for i, node_i in enumerate(nodes):
        for j, node_j in enumerate(nodes):
            if G.has_edge(node_i, node_j):
                conductance = G[node_i][node_j]['conductance']
                L[i, j] = -conductance
                L[i, i] += conductance

    # Compute pseudo-inverse of Laplacian
    try:
        L_pinv = pinv(L)
    except:
        # Fallback: use shortest path distances
        return compute_shortest_path_matrix(G, nodes)

    # Resistance distance: R_ij = L+_ii + L+_jj - 2*L+_ij
    R = np.zeros((n, n))
    for i in range(n):
        for j in range(n):
            R[i, j] = L_pinv[i, i] + L_pinv[j, j] - 2 * L_pinv[i, j]

    return R

def compute_shortest_path_matrix(G, nodes):
    """Compute shortest path distance matrix"""
    n = len(nodes)
    D = np.zeros((n, n))
    for i, node_i in enumerate(nodes):
        for j, node_j in enumerate(nodes):
            if i == j:
                D[i, j] = 0
            else:
                try:
                    path = nx.shortest_path(G, node_i, node_j, weight='resistance')
                    D[i, j] = sum(G[path[k]][path[k+1]]['resistance']
                                 for k in range(len(path)-1))
                except nx.NetworkXNoPath:
                    D[i, j] = 1000  # Large value for disconnected nodes
    return D

# Compute both matrices
R_resistance = compute_resistance_distance(G, connected_nodes)
R_shortest = compute_shortest_path_matrix(G, connected_nodes)

# Create visualization
fig, axes = plt.subplots(1, 2, figsize=(16, 7))

# Left: Resistance distance
ax_left = axes[0]
im_left = ax_left.imshow(R_resistance, cmap='YlOrRd', aspect='auto')
ax_left.set_xticks(range(len(connected_nodes)))
ax_left.set_yticks(range(len(connected_nodes)))
ax_left.set_xticklabels(connected_nodes, rotation=45, ha='right', fontsize=9)
ax_left.set_yticklabels(connected_nodes, fontsize=9)
ax_left.set_title('Resistance Distance Matrix\n(Considers all paths, weighted by conductance)',
                 fontsize=12, fontweight='bold')
cbar_left = plt.colorbar(im_left, ax=ax_left, fraction=0.046, pad=0.04)
cbar_left.set_label('Resistance Distance', rotation=270, labelpad=20)

# Add grid
ax_left.set_xticks(np.arange(len(connected_nodes))-0.5, minor=True)
ax_left.set_yticks(np.arange(len(connected_nodes))-0.5, minor=True)
ax_left.grid(which='minor', color='gray', linestyle='-', linewidth=0.5)

# Right: Shortest path distance
ax_right = axes[1]
im_right = ax_right.imshow(R_shortest, cmap='YlOrRd', aspect='auto')
ax_right.set_xticks(range(len(connected_nodes)))
ax_right.set_yticks(range(len(connected_nodes)))
ax_right.set_xticklabels(connected_nodes, rotation=45, ha='right', fontsize=9)
ax_right.set_yticklabels(connected_nodes, fontsize=9)
ax_right.set_title('Shortest Path Distance Matrix\n(Only considers single best path)',
                  fontsize=12, fontweight='bold')
cbar_right = plt.colorbar(im_right, ax=ax_right, fraction=0.046, pad=0.04)
cbar_right.set_label('Shortest Path Distance', rotation=270, labelpad=20)

# Add grid
ax_right.set_xticks(np.arange(len(connected_nodes))-0.5, minor=True)
ax_right.set_yticks(np.arange(len(connected_nodes))-0.5, minor=True)
ax_right.grid(which='minor', color='gray', linestyle='-', linewidth=0.5)

plt.suptitle('Resistance Distance vs Shortest Path Distance for Slipnet Concepts\n' +
             'Lower values = easier slippage between concepts',
             fontsize=14, fontweight='bold')
plt.tight_layout()

plt.savefig('figure3_resistance_distance.pdf', dpi=300, bbox_inches='tight')
plt.savefig('figure3_resistance_distance.png', dpi=300, bbox_inches='tight')
print("Generated figure3_resistance_distance.pdf and .png")
plt.close()

# Create additional plot: Slippability based on resistance distance
fig2, ax = plt.subplots(figsize=(10, 6))

# Select some interesting concept pairs
concept_pairs = [
    ('left', 'right', 'Opposite directions'),
    ('first', 'last', 'Opposite positions'),
    ('left', 'leftmost', 'Direction to position'),
    ('predecessor', 'successor', 'Sequential relations'),
    ('a', 'b', 'Adjacent letters'),
    ('a', 'c', 'Non-adjacent letters'),
]

# Compute slippability for different temperatures
temperatures = np.linspace(10, 90, 50)
alpha_values = 0.1 * (100 - temperatures) / 50  # Alpha increases as temp decreases

for src, dst, label in concept_pairs:
    if src in connected_nodes and dst in connected_nodes:
        i = connected_nodes.index(src)
        j = connected_nodes.index(dst)
        R_ij = R_resistance[i, j]

        # Proposed slippability: 100 * exp(-alpha * R_ij)
        slippabilities = 100 * np.exp(-alpha_values * R_ij)
        ax.plot(temperatures, slippabilities, linewidth=2, label=label, marker='o', markersize=3)

ax.set_xlabel('Temperature', fontsize=12)
ax.set_ylabel('Slippability', fontsize=12)
ax.set_title('Temperature-Dependent Slippability using Resistance Distance\n' +
             'Formula: slippability = 100 × exp(-α × R_ij), where α ∝ (100-T)',
             fontsize=12, fontweight='bold')
ax.legend(fontsize=10, loc='upper left')
ax.grid(True, alpha=0.3)
ax.set_xlim([10, 90])
ax.set_ylim([0, 105])

# Add annotations
ax.axvspan(10, 30, alpha=0.1, color='blue', label='Low temp (exploitation)')
ax.axvspan(70, 90, alpha=0.1, color='red', label='High temp (exploration)')
ax.text(20, 95, 'Low temperature\n(restrictive slippage)', fontsize=9, ha='center')
ax.text(80, 95, 'High temperature\n(liberal slippage)', fontsize=9, ha='center')

plt.tight_layout()
plt.savefig('slippability_temperature.pdf', dpi=300, bbox_inches='tight')
plt.savefig('slippability_temperature.png', dpi=300, bbox_inches='tight')
print("Generated slippability_temperature.pdf and .png")
plt.close()
