import numpy as np
import networkx as nx
import matplotlib.pyplot as plt

# === ENTITY CLASS: DRIFTING, POLARIZATION, AND CONNECTIONS ===
class MeshNode:
    def __init__(self, entity_id):
        self.entity_id = entity_id
        self.position = np.random.randn(3) * 5  # Position in N-space
        self.energy_level = np.random.uniform(0.5, 1.5)  # Flow potential
        self.connections = {}  # Stores linked nodes

    def update_energy(self):
        """Distribute energy based on network influence."""
        if not self.connections:
            return

        avg_energy = np.mean([node.energy_level for node in self.connections.values()])
        self.energy_level = (self.energy_level + avg_energy) / 2  # Flow stabilization

    def connect(self, other_node, weight=1.0):
        """Form a connection in the mesh."""
        self.connections[other_node.entity_id] = weight


# === BRIDGE CLASS: CONNECTIONS THAT REINFORCE OR WEAKEN ===
class MeshLink:
    def __init__(self, node_a, node_b):
        self.nodes = (node_a, node_b)
        self.elasticity = 1.0  # Stretching factor
        self.energy_transfer_rate = 0.1  # How much energy flows through the link

    def update_link(self):
        """Allow the link to stretch, strengthen, or dissolve dynamically."""
        node_a, node_b = self.nodes
        energy_delta = abs(node_a.energy_level - node_b.energy_level)

        if energy_delta > 0.5:
            self.elasticity *= 0.95  # Weakening under stress
        else:
            self.elasticity *= 1.05  # Strengthening under equilibrium

        if self.elasticity < 0.5:
            print(f"❌ Link between {node_a.entity_id} and {node_b.entity_id} has dissolved.")
            return False  # Remove link
        return True


# === POLARIZATION TRACKER: MEASURING STABILITY ===
class PolarizationTracker:
    def __init__(self):
        self.polarization_map = {}

    def update_polarization(self, node):
        """Track how much a node has shifted toward a stable state."""
        connection_strength = np.mean([weight for weight in node.connections.values()]) if node.connections else 0
        energy_bias = np.abs(node.energy_level - np.mean([n.energy_level for n in node.connections.values()])) if node.connections else 1

        polarization_factor = connection_strength / (energy_bias + 1e-6)  # Prevent divide by zero
        self.polarization_map[node.entity_id] = polarization_factor


# === BRIDGE REINFORCEMENT: ADAPTIVE STRENGTHENING ===
class BridgeReinforcement:
    def __init__(self):
        self.stable_bridges = {}

    def reinforce(self, bridge, tracker):
        """Strengthen bridges that exist in highly polarized zones."""
        polarization_a = tracker.polarization_map.get(bridge.nodes[0].entity_id, 0)
        polarization_b = tracker.polarization_map.get(bridge.nodes[1].entity_id, 0)
        avg_polarization = (polarization_a + polarization_b) / 2

        if avg_polarization > 0.7:
            bridge.elasticity *= 1.1  # Strengthen
            self.stable_bridges[bridge.nodes] = bridge.elasticity
            print(f"🔗 Bridge {bridge.nodes[0].entity_id} ↔ {bridge.nodes[1].entity_id} reinforced.")


# === VISUALIZATION FUNCTIONS ===
def plot_mesh(nodes, links):
    """Plot the N-space mesh network."""
    graph = nx.Graph()

    for node in nodes.values():
        graph.add_node(node.entity_id, pos=(node.position[0], node.position[1]))

    for link in links:
        graph.add_edge(link.nodes[0].entity_id, link.nodes[1].entity_id, weight=link.elasticity)

    pos = nx.get_node_attributes(graph, 'pos')
    nx.draw(graph, pos, with_labels=True, node_color="skyblue", edge_color="gray")
    plt.title("N-Space Mesh Network")
    plt.show()


def plot_polarization(tracker):
    """Visualizes how polarized different regions of the mesh have become."""
    x, y, polarization_values = [], [], []

    for node_id, factor in tracker.polarization_map.items():
        x.append(np.random.uniform(-10, 10))  # Assign random X for visualization
        y.append(np.random.uniform(-10, 10))  # Assign random Y
        polarization_values.append(factor)

    plt.scatter(x, y, c=polarization_values, cmap="plasma", alpha=0.7)
    plt.colorbar(label="Polarization Level")
    plt.xlabel("X-Dimension")
    plt.ylabel("Y-Dimension")
    plt.title("Mesh Polarization Heatmap")
    plt.show()