from ChildProject.projects import ChildProject
from ChildProject.annotations import AnnotationManager
from ChildProject.metrics import segments_to_grid, conf_matrix

import numpy as np
import matplotlib.pyplot as plt

speakers = ['CHI', 'OCH', 'FEM', 'MAL']
project = ChildProject('.')
am = AnnotationManager(project)
am.read()

SET_1 = 'vtc'
SET_2 = 'eaf'

intersection = AnnotationManager.intersection(am.annotations, [SET_1, SET_2])

segments = am.get_collapsed_segments(intersection)
segments = segments[segments['speaker_type'].isin(speakers)]

# Y
#vtc = segments_to_grid(segments[segments['set'] == 'vtc'], 0, segments['segment_offset'].max(), 100, 'speaker_type', speakers,none = False)
vtc = segments_to_grid(segments[segments['set'] == SET_1], 0, segments['segment_offset'].max(), 100, 'speaker_type', speakers)

# X
#its = segments_to_grid(segments[segments['set'] == 'its'], 0, segments['segment_offset'].max(), 100, 'speaker_type', speakers,none = False)
its = segments_to_grid(segments[segments['set'] == SET_2], 0, segments['segment_offset'].max(), 100, 'speaker_type', speakers)

confusion_counts = conf_matrix(vtc, its)
print(confusion_counts)
all_positive = np.delete(confusion_counts, -1, 0)
all_negative = np.delete(confusion_counts, -1, 1)

precision = np.delete(all_negative, -1, 0).trace() / all_positive.sum()
recall = np.delete(all_negative, -1, 0).trace() / all_negative.sum()
fscore = (2 * precision * recall) / (precision + recall)

scores = {}
i=0

with open('extra/scores.txt','w') as f:
    for label in speakers:
        rec = confusion_counts[i,i] / confusion_counts[ :,i].sum()
        preci = confusion_counts[i,i] / confusion_counts[i,: ].sum()
        fsc = (2 * preci * rec) / (preci + rec)
        #scores[label] = (preci, rec, fsc)
        f.write(f"{label}: precision {preci}; recall {rec}; F-score {fsc}\n")
        i+=1

    f.write(f"General: precision {precision}; recall {recall}; F-score {fscore}\n")
    #print(f"General: precision {precision}; recall {recall}; F-score {fscore}")

print(f"Results written to scores.txt")

normalized = confusion_counts

speakers.append("None")
speakers = [""] + speakers

print(normalized)
print(speakers)

fig, ax = plt.subplots(figsize=(7.5, 7.5))
ax.set_xticklabels(speakers)  
ax.set_yticklabels(speakers)
ax.matshow(normalized, cmap=plt.cm.Blues, alpha=0.3)
for i in range(normalized.shape[0]):
    for j in range(normalized.shape[1]):
        ax.text(x=j, y=i,s=round(normalized[i, j],3), va='center', ha='center', size='xx-large')
 
ax.xaxis.set_label_position("top")
# set Y and X
plt.ylabel(SET_1, fontsize=18)
plt.xlabel(SET_2, fontsize=18)
plt.title('Confusion Matrix', fontsize=18)
plt.savefig('extra/conf_matrix.png')

normalized = confusion_counts/(np.sum(vtc, axis = 0)[:,None])

fig, ax = plt.subplots(figsize=(7.5, 7.5))
ax.set_xticklabels(speakers)  
ax.set_yticklabels(speakers)
ax.matshow(normalized, cmap=plt.cm.Blues, alpha=0.3)
for i in range(normalized.shape[0]):
    for j in range(normalized.shape[1]):
        ax.text(x=j, y=i,s=round(normalized[i, j],3), va='center', ha='center', size='xx-large')
 
ax.xaxis.set_label_position("top")
plt.ylabel(SET_1, fontsize=18)
plt.xlabel(SET_2, fontsize=18)
plt.title('Confusion Matrix', fontsize=18)
plt.savefig('extra/conf_matrix_normalized.png')