import os
from functools import singledispatch

from scinexus import get_id_from_source
from scinexus.composable import ComposableApp
from scinexus.io_util import is_url, path_exists

from cogent3 import load_tree, make_tree
from cogent3.core.tree import PhyloNode
from cogent3.phylo.nj import gnj

from ._citations import cite_cogent3
from .data_store import get_data_source
from .typing import PairwiseDistanceType, SerialisableType, TreeType

NoneType = type(None)

get_data_source = get_id_from_source()


class scale_branches(ComposableApp, cite=cite_cogent3):
    """Transforms tree branch lengths from nucleotide to codon, or the
    converse. Returns a new tree with lengths divided by scalar"""

    def __init__(
        self,
        nuc_to_codon=None,
        codon_to_nuc=None,
        scalar=1,
        min_length=1e-6,
    ) -> None:
        """
        Parameters
        ----------
        nuc_to_codon : bool
            scalar is 3
        codon_to_nuc : bool
            scalar is 1/3
        scalar
            numerical value. Overridden by either nuc_to_codon or
            codon_to_nuc
        min_length
            set branch length to this value if it's not defined,
            or <= zero
        """
        assert not all([nuc_to_codon, codon_to_nuc])
        if nuc_to_codon:
            scalar = 3
        elif codon_to_nuc:
            scalar = 1 / 3
        else:
            assert scalar is not None

        self._scalar = scalar
        self._min_length = min_length

    def main(self, tree: TreeType) -> SerialisableType | TreeType:
        source = tree.source
        scalar = self._scalar
        min_length = self._min_length
        tree = tree.deepcopy()
        for edge in tree.postorder():
            if edge.name == "root":
                continue  # root has no length
            length = edge.length
            length = length / scalar if length else min_length
            edge.length = abs(length)  # force to be positive

        tree.source = source
        return tree


class uniformize_tree(ComposableApp, cite=cite_cogent3):
    """Standardises the orientation of unrooted trees."""

    def __init__(
        self, root_at: str = "midpoint", ordered_names: list[str] | None = None
    ) -> None:
        """
        Parameters
        ----------
        root_at
            edge to root at, defaults to midpoint
        ordered_names
            ordering of names, if not provided, taken from first
            tree
        """
        self._root_at = root_at
        self._ordered_names = ordered_names

    def main(self, tree: TreeType) -> SerialisableType | TreeType:
        source = tree.source
        if self._root_at == "midpoint":
            new = tree.root_at_midpoint()
        else:
            new = tree.rooted_with_tip(self._root_at)

        if self._ordered_names is None:
            self._ordered_names = tree.get_tip_names()

        result = new.sorted(self._ordered_names)
        result.source = source
        return result


class quick_tree(ComposableApp, cite=cite_cogent3):
    """Computes a Neighbour Joining tree from pairwise distances."""

    def __init__(self, drop_invalid: bool = False) -> None:
        """
        Parameters
        ----------
        drop_invalid
            drops all rows / columns with an invalid entry. If True, sequences
            for which a distance could not be calculated are excluded and
            the resulting tree will be for the subset of labels with strictly
            valid distances. If False, an ArithmeticError is raised if a
            distance is invalid.
        """
        self._drop_invalid = drop_invalid

    def main(self, dists: PairwiseDistanceType) -> SerialisableType | TreeType:
        """estimates a neighbor joining tree"""
        size = dists.shape[0]
        dists = dists.drop_invalid() if self._drop_invalid else dists
        if dists is None or (dists.shape[0] != size and not self._drop_invalid):
            msg = "invalid pairwise distances"
            raise ValueError(msg)

        if size == 2:
            # use code from dist matrix for this, no hook applied
            tree = dists.quick_tree()
        else:
            # we hard-code this call to allow for over-riding third-party
            # hooks
            (result,) = gnj(dists.to_dict(), keep=1, show_progress=False)
            _, tree = result

        tree.source = get_data_source(dists)
        return tree


@singledispatch
def interpret_tree_arg(tree) -> None | PhyloNode:
    msg = f"invalid tree type {type(tree)}"
    raise TypeError(msg)


@interpret_tree_arg.register
def _(tree: None) -> None:
    return tree


@interpret_tree_arg.register
def _(tree: os.PathLike) -> PhyloNode:
    return load_tree(filename=tree, underscore_unmunge=True)


@interpret_tree_arg.register
def _(tree: str) -> PhyloNode:
    if path_exists(tree) or is_url(tree):
        return load_tree(filename=tree, underscore_unmunge=True)
    return make_tree(tree, underscore_unmunge=True)


@interpret_tree_arg.register
def _(tree: PhyloNode) -> PhyloNode:
    return tree
