import os, sys
from Bio import AlignIO, Phylo
from .vcf_utils import read_vcf, write_vcf
from .seq_utils import alphabets
from Bio import __version__ as bioversion
from . import version as treetime_version
import numpy as np

def get_outdir(params, suffix='_treetime'):
    if params.outdir:
        if os.path.exists(params.outdir):
            if os.path.isdir(params.outdir):
                return params.outdir.rstrip('/') + '/'
            else:
                print("designated output location %s is not a directory"%params.outdir, file=sys.stderr)
        else:
            os.makedirs(params.outdir)
            return params.outdir.rstrip('/') + '/'

    from datetime import datetime
    outdir_stem = datetime.now().date().isoformat()
    outdir = outdir_stem + suffix.rstrip('/')+'/'
    count = 1
    while os.path.exists(outdir):
        outdir = outdir_stem + '-%04d'%count + suffix.rstrip('/')+'/'
        count += 1

    os.makedirs(outdir)
    return outdir

def get_basename(params, outdir):
    # if params.aln:
    #     basename = outdir + '.'.join(params.aln.split('/')[-1].split('.')[:-1])
    # elif params.tree:
    #     basename = outdir + '.'.join(params.tree.split('/')[-1].split('.')[:-1])
    # else:
    basename = outdir
    return basename

def read_in_DRMs(drm_file, offset):
    import pandas as pd

    DRMs = {}
    drmPositions = []

    df = pd.read_csv(drm_file, sep='\t')
    for mi, m in df.iterrows():
        pos = m.GENOMIC_POSITION-1+offset #put in correct numbering
        drmPositions.append(pos)

        if pos in DRMs:
            DRMs[pos]['alt_base'][m.ALT_BASE] = m.SUBSTITUTION
        else:
            DRMs[pos] = {}
            DRMs[pos]['drug'] = m.DRUG
            DRMs[pos]['alt_base'] = {}
            DRMs[pos]['alt_base'][m.ALT_BASE] = m.SUBSTITUTION
            DRMs[pos]['gene'] = m.GENE

    drmPositions = np.array(drmPositions)
    drmPositions = np.unique(drmPositions)
    drmPositions = np.sort(drmPositions)

    DRM_info = {'DRMs': DRMs,
            'drmPositions': drmPositions}

    return DRM_info


def read_if_vcf(params):
    """
    Checks if input is VCF and reads in appropriately if it is
    """
    ref = None
    aln = params.aln
    fixed_pi = None
    if hasattr(params, 'aln') and params.aln is not None:
        if any([params.aln.lower().endswith(x) for x in ['.vcf', '.vcf.gz']]):
            if not params.vcf_reference:
                print("ERROR: a reference Fasta is required with VCF-format alignments")
                return -1
            compress_seq = read_vcf(params.aln, params.vcf_reference)
            sequences = compress_seq['sequences']
            ref = compress_seq['reference']
            aln = sequences

            if not hasattr(params, 'gtr') or params.gtr=="infer": #if not specified, set it:
                alpha = alphabets['aa'] if params.aa else alphabets['nuc']
                fixed_pi = [ref.count(base)/len(ref) for base in alpha]
                if fixed_pi[-1] == 0:
                    fixed_pi[-1] = 0.05
                    fixed_pi = [v-0.01 for v in fixed_pi]

    return aln, ref, fixed_pi


def plot_rtt(tt, fname):
    tt.plot_root_to_tip()

    from matplotlib import pyplot as plt
    plt.savefig(fname)
    print("--- root-to-tip plot saved to  \n\t"+fname)


def export_sequences_and_tree(tt, basename, is_vcf=False, zero_based=False,
                              report_ambiguous=False, timetree=False, confidence=False,
                              reconstruct_tip_states=False, tree_suffix=''):
    seq_info = is_vcf or tt.aln
    if is_vcf:
        outaln_name = basename + f'ancestral_sequences{tree_suffix}.vcf'
        write_vcf(tt.get_reconstructed_alignment(reconstruct_tip_states=reconstruct_tip_states), outaln_name)
    elif tt.aln:
        outaln_name = basename + f'ancestral_sequences{tree_suffix}.fasta'
        AlignIO.write(tt.get_reconstructed_alignment(reconstruct_tip_states=reconstruct_tip_states), outaln_name, 'fasta')
    if seq_info:
        print("\n--- alignment including ancestral nodes saved as  \n\t %s\n"%outaln_name)

    # decorate tree with inferred mutations
    terminal_count = 0
    offset = 0 if zero_based else 1
    if timetree:
        dates_fname = basename + f'dates{tree_suffix}.tsv'
        fh_dates = open(dates_fname, 'w', encoding='utf-8')
        if confidence:
            fh_dates.write('#Lower and upper bound delineate the 90% max posterior region\n')
            fh_dates.write('#node\tdate\tnumeric date\tlower bound\tupper bound\n')
        else:
            fh_dates.write('#node\tdate\tnumeric date\n')

    mutations_out = open(basename + "branch_mutations.txt", "w")
    mutations_out.write("node\tstate1\tpos\tstate2\n")
    for n in tt.tree.find_clades():
        if timetree:
            if confidence:
                if n.bad_branch:
                    fh_dates.write('%s\t--\t--\t--\t--\n'%(n.name))
                else:
                    conf = tt.get_max_posterior_region(n, fraction=0.9)
                    fh_dates.write('%s\t%s\t%f\t%f\t%f\n'%(n.name, n.date, n.numdate,conf[0], conf[1]))
            else:
                if n.bad_branch:
                    fh_dates.write('%s\t--\t--\n'%(n.name))
                else:
                    fh_dates.write('%s\t%s\t%f\n'%(n.name, n.date, n.numdate))

        n.confidence=None
        # due to a bug in older versions of biopython that truncated filenames in nexus export
        # we truncate them by hand and make them unique.
        if n.is_terminal() and len(n.name)>40 and bioversion<"1.69":
            n.name = n.name[:35]+'_%03d'%terminal_count
            terminal_count+=1
        n.comment=''
        if seq_info and len(n.mutations):
            if n.mask is None:
                if report_ambiguous:
                    n.comment= '&mutations="' + ','.join([a+str(pos + offset)+d for (a,pos, d) in n.mutations])+'"'
                else:
                    n.comment= '&mutations="' + ','.join([a+str(pos + offset)+d for (a,pos, d) in n.mutations
                                                        if tt.gtr.ambiguous not in [a,d]])+'"'
            else:
                if report_ambiguous:
                    n.comment= '&mutations="' + ','.join([a+str(pos + offset)+d for (a,pos, d) in n.mutations if n.mask[pos]>0])+f'",mcc="{n.mcc}"'
                else:
                    n.comment= '&mutations="' + ','.join([a+str(pos + offset)+d for (a,pos, d) in n.mutations
                                                        if tt.gtr.ambiguous not in [a,d] and n.mask[pos]>0])+f'",mcc="{n.mcc}"'

            for (a, pos, d) in n.mutations:
                if tt.gtr.ambiguous not in [a,d] or report_ambiguous:
                    mutations_out.write("%s\t%s\t%s\t%s\n" %(n.name, a, pos + 1, d))
        if timetree:
            n.comment+=(',' if n.comment else '&') + 'date=%1.2f'%n.numdate
    mutations_out.close()

    # write tree to file
    fmt_bl = "%1.7f" if tt.data.full_length<1e6 else "%1.9e"
    if timetree:
        outtree_name = basename + f'timetree{tree_suffix}.nexus'
        print("--- saved divergence times in \n\t %s\n"%dates_fname)
        Phylo.write(tt.tree, outtree_name, 'nexus')
    else:
        outtree_name = basename + f'annotated_tree{tree_suffix}.nexus'
        Phylo.write(tt.tree, outtree_name, 'nexus', format_branch_length=fmt_bl)
    print("--- tree saved in nexus format as  \n\t %s\n"%outtree_name)

    # Only create auspice json if there is sequence information
    auspice = create_auspice_json(tt, timetree=timetree, confidence=confidence, seq_info=seq_info)
    outtree_name_json = basename + f'auspice_tree{tree_suffix}.json'
    with open(outtree_name_json, 'w') as fh:
        import json
        json.dump(auspice, fh, indent=0)
        print("--- tree saved in auspice json format as  \n\t %s\n"%outtree_name_json)

    if timetree:
        for n in tt.tree.find_clades():
            n.branch_length = n.mutation_length
        outtree_name = basename + f'divergence_tree{tree_suffix}.nexus'
        Phylo.write(tt.tree, outtree_name, 'nexus', format_branch_length=fmt_bl)
        print("--- divergence tree saved in nexus format as  \n\t %s\n"%outtree_name)

    if hasattr(tt, 'outliers') and tt.outliers is not None:
        print("--- saved detected outliers as " + basename + 'outliers.tsv')
        tt.outliers.to_csv(basename + 'outliers.tsv', sep='\t')

def print_save_plot_skyline(tt, n_std=2.0, screen=True, save='', plot='', gen=50):
    if plot:
        import matplotlib.pyplot as plt

    skyline, conf = tt.merger_model.skyline_inferred(gen=gen, confidence=n_std)
    if save: fh = open(save, 'w', encoding='utf-8')
    header1 = "Skyline assuming "+ str(gen)+" gen/year and approximate confidence bounds (+/- %f standard deviations of the LH)\n"%n_std
    header2 = "date \tN_e \tlower \tupper"
    if screen: print('\t'+header1+'\t'+header2)
    if save: fh.write("#"+ header1+'#'+header2+'\n')
    for (x,y, y1, y2) in zip(skyline.x, skyline.y, conf[0], conf[1]):
        if screen: print("\t%1.3f\t%1.3e\t%1.3e\t%1.3e"%(x,y, y1, y2))
        if save: fh.write("%1.3f\t%1.3e\t%1.3e\t%1.3e\n"%(x,y, y1, y2))

    if save:
        print("\n --- written skyline to %s\n"%save)
        fh.close()

    if plot:
        plt.figure()
        plt.fill_between(skyline.x, conf[0], conf[1], color=(0.8, 0.8, 0.8))
        plt.plot(skyline.x, skyline.y, label='maximum likelihood skyline')
        plt.yscale('log')
        plt.legend()
        plt.ticklabel_format(axis='x',useOffset=False)
        plt.savefig(plot)



def create_auspice_json(tt, timetree=False, confidence=False, seq_info=False):
    # mock up meta data for auspice json
    from datetime import datetime
    meta = {
        "title": f"Auspice visualization of TreeTime (v{treetime_version}) analysis",
        "build_url": "https://github.com/neherlab/treetime",
        "last_updated": datetime.now().strftime("%Y-%m-%d"),
        "treetime_version": treetime_version,
        "genome_annotations": {
            "nuc":{"start":1, "end":int(tt.data.full_length), "type":"source", "strand":"+:"}
        },
        "panels":["tree", "entropy"],
        "colorings": [
            {
                "title": "Date",
                "type": "continuous",
                "key": "num_date",
            },
            {
                "title": "Genotype",
                "type": "categorical",
                "key": "gt",
            },
            {
                "title": "Excluded",
                "type": "categorical",
                "key": "bad_branch"
            },
            {
                "title": "Branch Support",
                "type": "continuous",
                "key": "confidence"
            }
        ],
        "display_defaults": {"color_by":"bad_branch"},
        "filters": ["bad_branch"]
    }

    def node_to_json(n, pdiv=0.0):
        j = {"name":n.name, "node_attrs":{}, "branch_attrs":{}}
        if n.clades:
            j["children"] = []

        if timetree:
            j["node_attrs"]["num_date"] = {"value":float(n.numdate)}
            if confidence:
                conf = tt.get_max_posterior_region(n, fraction=0.9)
                j["node_attrs"]["num_date"]["confidence"] = (float(conf[0]), float(conf[1]))
        j["node_attrs"]["div"] = float(pdiv + n.mutation_length)
        j["node_attrs"]["bad_branch"] = {"value": "Yes" if n.bad_branch else "No"}

        if seq_info: # only add mutations to the json if run with sequence data (fasta or vcf)
            j["branch_attrs"]["mutations"] = {"nuc": [f"{a}{pos+1}{d}" for a,pos,d in n.mutations if d in "ACGT-"]}
            # generate bootstrap confidence substitute via the negative exponential of the number of mutations
            # this is the bootstrap confidence for iid mutations (only ACGT mutations)
            j["node_attrs"]["confidence"] = {"value":round(1-np.exp(-len([pos for a,pos,d in n.mutations if d in "ACGT"])),3)
                                            if not n.is_terminal() else 1.0}
        return j

    # create the tree data structure from the Biopython tree
    tree = node_to_json(tt.tree.root, 0.0)
    # dictionary to look up nodes by name
    node_lookup = {tt.tree.root.name: tree}
    for n in tt.tree.get_nonterminals():
        n_json = node_lookup[n.name]
        for c in n.clades:
            # generate node jsons for all children and attach them the to parent
            n_json["children"].append(node_to_json(c, n_json["node_attrs"]["div"]))
            node_lookup[c.name] = n_json["children"][-1]

    return {"meta":meta, "tree":tree}