import time
from collections.abc import Mapping
from pathlib import Path
from typing import Any, Callable

import joblib
import pandas as pd
from numpy.random import default_rng
from tqdm import tqdm

from phylogenie.io import dump_newick
from phylogenie.tree_node import TreeNode
from phylogenie.treesimulator.core import Model


def simulate_tree(
    model: Model,
    n_leaves: int | None = None,
    max_time: float | None = None,
    timeout: float | None = None,
    acceptance_criterion: Callable[[TreeNode], bool] | None = None,
    logs: dict[str, Callable[[TreeNode], Any]] | None = None,
) -> tuple[TreeNode, dict[str, Any]]:
    start_clock = time.perf_counter()
    while True:
        model.reset()
        while model.step(max_time) and (n_leaves is None or model.tree_size < n_leaves):
            if timeout is not None and time.perf_counter() - start_clock > timeout:
                raise TimeoutError("Simulation timed out.")

        # If the simulation stopped because the process stopped,
        # and not because we reached the desired number of leaves,
        # we restart the simulation.
        if n_leaves is not None and model.tree_size < n_leaves:
            continue

        tree = model.get_tree()

        if (
            tree is None
            or acceptance_criterion is not None
            and not acceptance_criterion(tree)
        ):
            continue

        metadata: dict[str, Any] = {}
        if logs is not None:
            for key, func in logs.items():
                metadata[key] = func(tree)

        return (tree, metadata)


def generate_trees(
    output_dir: str | Path,
    n_trees: int,
    model: Model,
    n_leaves: int | None = None,
    max_time: float | None = None,
    node_features: Mapping[str, str] | None = None,
    seed: int | None = None,
    n_jobs: int = -1,
    timeout: float | None = None,
    acceptance_criterion: Callable[[TreeNode], bool] | None = None,
    logs: dict[str, Callable[[TreeNode], Any]] | None = None,
) -> pd.DataFrame:
    if isinstance(output_dir, str):
        output_dir = Path(output_dir)
    if output_dir.exists():
        raise FileExistsError(f"Output directory {output_dir} already exists")
    output_dir.mkdir(parents=True)

    def _simulate_tree(i: int, seed: int) -> dict[str, Any]:
        while True:
            try:
                model.rng.seed(seed)
                tree, metadata = simulate_tree(
                    model=model,
                    n_leaves=n_leaves,
                    max_time=max_time,
                    timeout=timeout,
                    acceptance_criterion=acceptance_criterion,
                    logs=logs,
                )
                metadata["file_id"] = i
                if node_features is not None:
                    for name, feature in node_features.items():
                        mapping = getattr(tree, feature)
                        for node in tree:
                            node[name] = mapping[node]
                dump_newick(tree, output_dir / f"{i}.nwk")
                return metadata
            except TimeoutError as e:
                print(f"{e}. Retrying with a different seed...")
            seed += 1

    rng = default_rng(seed)
    jobs = joblib.Parallel(n_jobs=n_jobs, return_as="generator_unordered")(
        joblib.delayed(_simulate_tree)(i=i, seed=int(rng.integers(2**32)))
        for i in range(n_trees)
    )

    return pd.DataFrame(
        [md for md in tqdm(jobs, f"Generating trees in {output_dir}...", n_trees)]
    )
