import torch
import numpy as np
import matplotlib.pyplot as plt

# Configuration
CONFIG = {
    "n_agents": 10,
    "steps": 50,
    "grid_size": 10,
    "target": torch.tensor([5.0, 5.0])
}

# Tracker for entropy and convergence
class MetricsTracker:
    def __init__(self):
        self.entropies = []
        self.distances = []

    def record(self, positions, target):
        center = positions.mean(dim=0)
        distances = torch.norm(positions - center, dim=1)
        entropy = distances.std().item()
        self.entropies.append(entropy)

        dist_to_target = torch.norm(positions - target, dim=1).mean().item()
        self.distances.append(dist_to_target)

    def plot(self):
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
        ax1.plot(self.entropies, label="Entropy")
        ax1.set_title("Entropy Over Time")
        ax1.set_xlabel("Steps")
        ax1.set_ylabel("Entropy")

        ax2.plot(self.distances, label="Convergence")
        ax2.set_title("Convergence to Target")
        ax2.set_xlabel("Steps")
        ax2.set_ylabel("Avg Distance to Target")
        plt.tight_layout()
        plt.show()

# Agent class
class Agent:
    def __init__(self, position):
        self.position = torch.tensor(position, dtype=torch.float32)
        self.dna = ["∑", "Ω", "ε₀"]
        self.entropy_thresholds = [1.5, 1.0]  # Dynamic thresholds for evolution

    def emit_symbol(self):
        return self.dna[int(torch.norm(self.position - CONFIG["target"]).item()) % 3]

    def evolve_dna(self, entropy):
        if entropy > self.entropy_thresholds[0]:
            self.dna[0] = "∑"
        elif entropy > self.entropy_thresholds[1]:
            self.dna[0] = "Ω"
        else:
            self.dna[0] = "ε₀"

    def update(self, field, target):
        self.position += 0.3 * (field - self.position) + 0.3 * (target - self.position)

# Morphic Memory
class MorphicMemoryCore:
    def __init__(self, decay_rate=0.1):
        self.memory = []
        self.decay_rate = decay_rate

    def update(self, state):
        self.memory.append(state.clone())
        if len(self.memory) > 50:
            self.memory.pop(0)

    def compute_field(self):
        if not self.memory:
            return torch.zeros(2)
        weights = torch.exp(-self.decay_rate * torch.arange(len(self.memory), 0, -1, dtype=torch.float32))
        weights /= weights.sum()
        field = sum(w * s.mean(dim=0) for w, s in zip(weights, self.memory))
        return field

# TMFS system
class TMFS:
    def __init__(self, use_field=True):
        self.use_field = use_field
        self.agents = [Agent([np.random.rand()*CONFIG["grid_size"], np.random.rand()*CONFIG["grid_size"]]) for _ in range(CONFIG["n_agents"])]
        self.memory = MorphicMemoryCore()
        self.target = CONFIG["target"]
        self.tracker = MetricsTracker()

    def step(self):
        positions = torch.stack([a.position for a in self.agents])
        entropy_val = torch.norm(positions - positions.mean(dim=0), dim=1).std().item()

        if self.use_field:
            self.memory.update(positions)
            field = self.memory.compute_field()
        else:
            field = torch.zeros(2)

        for agent in self.agents:
            agent.evolve_dna(entropy_val)
            agent.update(field, self.target)

        self.tracker.record(torch.stack([a.position for a in self.agents]), self.target)

    def run(self):
        trajectories = []
        for _ in range(CONFIG["steps"]):
            self.step()
            trajectories.append(torch.stack([a.position for a in self.agents]).clone())
        return torch.stack(trajectories), self.tracker

# Run and visualize
def run_simulation():
    sim = TMFS(use_field=True)
    traj, tracker = sim.run()

    # Trajectories
    plt.figure(figsize=(12, 6))
    for i in range(CONFIG["n_agents"]):
        plt.plot(traj[:, i, 0], traj[:, i, 1], label=f"Agent {i+1}")
    plt.scatter([CONFIG["target"][0]], [CONFIG["target"][1]], c="red", label="Target")
    plt.title("Agent Trajectories With Morphic Field")
    plt.legend()
    plt.grid(True)
    plt.show()

    tracker.plot()

    print("Simulation complete. Final symbolic DNA:")
    print([agent.dna for agent in sim.agents])

if __name__ == "__main__":
    run_simulation()