import random
import matplotlib.pyplot as plt

class EvoTransformer:
    def __init__(self, config=None):
        self.config = config or {
            "layers": 4,
            "attention_heads": 4,
            "ffn_dim": 1024,
            "dropout": 0.1,
            "memory": False,
        }
        self.history = [self.config.copy()]

    def mutate(self):
        new_config = self.config.copy()
        trait = random.choice(list(new_config.keys()))

        if trait == "layers":
            new_config[trait] = max(1, new_config[trait] + random.choice([-1, 1]))
        elif trait == "attention_heads":
            new_config[trait] = random.choice([2, 4, 6, 8])
        elif trait == "ffn_dim":
            new_config[trait] = random.choice([512, 1024, 2048])
        elif trait == "dropout":
            new_config[trait] = round(min(max(0.0, new_config[trait] + random.uniform(-0.05, 0.05)), 0.5), 2)
        elif trait == "memory":
            new_config[trait] = not new_config[trait]

        self.config = new_config
        self.history.append(new_config.copy())

    def evolve(self, generations=5):
        for _ in range(generations):
            self.mutate()

    def get_history(self):
        return self.history

    def evaluate(self):
        # Simulated accuracy for demo
        score = round(random.uniform(0.85, 0.95), 4)
        return {
            "accuracy": score,
            "params": self.estimate_params()
        }

    def estimate_params(self):
        return round(10 + self.config["layers"] * self.config["ffn_dim"] * 0.001, 2)

    def plot_evolution(self):
        layers = [cfg["layers"] for cfg in self.history]
        heads = [cfg["attention_heads"] for cfg in self.history]
        ffn_dims = [cfg["ffn_dim"] for cfg in self.history]

        plt.figure(figsize=(10, 6))
        plt.plot(layers, label="Layers", marker='o', color='orange')
        plt.plot(heads, label="Attention Heads", marker='s', color='blue')
        plt.plot(ffn_dims, label="FFN Dim", marker='^', color='green')
        plt.xlabel("Generation")
        plt.ylabel("Value")
        plt.title("Evolution of Traits")
        plt.legend()
        plt.grid(True)
        plt.tight_layout()
        plt.show()


# Run test locally
if __name__ == "__main__":
    evo = EvoTransformer()
    evo.evolve(generations=8)
    evo.plot_evolution()
    print("Final Config:", evo.config)
    print("Evaluation:", evo.evaluate())