from typing import OrderedDict


def species_gene_mapping(gene_tree,
                         species_tree):
    """ Find the association between the leaves in species and gene trees.

    Parameters
    ----------
    gene_tree: skbio.TreeNode
        TreeNode instance for gene tree
    species_tree_fp: skbio.TreeNode
        TreeNode instance for species tree

    Returns
    -------
    mapping_leaves_t: OrderedDict
        Mapping between the species tree leaves and the gene tree leaves;
        species tips are the keys and gene tips are the values

    See Also
    --------
    skbio.TreeNode

    Notes
    -----
    Given the label format "SPECIES" for the species leaves and
    "SPECIES_GENE" in the gene leaves, report the associations between all
    species and gene leaves. Only one instance of the '_' delimiter is
    allowed in the gene leaves and this is used as a separator between the
    species name and the gene name.

    Ex.

    mapping = {"SE001":["SE001_1", "SE001_2"],
               "SE002":["SE002_1"]}
    """
    mapping_leaves = {}
    for node in species_tree.tips():
        if node.name not in mapping_leaves:
            mapping_leaves[node.name] = []
        else:
            raise ValueError(
                "Species tree leaves must be uniquely labeled: %s" % node.name)
    for node in gene_tree.tips():
        species, gene = node.name.split()
        if species in mapping_leaves:
            mapping_leaves[species].append("%s_%s" % (species, gene))
        else:
            raise ValueError(
                "Species %s does not exist in the species tree" % species)
    return OrderedDict(sorted(mapping_leaves.items(),
                       key=lambda x: x[1], reverse=True))