"""
Reconstruct alignments from mutations inferred on the tree
"""

import os, sys
import numpy as np
from collections import defaultdict
from Bio import SeqIO, Seq, SeqRecord, Phylo
from .utils import read_node_data, write_json
from treetime.vcf_utils import read_vcf



def register_arguments(parser):
    parser.add_argument('--tree', required=True, help="tree as Newick file")
    parser.add_argument('--gene', type=str, help="gene to translate (list or file containing list)")
    parser.add_argument('--mutations', required=True, type=str, help="json file containing mutations "
                            "mapped to each branch and the sequence of the root.")
    parser.add_argument('--vcf-aa-reference', type=str, help='fasta file of the reference gene translations for VCF format')
    parser.add_argument('--internal-nodes', action='store_true', help="include sequences of internal nodes in output")
    parser.add_argument('--output', type=str)


def get_sequence(pseq, muts):
    """reconstruct a child sequence from that of its parent and the mutations

    Parameters
    ----------
    pseq : str
        sequence of the parent as string, won't be changed!
    muts : list
        mutations of the type K135T where K is the parental state, 135 is the
        position in 1-based numbering, and T is the new state

    Returns
    -------
    str
        reconstructed sequence
    """
    pseq_list = list(pseq)
    for mut in muts:
        new_state = mut[-1]
        pos = int(mut[1:-1])-1
        assert pseq_list[pos]==mut[0]
        pseq_list[pos]=new_state

    return "".join(pseq_list)


def load_alignments(sequence_files, gene_names):
    from Bio import AlignIO
    alignments = {}
    for fname, gene in zip(sequence_files, gene_names):
        alignments[gene] = AlignIO.read(fname, 'fasta')
    return alignments


def run(args):
    ## read tree and data, if reading data fails, return with error code
    tree = Phylo.read(args.tree, 'newick')

    #check whether VCF - by existance of VCF reference file
    is_vcf = args.vcf_aa_reference is not None

    ## check file format and read in sequences
    node_data = read_node_data(args.mutations, args.tree)
    if node_data is None:
        print("ERROR: could not read mutation data "+ ("(incl sequences)" if not is_vcf else ""))
        return 1

    root_node = tree.root.name

    #if VCF, read in the reference seq for each gene, put on root
    if(is_vcf):
        node_data["nodes"][root_node]['aa_sequences'] = {}
        with open(args.vcf_aa_reference, encoding='utf-8') as handle:
            for record in SeqIO.parse(handle, "fasta"):
                if record.id==args.gene:
                    #'root' may not be same as 'reference', so apply any mutations at root here!
                    node_data["nodes"][root_node]['aa_sequences'][record.id] = get_sequence(str(record.seq), node_data["nodes"][root_node]["aa_muts"][record.id])

    # check that root node has sequences for each requested gene
    if "aa_sequences" not in node_data["nodes"][root_node]:
        print("ERROR: ancestral sequences are not provided")
        return 1
    if not args.gene in node_data["nodes"][root_node]['aa_sequences']:
        print("ERROR: ancestral sequences missing for gene",args.gene)
        return 1

    # gather all reconstructed sequences
    sequences = dict()
    is_terminal = {}

    sequences[root_node] = node_data["nodes"][root_node]['aa_sequences'][args.gene]
    is_terminal[root_node] = False
    for node in tree.get_nonterminals(order='preorder'):
        parent_sequence = sequences[node.name]
        for child in node:
            sequences[child.name] = get_sequence(parent_sequence, node_data["nodes"][child.name]["aa_muts"][args.gene])
            is_terminal[child.name] = child.is_terminal()

    # write alignments to file
    seqs = [SeqRecord.SeqRecord(seq=Seq.Seq(sequences[strain]), id=strain, name=strain, description='')
            for strain in sequences if is_terminal[strain] or args.internal_nodes]
    SeqIO.write(seqs, args.output, 'fasta')
