"""
Simulate and visualize activation spreading in the Slipnet (Figure 2)
Shows differential decay rates based on conceptual depth
"""

import matplotlib.pyplot as plt
import numpy as np
import networkx as nx
from matplotlib.gridspec import GridSpec

# Define simplified Slipnet structure
nodes_with_depth = {
    'sameness': 80,  # Initial activation source
    'samenessGroup': 80,
    'identity': 90,
    'letterCategory': 30,
    'a': 10, 'b': 10, 'c': 10,
    'predecessor': 50,
    'successor': 50,
    'bondCategory': 80,
    'left': 40,
    'right': 40,
}

edges_with_strength = [
    ('sameness', 'samenessGroup', 30),
    ('sameness', 'identity', 50),
    ('sameness', 'bondCategory', 40),
    ('samenessGroup', 'letterCategory', 50),
    ('letterCategory', 'a', 97),
    ('letterCategory', 'b', 97),
    ('letterCategory', 'c', 97),
    ('predecessor', 'bondCategory', 60),
    ('successor', 'bondCategory', 60),
    ('sameness', 'bondCategory', 30),
    ('left', 'right', 80),
]

# Create graph
G = nx.Graph()
for node, depth in nodes_with_depth.items():
    G.add_node(node, depth=depth, activation=0.0, buffer=0.0)

for src, dst, link_len in edges_with_strength:
    G.add_edge(src, dst, length=link_len, strength=100-link_len)

# Initial activation
G.nodes['sameness']['activation'] = 100.0

# Simulate activation spreading with differential decay
def simulate_spreading(G, num_steps):
    history = {node: [] for node in G.nodes()}

    for step in range(num_steps):
        # Record current state
        for node in G.nodes():
            history[node].append(G.nodes[node]['activation'])

        # Decay phase
        for node in G.nodes():
            depth = G.nodes[node]['depth']
            activation = G.nodes[node]['activation']
            decay_rate = (100 - depth) / 100.0
            G.nodes[node]['buffer'] -= activation * decay_rate

        # Spreading phase (if fully active)
        for node in G.nodes():
            if G.nodes[node]['activation'] >= 95.0:
                for neighbor in G.neighbors(node):
                    strength = G[node][neighbor]['strength']
                    G.nodes[neighbor]['buffer'] += strength

        # Apply buffer
        for node in G.nodes():
            G.nodes[node]['activation'] = max(0, min(100,
                G.nodes[node]['activation'] + G.nodes[node]['buffer']))
            G.nodes[node]['buffer'] = 0.0

    return history

# Run simulation
history = simulate_spreading(G, 15)

# Create visualization
fig = plt.figure(figsize=(16, 10))
gs = GridSpec(2, 3, figure=fig, hspace=0.3, wspace=0.3)

# Time snapshots: t=0, t=5, t=10
time_points = [0, 5, 10]
positions = nx.spring_layout(G, k=1.5, iterations=50, seed=42)

for idx, t in enumerate(time_points):
    ax = fig.add_subplot(gs[0, idx])

    # Get activations at time t
    node_colors = [history[node][t] for node in G.nodes()]

    # Draw graph
    nx.draw_networkx_edges(G, positions, alpha=0.3, width=2, ax=ax)
    nodes_drawn = nx.draw_networkx_nodes(G, positions,
                                         node_color=node_colors,
                                         node_size=800,
                                         cmap='hot',
                                         vmin=0, vmax=100,
                                         ax=ax)
    nx.draw_networkx_labels(G, positions, font_size=8, font_weight='bold', ax=ax)

    ax.set_title(f'Time Step {t}', fontsize=12, fontweight='bold')
    ax.axis('off')

    if idx == 2:  # Add colorbar to last subplot
        cbar = plt.colorbar(nodes_drawn, ax=ax, fraction=0.046, pad=0.04)
        cbar.set_label('Activation', rotation=270, labelpad=15)

# Bottom row: activation time series for key nodes
ax_time = fig.add_subplot(gs[1, :])

# Plot activation over time for nodes with different depths
nodes_to_plot = [
    ('sameness', 'Deep (80)', 'red'),
    ('predecessor', 'Medium (50)', 'orange'),
    ('letterCategory', 'Shallow (30)', 'blue'),
    ('a', 'Very Shallow (10)', 'green'),
]

time_steps = range(15)
for node, label, color in nodes_to_plot:
    ax_time.plot(time_steps, history[node], marker='o', label=label,
                linewidth=2, color=color)

ax_time.set_xlabel('Time Steps', fontsize=12)
ax_time.set_ylabel('Activation Level', fontsize=12)
ax_time.set_title('Activation Dynamics: Differential Decay by Conceptual Depth',
                 fontsize=13, fontweight='bold')
ax_time.legend(title='Node (Depth)', fontsize=10)
ax_time.grid(True, alpha=0.3)
ax_time.set_xlim([0, 14])
ax_time.set_ylim([0, 105])

# Add annotation
ax_time.annotate('Deep nodes decay slowly\n(high conceptual depth)',
                xy=(10, history['sameness'][10]), xytext=(12, 70),
                arrowprops=dict(arrowstyle='->', color='red', lw=1.5),
                fontsize=10, color='red')
ax_time.annotate('Shallow nodes decay rapidly\n(low conceptual depth)',
                xy=(5, history['a'][5]), xytext=(7, 35),
                arrowprops=dict(arrowstyle='->', color='green', lw=1.5),
                fontsize=10, color='green')

fig.suptitle('Activation Spreading with Differential Decay\n' +
             'Formula: decay = activation × (100 - conceptual_depth) / 100',
             fontsize=14, fontweight='bold')

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