# trainer.py
import asyncio
import time
from typing import Any, Dict, List, Tuple, Optional

import morphology  # dein morphology.py
from genotype import construct, save_genotype, print_genotype
from exoself import Exoself  # deine Actor-basierte Exoself-Implementierung


class Trainer:
    """
    Stochastischer Hillclimber wie im Erlang-Original.
    - morphology_spec: entweder das morphology-Modul, ein String-Key, oder eine Callable-Morphologie
    - hidden_layer_densities: z.B. [4, 3]
    - max_attempts / eval_limit / fitness_target: Stoppbedingungen
    - experimental_file / best_file: Pfade, falls du wie im Buch speichern willst
    """

    def __init__(
            self,
            morphology_spec=morphology,
            hidden_layer_densities: List[int] = None,
            *,
            max_attempts: int = 5,
            eval_limit: float = float("inf"),
            fitness_target: float = float("inf"),
            experimental_file: Optional[str] = "experimental.json",
            best_file: Optional[str] = "best.json",
            exoself_steps_per_eval: int = 0,
    ):
        self.morphology_spec = morphology_spec
        self.hds = hidden_layer_densities or []
        self.max_attempts = max_attempts
        self.eval_limit = eval_limit
        self.fitness_target = fitness_target
        self.experimental_file = experimental_file
        self.best_file = best_file
        # Wenn deine Exoself/Cortex eine feste Anzahl Zyklen pro „Evaluation“ braucht,
        # kannst du hier default 0 lassen (Cortex entscheidet über Halt-Flag),
        # oder einen Wert setzen.
        self.exoself_steps_per_eval = exoself_steps_per_eval

        # Laufende Akkus (wie im Erlang-Code)
        self.best_fitness = float("-inf")
        self.best_genotype: Optional[Dict[str, Any]] = None

        self.eval_acc = 0
        self.cycle_acc = 0
        self.time_acc = 0.0

    async def _run_one_attempt(self) -> Tuple[float, int, int, float, Dict[str, Any]]:
        """
        Ein Trainingsversuch:
        - Genotyp konstruieren
        - Exoself laufen lassen
        - Fitness/Evals/Cycles/Time zurückgeben + den verwendeten Genotyp
        """
        print("constructing genotype...")
        geno = construct(self.morphology_spec, self.hds, file_name=self.experimental_file, add_bias=True)

        # Exoself starten und bis zum evaluation_completed warten
        fitness, evals, cycles, elapsed = await self._evaluate_with_exoself(geno)

        return fitness, evals, cycles, elapsed, geno

    async def _evaluate_with_exoself(self, genotype: Dict[str, Any]) -> Tuple[float, int, int, float]:
        """
        Startet Exoself (deine Actor-basierte Variante) und wartet,
        bis der Cortex die Evaluation abgeschlossen hat.
        Erwartete Rückgabe: fitness, evals, cycles, elapsed
        """
        print("creating exoself...")
        ex = Exoself(genotype)
        # Exoself.run() sollte idealerweise einen Tuple (fitness, evals, cycles, time)
        # liefern. Falls deine Version aktuell „backup“-Listen liefert, ersetze das hier
        # mit der passenden Logik oder benutze das „evaluation_completed“-Signal aus dem Cortex.
        fitness, evals, cycles, elapsed = await ex.run_evaluation()
        return fitness, evals, cycles, elapsed

    async def go(self):
        """
        Entspricht dem Erlang loop/…:
        Wiederholt Versuche, bis Stoppbedingung erfüllt.
        """
        attempt = 1
        while True:
            print(".........")
            print("current attempt: ", attempt)
            print(".........")
            # Stoppbedingung vor Versuch?
            if attempt > self.max_attempts or self.eval_acc >= self.eval_limit or self.best_fitness >= self.fitness_target:
                # Ausgabe wie im Buch
                if self.best_genotype and self.best_file:
                    # bestes Genotypfile ausgeben/„drucken“
                    save_genotype(self.best_file, self.best_genotype)
                    print_genotype(self.best_file)
                print(
                    f" Morphology: {getattr(self.morphology_spec, '__name__', str(self.morphology_spec))} | "
                    f"Best Fitness: {self.best_fitness} | EvalAcc: {self.eval_acc}"
                )
                # Optional: an „Benchmarker“ melden – bei dir vermutlich nicht nötig
                return {
                    "best_fitness": self.best_fitness,
                    "eval_acc": self.eval_acc,
                    "cycle_acc": self.cycle_acc,
                    "time_acc": self.time_acc,
                    "best_file": self.best_file,
                }

            print("RUN ONE ATTEMPT!")
            # --- Ein Versuch ---
            fitness, evals, cycles, elapsed, geno = await self._run_one_attempt()

            print("update akkus...")

            # Akkus updaten
            self.eval_acc += evals
            self.cycle_acc += cycles
            self.time_acc += elapsed

            # Besser als bisher?
            if fitness > self.best_fitness:
                # „experimental.json“ → „best.json“ (semantisch wie file:rename(...))
                self.best_fitness = fitness
                self.best_genotype = geno
                if self.best_file:
                    save_genotype(self.best_file, geno)
                # Reset Attempt-Zähler (wie Erlang: Attempt=0 nach Verbesserung)
                attempt = 1
            else:
                attempt += 1


# --------------------------
# Beispiel: ausführen
# --------------------------
if __name__ == "__main__":
    # Beispielkonfiguration (XOR-Morphologie, kleine Topologie)
    trainer = Trainer(
        morphology_spec=morphology,  # oder morphology.xor_mimic
        hidden_layer_densities=[2],  # wie im Buch-Beispiel
        max_attempts=1000,
        eval_limit=float("inf"),
        fitness_target=float("inf"),
        experimental_file="experimental.json",
        best_file="best.json",
        exoself_steps_per_eval=0,  # 0 => Cortex/Scape steuern Halt-Flag
    )

    asyncio.run(trainer.go())