import os
from Bio import Phylo
from typing import Optional
from matplotlib.pyplot import Axes
from .fasta import FastaParser
from .lowlevel import call


def muscle(
        fasta: str,
        output: str,
        maxiters: int = 16,
        outformat: str = 'fasta',
        log: Optional[str] = None):
    """
    Wrapper function for MSUCLE sequence alignment

    Args:
        fasta: path-like
            The input fasta file

        output: path-like
            The output file

        maxiters:
            Maximum number of iterations (default 16)

        outformat:
            'fasta' (default)
            'html'
            'msf' (GCF MSF)
            'clw' (ClustalW)

        log: path-like
            The log file, default None -> <output>.log
    """
    outformat = {'fasta': '', 'html': '-html', 'msf': '-msf', 'clw': '-clw'}[outformat]

    if log is None: log = output + '.log'
    call(f'muscle -in {fasta} -out {output} -maxiters {maxiters} -log {log} {outformat}')


def simple_phylogeny(
        fasta: str,
        output: str,
        clustering: str = 'NJ',
        log: Optional[str] = None):
    """
    Make a tree file (newick/phylip format) from the aligned <fasta>
        using the command "clustalw -tree"

    Args:
        fasta: path-like
            The input fasta file, sequences must be aligned

        output: path-like
            The output tree file

        clustering:
            clustering method, 'NJ' or 'UPGMA'

        log: path-like
            The log file, default None -> <output>.log
    """
    with FastaParser(fasta) as parser:
        head, seq = parser.next()
        if set(seq.upper()) == set('ACGT-'): type_ = 'DNA'
        else: type_ = 'PROTEIN'

    if log is None: log = output + '.log'
    call(f'clustalw -tree -infile={fasta} -type={type_} -CLUSTERING={clustering} > {log}')

    # Rename the .ph file automatically generated by clustalw
    ph_file = fasta[:fasta.rfind('.')] + '.ph'
    os.rename(ph_file, output)


def draw_tree(file: str, ax: Axes):
    """
    Wrapper function for Bio.Phylo.read() and Bio.Phylo.draw()
        to draw a phylogenetic tree from the input tree <file>

    Args:
        file: path-like
            The input tree file in newick/phylip format

        ax: matplotlib Axes object
            The ax object in which the tree will be drawn
    """
    tree = Phylo.read(file, 'newick')
    Phylo.draw(tree, axes=ax)
