from __future__ import annotations

from collections.abc import Hashable, Iterator, Sequence
from dataclasses import dataclass
from datetime import timedelta
from typing import Optional, Sized, SupportsFloat

from ..b3 import Callback, Operator, Prior
from ..b3.likelihoods import Likelihood
from ..b3.parameters import Node, Parameter, Scalable
from ..b3.substitutions import Substiution4
from ..data.msa import MSA
from ..rng import RNG
from ..stats.distributions import Sample

class Tree:
    def __init__(self, names: list[str], rng: RNG): ...
    def set_random_topology(self, rng: RNG): ...
    def set_random_edges(self, rng: RNG): ...
    def set_random_heights(self, diff, rng: RNG): ...
    @property
    def names(self) -> list[str]: ...
    def scale(self, scale: float): ...
    def replace_child(
        self, parent: Internal, old_child: Node, new_child: Node
    ) -> None: ...
    def spr(self, node: Node, other: Node) -> None: ...
    def set_height(self, node: Node, weigth: float) -> None: ...
    def set_root(self, node: Node) -> None: ...
    def swap_parents(self, a: Node, b: Node) -> None: ...
    @property
    def num_nodes(self) -> int: ...
    @property
    def num_internals(self) -> int: ...
    @property
    def num_leaves(self) -> int: ...
    @property
    def num_edges(self) -> int: ...
    def is_internal(self, node: Node) -> bool: ...
    def is_leaf(self, node: Node) -> bool: ...
    @property
    def root(self) -> Internal: ...
    def height_of(self, node: Node) -> float: ...
    def children_of(self, node: Internal) -> tuple[Node, Node]: ...
    def other_child(self, parent: Internal, child: Node) -> Node: ...
    def edge_index(self, child: Node) -> int: ...
    def edge_length(self, edge: int) -> float: ...
    def edge_nodes(self, edge: int) -> tuple[Node, Internal]: ...
    def parent_of(self, node: Node) -> Optional[Internal]: ...
    def is_grandparent(self, node: Internal) -> bool: ...
    def num_grandparents(self) -> int: ...
    def nodes(self) -> Iterator[Node]: ...
    def internals(self) -> Iterator[Internal]: ...
    def leaves(self) -> Iterator[Leaf]: ...
    def random_node(self, rng: RNG) -> Node: ...
    def random_nonroot_node(self, rng: RNG) -> tuple[Internal, Internal]: ...
    def random_internal(self, rng: RNG) -> Internal: ...
    def random_nonroot_internal(self, rng: RNG) -> tuple[Internal, Internal]: ...
    def random_leaf(self, rng: RNG) -> Leaf: ...
    def leaf_by_name(self, name: str) -> Optional[Leaf]: ...
    def total_length(self) -> float: ...
    def has_dated_tips(self) -> bool: ...
    def validate(self) -> None: ...
    def to_newick(self, internal_ids=False) -> str: ...
    def accept(self) -> None: ...
    def reject(self) -> None: ...
    def load(self, bytes: bytes) -> None: ...

class Leaf(Hashable): ...
class Internal(Hashable): ...

class Proposal:
    @classmethod
    def Reject(cls) -> Proposal: ...
    @classmethod
    def Hastings(cls, ratio: float) -> Proposal: ...
    @classmethod
    def Accept(cls) -> Proposal: ...

class CPU4Likelihood(Likelihood):
    def __init__(
        self,
        msa: MSA,
        substitution: Substiution4,
        clock: Clock,
        tree: Tree,
        scale_ln: int = 30,
    ): ...

class CUDALikelihood(Likelihood):
    def __init__(
        self,
        msa: MSA,
        substitution: Substiution4,
        clock: Clock,
        tree: Tree,
        *,
        cuda_device: int = 0,
    ): ...

class HeteroLikelihood(Likelihood):
    def __init__(self, likelihoods: list[Likelihood]): ...
    @property
    def class_vector(self) -> ClassVector: ...

class MCMC:
    def __init__(
        self,
        state: Sequence[Parameter],
        priors: Sequence[Prior],
        operators: Sequence[Operator],
        likelihood: Likelihood,
        callbacks: Sequence[Callback],
        rng: RNG,
    ): ...
    @property
    def current_step(self) -> int: ...
    @property
    def parameters(self) -> list[Parameter]: ...
    @property
    def priors(self) -> list[Prior]: ...
    @property
    def operators(self) -> list[Operator]: ...
    @property
    def likelihood(self) -> Likelihood: ...
    @property
    def callbacks(self) -> list[Callback]: ...
    @property
    def rng(self) -> RNG: ...
    @property
    def posterior(self) -> float: ...
    @property
    def likelihood_value(self) -> float: ...
    @property
    def prior(self) -> float: ...
    @property
    def operator_statistics(
        self,
    ) -> list[tuple[Operator, list[int], timedelta, timedelta]]: ...
    def run(self, n: int) -> None: ...
    def dump_state(self) -> bytes: ...
    def load_state(self, state: bytes) -> None: ...

@dataclass
class EpochScale(Operator):
    tree: Tree
    factor: float
    distribution: Sample[float]
    rng: RNG
    weight: float = 1

@dataclass(slots=True)
class SubtreeLeap(Operator):
    tree: Tree
    distribution: Sample[float]
    rng: RNG
    weight: float = 1

@dataclass(slots=True)
class FixedHeightSPR(Operator):
    tree: Tree
    rng: RNG
    weight: float = 1

class ClassvecFlip(Operator):
    def __init__(self, classvec: ClassVector, rng: RNG, weight: float): ...

@dataclass
class SymmetricDirichlet(Prior):
    vec: RealVector
    sum: float

@dataclass
class ConstantPopulation(Prior):
    tree: Tree
    population: Real

@dataclass
class ExponentialGrowth(Prior):
    tree: Tree
    population: Real
    growth_rate: Real

@dataclass
class Monophyly(Prior):
    tree: Tree
    leaves: Sequence[Leaf]

@dataclass
class Yule(Prior):
    tree: Tree
    birth_rate: Real

class JC: ...

@dataclass
class K80:
    kappa: Real

@dataclass
class HKY:
    frequencies: RealVector
    kappa: Real

@dataclass
class GTR:
    frequencies: RealVector
    rates: RealVector

class ClassVector(Sized):
    def into_list(self) -> list[int]: ...

class Real(SupportsFloat, Scalable):
    def __init__(self, value: SupportsFloat): ...
    def set(self, value: SupportsFloat) -> None: ...
    def __lt__(self, other) -> bool: ...
    def __le__(self, other) -> bool: ...
    def __gt__(self, other) -> bool: ...
    def __ge__(self, other) -> bool: ...
    def __eq__(self, other) -> bool: ...
    def __ne__(self, other) -> bool: ...

class RealVector(Sized):
    def __init__(self, *values: SupportsFloat): ...
    def __getitem__(self, index: int) -> float: ...
    def __setitem__(self, index: int, value: float): ...
    def __lt__(self, other) -> bool: ...
    def __le__(self, other) -> bool: ...
    def __gt__(self, other) -> bool: ...
    def __ge__(self, other) -> bool: ...
    def __eq__(self, other) -> bool: ...
    def __ne__(self, other) -> bool: ...

class Clock:
    @classmethod
    def Strict(_cls, rate: Real) -> Clock: ...

class TraceWriter(Callback):
    def __init__(
        self,
        items: dict[str, Parameter | Prior],
        path: str,
        *,
        overwrite: bool = False,
        zstd: bool = False,
        every: int,
    ): ...
