import random
import json
import csv
import io

class EvoTransformer:
    def __init__(self):
        self.history = []
        self.current = {}

    def reset(self):
        self.history = []
        self.current = {}

    def mutate(self, config):
        config["layers"] = max(1, config["layers"] + random.choice([-1, 1]))
        config["attention_heads"] = max(1, config["attention_heads"] + random.choice([-1, 1]))
        config["ffn_dim"] = max(64, config["ffn_dim"] + random.choice([-256, 256]))
        config["dropout"] = min(0.5, max(0.01, config["dropout"] + random.choice([-0.02, 0.02])))
        config["memory"] = random.choice([True, False])
        return config

    def simulate_accuracy(self, config):
        return round(random.uniform(70, 90), 2)

    def estimate_parameters(self, config):
        return round(
            config["layers"] * config["attention_heads"] * config["ffn_dim"] / 1e6,
            2
        )

    def evolve(self, generations):
        base_config = {
            "layers": 4,
            "attention_heads": 4,
            "ffn_dim": 1024,
            "dropout": 0.1,
            "memory": False
        }

        for _ in range(generations):
            base_config = self.mutate(base_config.copy())
            acc = self.simulate_accuracy(base_config)
            params = self.estimate_parameters(base_config)
            record = {"traits": base_config.copy(), "accuracy": acc, "parameters": params}
            self.history.append(record)

        self.current = self.history[-1]

    def get_best_config(self):
        return self.current

    def export_csv(self):
        output = io.StringIO()
        writer = csv.writer(output)
        writer.writerow(["Generation", "Layers", "Attention Heads", "FFN Dim", "Dropout", "Memory", "Accuracy", "Params"])
        for i, entry in enumerate(self.history):
            t = entry["traits"]
            writer.writerow([
                i+1,
                t["layers"], t["attention_heads"], t["ffn_dim"],
                t["dropout"], t["memory"],
                entry["accuracy"], entry["parameters"]
            ])
        return output.getvalue()

    def export_json(self):
        return json.dumps(self.history, indent=2)