import numpy as np
import math
class DistanceMatrix:


    def __init__(self, *args):
        self.arr = np.array(*args)

    def __repr__(self):
        return 'DistanceMatrix(' + str(self.arr.tolist()) + ')'

    def __str__(self):
        return str(self.arr.tolist())

    @staticmethod
    def loadtxt(args):
        return DistanceMatrix(np.loadtxt(args))

    def savetxt(self, *args, **kwargs):
        np.savetxt(*args, self.arr, **kwargs)

    def neighbour_joining(self):
        nodes = [x for x in range(len(self.arr))]
        mapper = {x:x for x in range(len(self.arr))}
        return self.neighbour_joiningcalc(self.arr.copy(), len(self.arr), nodes, len(nodes), mapper)

    def neighbour_joiningcalc(self, d, n, nodes, counter, mapper):
        if n == 2:
            return UnrootedTree((0, 1, d[0][1]))

        totalDistance= [sum(row) for row in d]
        dster = [[(n-2)*d[i][j] - totalDistance[i] - totalDistance[j] for j in range(n)] for i in range(n)]

        min = math.inf
        row = 0
        col = 0
        nodesRow = 0
        nodesCol = 0
        for i in range(len(nodes)):
            first = nodes[i]
            firstVal = mapper[first]
            for j in range(len(nodes)):
                second = nodes[j]
                secondVal = mapper[second]
                if i != j and dster[firstVal][secondVal] < min:
                    min = dster[firstVal][secondVal]
                    row = firstVal
                    col = secondVal
                    nodesRow = first
                    nodesCol = second

        delta = (totalDistance[row] - totalDistance[col])/(n-2)
        limbLengthi = 1/2*(d[row][col] + delta)
        limbLengthj = 1/2*(d[row][col] - delta)

        d = np.hstack((d, [[1/2*(d[k][row] + d[k][col] - d[row][col])] for k in range(n)]))
        d = np.vstack((d, [1/2*(d[k][row] + d[k][col] - d[row][col]) if k != n else 0 for k in range(n+1)]))

        nodes.append(counter)
        mapper[counter] = counter

        d = np.delete(d, row, 0)
        d = np.delete(d, row, 1)
        d = np.delete(d, col, 0)
        d = np.delete(d, col, 1)

        for node in mapper:
            if node > row:
                mapper[node] -= 1
            if node > col:
                mapper[node] -= 1


        nodes.remove(nodesRow)
        nodes.remove(nodesCol)

        t = self.neighbour_joiningcalc(d, n-1, nodes, counter + 1, mapper)

        t.add((n, row, limbLengthi))
        t.add((n, col, limbLengthj))
        return t

class UnrootedTree:
    def __init__(self, *args):
        self.edges = args

    def __str__(self):
        return str(self.edges)

    def __repr__(self):
        return 'UnrootedTree' + str(self.edges)


    @staticmethod
    def loadtxt(file):
        openFile = open(file, 'r')
        opl = list()
        for line in openFile:
            line = line.rstrip()
            f = line.split('<->')
            f[1] = f[1].split(':')
            opl.append((float(f[0]), float(f[1][0]), float(f[1][1])))
        return UnrootedTree(*tuple(opl))

    def add(self, lim):
        self.edges += (lim,)