#!/usr/bin/env python

import sys
import csv
import numpy as np

import Bio.Phylo
from Bio.Phylo.TreeConstruction import DistanceMatrix, DistanceTreeConstructor

def usage():
  print("usage: mentalist_tree <input.tsv>\n")

def process_input_matrix(input_matrix):
    """ Converts an array-of-arrays containting sample IDs and distances
        into a BioPython DistanceMatrix object
    """
    input_matrix.pop(0)
    sample_names = [row[0] for row in input_matrix]
    for row in input_matrix:
        row.pop(0)
    distance_matrix = []
    for input_matrix_row in input_matrix:
        distance_matrix.append([int(i) for i in input_matrix_row])
    distance_matrix = np.tril(np.array(distance_matrix))
    num_rows = distance_matrix.shape[0]
    lower_triangular_idx_mask = np.tril_indices(num_rows)
    linear_distance_matrix = distance_matrix[lower_triangular_idx_mask]
    distance_matrix = []
    min = 0
    max = 1
    for i in range(num_rows):
        distance_matrix.append(linear_distance_matrix[min:max].tolist())
        min = max
        max = max + (i + 2) 
    distance_matrix = DistanceMatrix(names=sample_names, matrix=distance_matrix)
    return distance_matrix

def main():
    if len(sys.argv) < 2:
        usage()
        sys.exit(1)

    input_file = sys.argv[1]
    reader = csv.reader(open(input_file, "r"), delimiter="\t")
    input_matrix = list(reader)
    distance_matrix = process_input_matrix(input_matrix)
    constructor = DistanceTreeConstructor()
    tree = constructor.nj(distance_matrix)
    Bio.Phylo.write(tree, sys.stdout, 'newick')

if __name__ == '__main__':
    main()