#! /usr/bin/env python3

from Bio import Phylo
from Bio.Phylo.TreeConstruction import DistanceCalculator
from Bio.Phylo.TreeConstruction import DistanceTreeConstructor
from Bio import AlignIO
from Bio.Phylo.TreeConstruction import DistanceMatrix
import numpy as np
import pickle
import os


def get_matrix_from_anchors():

    dir_path = os.path.dirname(os.path.realpath(__file__))

    orgs = []
    # Check both paths for orgs file
    if os.path.exists('orgs'):
        orgs_path = 'orgs'
    elif os.path.exists(dir_path + '/../utils/orgs'):
        orgs_path = dir_path + '/../utils/orgs'
    else:
        raise FileNotFoundError("orgs file not found")

    with open(orgs_path) as f:
        for line in f:
            orgs.append(line.strip())

    ams = {}
    for org in orgs:
        # Check both paths for compressed_maps_multis_to_one
        if os.path.exists('compressed_maps_multis_to_one/{}'.format(org)):
            path = 'compressed_maps_multis_to_one/{}'.format(org)
        elif os.path.exists(dir_path + '/../utils/compressed_maps_multis_to_one/{}'.format(org)):
            path = dir_path + '/../utils/compressed_maps_multis_to_one/{}'.format(org)
        else:
            raise FileNotFoundError(f"compressed_maps_multis_to_one/{org} not found")
        with open(path, 'rb') as f:
            ams[org] = pickle.load(f)

    lens = {}
    for org in orgs:
        # Check both paths to work from downstream/ and from out/scripts_which_produced_these_results/
        if os.path.exists(dir_path + '/../utils/small_meta/{}'.format(org)):
            path = dir_path + '/../utils/small_meta/{}'.format(org)
        elif os.path.exists(dir_path + '/../../utils/small_meta/{}'.format(org)):
            path = dir_path + '/../../utils/small_meta/{}'.format(org)
        else:
            raise FileNotFoundError(f"small_meta file not found for {org}")
        with open(path, 'rb') as f:
            x = pickle.load(f)
            lens[org] = x[1][-1]

    matrix = [[]]
    for enu, org in enumerate(orgs):
        if enu == 0:
            continue
        loc_dists = [0 for i in range(enu)]
        for i in range(0, enu):
            org2 = orgs[i]
            if lens[org] < lens[org2]:
                dist_org = org
                other = org2
            else:
                dist_org = org2
                other = org
            for j, bib in ams[dist_org].items():
                if other not in bib['matches']:
                    continue
                else:
                    for k, t in bib['matches'][other].items():
                        loc_dists[i] += (t[1][1] - t[1][0])
            for j, bib in ams[other].items():
                if dist_org not in bib['matches']:
                    continue
                else:
                    for k, t in bib['matches'][dist_org].items():
                        loc_dists[i] += (t[1][1] - t[1][0])

        dists = []
        for i in range(0, enu):
            d = loc_dists[i]
            org2 = orgs[i]
            if lens[org] < lens[org2]:
                dist_org = org
                other = org2
            else:
                dist_org = org2
                other = org
            dists.append(1 - (d / (lens[dist_org]+lens[other])))
        matrix.append(dists)
    return matrix, orgs

if __name__ == "__main__":

    matrix,orgs = get_matrix_from_anchors()

    for i in range(len(matrix)):
        matrix[i].append(0)

    distMatrix = DistanceMatrix(orgs, matrix)

    constructor = DistanceTreeConstructor()

    UPGMATree = constructor.upgma(distMatrix)

    #Phylo.draw(UPGMATree)

    #Phylo.draw_ascii(UPGMATree)

    NJTree = constructor.nj(distMatrix)

    #Phylo.draw(NJTree)

    #Phylo.draw_ascii(NJTree)

    Phylo.write(UPGMATree, "UPGMATree.nwk", "newick")
    Phylo.write(NJTree, "NJTree.nwk", "newick")
