#!/usr/bin/env python3

"""Input: CoNLL-2000 style labeled file where columns are separated by a single
space, the second-to-last column contains the gold label and the last column
contains the predicted label. Output: a confusion matrix."""

import sys
from collections import OrderedDict, Counter

count = 0
errors = 0
matrix = {}

# Read and count
for line in sys.stdin:
    line = line.rstrip()
    if not line:
        continue
    fields = line.split(' ')
    gold = fields[-2]
    predicted = fields[-1]
    count += 1
    if gold != predicted:
        errors += 1
    if gold not in matrix:
        matrix[gold] = {}
    row = matrix[gold]
    if predicted not in row:
        row[predicted] = 0
    row[predicted] += 1

labels = list(matrix.keys())
freq_by_label = {label: sum(matrix[label].values()) for label in labels}
sortkey = lambda x: -freq_by_label[x]
labels.sort(key=sortkey)

# Error rate
print(f'Annotated units: {count}')
print(f'Errors:          {errors}')
print(f'Error rate:      {errors / count:.6f}\n')

# Print confusion matrix
def print_matrix_table(matrix, rowsortkey=None, columnsortkey=None, defaultvalue=0):
    row_keys = list(matrix.keys())
    if rowsortkey:
        row_keys.sort(key=rowsortkey)

    all_columns = set()
    for row in matrix.values():
        all_columns.update(row.keys())
    col_keys = list(all_columns)
    if columnsortkey:
        col_keys.sort(key=columnsortkey)

    header = [""] + col_keys
    col_widths = [max(len(str(x)) for x in col) for col in zip(*([header] + [
        [row] + [matrix.get(row, {}).get(col, defaultvalue) for col in col_keys]
        for row in row_keys
    ]))]

    # Print header
    header_row = "  ".join(f"{col:<{col_widths[i]}}" for i, col in enumerate(header))
    print(header_row)

    # Print rows
    for row in row_keys:
        row_data = [row] + [matrix.get(row, {}).get(col, defaultvalue) for col in col_keys]
        print("  ".join(f"{str(cell):<{col_widths[i]}}" for i, cell in enumerate(row_data)))
    print()

print_matrix_table(matrix, rowsortkey=sortkey, columnsortkey=sortkey)

# Scores for individual labels
countstable = OrderedDict((label, Counter()) for label in labels)

for gold in labels:
    for predicted in matrix[gold]:
        c = matrix[gold][predicted]
        if gold == predicted:
            countstable[gold]['tp'] += c
        else:
            countstable[gold]['fn'] += c
            countstable[predicted]['fp'] += c

for label in labels:
    row = countstable[label]
    for m in ('tp', 'fn', 'fp'):
        if m not in row:
            row[m] = 0
    try:
        row['prec'] = row['tp'] / (row['tp'] + row['fp'])
    except ZeroDivisionError:
        row['prec'] = 'N/A'
    try:
        row['rec'] = row['tp'] / (row['tp'] + row['fn'])
    except ZeroDivisionError:
        row['rec'] = 'N/A'
    try:
        if isinstance(row['prec'], float) and isinstance(row['rec'], float):
            row['f1'] = 2 * row['prec'] * row['rec'] / (row['prec'] + row['rec'])
        else:
            row['f1'] = 'N/A'
    except:
        row['f1'] = 'N/A'

def print_aligned_table(table):
    # Get all row and column labels
    row_labels = list(table.keys())
    column_labels = set()
    for row in table.values():
        column_labels.update(row.keys())
    column_labels = sorted(column_labels)

    # Determine the max width of each column
    col_widths = {}
    for col in column_labels:
        max_width = len(col)
        for row in table.values():
            val = row.get(col, "")
            if isinstance(val, float):
                val_str = str(val)
            else:
                val_str = str(val)
            max_width = max(max_width, len(val_str))
        col_widths[col] = max_width

    # Width for row label column
    row_label_width = max(len("Label"), max(len(row) for row in row_labels))

    # Print header
    header = f"{'Label':<{row_label_width}}"
    for col in column_labels:
        header += f"  {col:<{col_widths[col]}}"
    print(header)

    # Print rows
    for row_label in row_labels:
        row = table[row_label]
        line = f"{row_label:<{row_label_width}}"
        for col in column_labels:
            val = row.get(col, "")
            val_str = str(val)
            line += f"  {val_str:<{col_widths[col]}}"
        print(line)

print_aligned_table(countstable)
