#! /usr/bin/env python
# __author__ = "Parisa Kianmajd"
#__version__ = "2.0.0"
# input: a 3 column csv file <Name, Year, Rank> generated by addRank and a MIR file
# output: A csv file showing some statistics about MIR file

import csv
import argparse
import sys


def main(inFile, rankFile):
    art2symbol = {'equals':'==',
                      'is_included_in':'<',
                      'includes':'>',
                      'overlaps':'><',
                  'disjoint':'!'}
    elements = dict()
    table = dict()
    ranks = ['Species', 'Genus', 'Subfamily', 'Family', 'Superfamily', 'Order', 'Suborder','Infraorder', 'Parvorder','?']
    for rank in ranks:
        table.update({rank:{'Totals':0}})
    rankData = csv.reader(open(rankFile,'rU'))
    next(rankData, None) #skip header row
    rankDict = dict()
    # extract the year from author's name
    for d in rankData:
        if "(" in d[1]:
            y =d[1].split("(")[1].split(")")[0]
        else:
            y = d[1]
        rankDict.update({y + "." + d[0].replace(" ", "_") : d[2]})
    inputData = csv.reader(open(inFile))
    for line in inputData:
        #second half of the table
        if line[1] in art2symbol:
            col = art2symbol[line[1]] + " / "
        else:
            col = line[1] + " / "
        name1 = line[0].split(".")[1]
        name2 = line[2].split(".")[1]
        if name1 == name2:
            col += "="
        else:
            col += "<>"
        # if there is a "_" in name, the rank is species, otherwise refer to rank file to get the rank information
        if "_" in line[0]:
            rank1 = 'Species'
        else:
            rank1 = rankDict[line[0]]
        if "_" in line[2]:
            rank2 = 'Species'
        else:       
            rank2 = rankDict[line[2]]
        if rank1 == rank2:
            row = rank1
            if col not in table[row]:
                table[row].update({col:0})
            table[row][col] += 1
            #first half of the table
            g1 = line[0].split(".")[0]
            g2 = line[2].split(".")[0]
            for g in [g1,g2]:
                if g not in elements:
                    elements.update({g:dict()})
            if rank1 not in elements[g1]:
                elements[g1].update({rank1:set()})
            if rank2 not in elements[g2]:
                elements[g2].update({rank2:set()})
            elements[g1][rank1].add(line[0].split(".")[1])
            elements[g2][rank2].add(line[2].split(".")[1])
    for g in elements:
        for e in elements[g]:
            table[e][g] = len(elements[g][e])
    for k in table:
        total = 0
        for v in table[k]:
            if v in ['== / =', '== / <>', '> / =', '< / =' , '>< / =']:
                total += table[k][v]
            table[k]["Totals"] = total
    writer = csv.writer(sys.stdout)
    row = ['Rank', g1, g2, '== / =', '== / <>', '> / =', '> / <>', '< / =' , '< / <>', '>< / =', '>< / <>', '! / <>' ,'Totals']
    writer.writerow(row)
    for r in ranks:
        line = [r]
        for i in row[1:]:
            if i in table[r]:
                line.append(table[r][i])
            elif i != '':
                line.append("0")
        writer.writerow(line)
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-i","--iFile")
    parser.add_argument("-r","--rFile")
    args = parser.parse_args()
    main(args.iFile, args.rFile)
