import asyncio
import os
import time
from typing import Any, Dict, List, Tuple, Optional

import morphology
from genotype import construct, save_genotype, print_genotype
from exoself import Exoself


class Trainer:
    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
        self.exoself_steps_per_eval = exoself_steps_per_eval

        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]:
        print("constructing genotype...")
        geno = construct(
            self.morphology_spec,
            self.hds,
            file_name=self.experimental_file,  # <-- schreibt Startnetz nach experimental.json
            add_bias=True
        )
        fitness, evals, cycles, elapsed = await self._evaluate_with_exoself(geno)
        return fitness, evals, cycles, elapsed

    async def _evaluate_with_exoself(self, genotype: Dict[str, Any]) -> Tuple[float, int, int, float]:
        print("creating exoself...")
        ex = Exoself(genotype, file_name=self.experimental_file)
        best_fitness, evals, cycles, elapsed = await ex.train_until_stop()
        return best_fitness, evals, cycles, elapsed

    async def go(self):
        attempt = 1
        while True:
            print(".........")
            print("current attempt: ", attempt)
            print(".........")

            if attempt > self.max_attempts or self.eval_acc >= self.eval_limit or self.best_fitness >= self.fitness_target:
                # Abschlussausgabe wie im Buch
                if self.best_file and os.path.exists(self.best_file):
                    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}"
                )
                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!")
            fitness, evals, cycles, elapsed = await self._run_one_attempt()

            print("update akkus...")

            self.eval_acc += evals
            self.cycle_acc += cycles
            self.time_acc += elapsed

            # Besser als bisher?
            if fitness > self.best_fitness:
                self.best_fitness = fitness
                if self.best_file and self.experimental_file and os.path.exists(self.experimental_file):
                    os.replace(self.experimental_file, self.best_file)
                attempt = 1
            else:
                attempt += 1



if __name__ == "__main__":
    trainer = Trainer(
        morphology_spec=morphology,
        hidden_layer_densities=[2],
        max_attempts=200,
        eval_limit=float("inf"),
        fitness_target=99.9,
        experimental_file="experimental.json",
        best_file="best.json",
        exoself_steps_per_eval=0,
    )

    asyncio.run(trainer.go())