import os
from matplotlib import pyplot as plt
import numpy as np
import json

def mean(l):
    return sum(l)/len(l)
def stdev(l):
    m = mean(l)
    return (sum([(x-m)**2 for x in l])/len(l))**0.5
subdirs = os.listdir('/ROOT')
subdirs = [x for x in subdirs if os.path.isdir(f'/ROOT/{x}')]
# ignore __pycache__
subdirs = [x for x in subdirs if x != '__pycache__']

aesthetic_jsons = {}
corrupt_jsons = {}
for subdir in subdirs:
    aesthetic_jsons[subdir] = f'/ROOT/{subdir}/aesthetic.json'
    corrupt_jsons[subdir] = f'/ROOT/{subdir}/corrupt.json'

# {k:score}
# get score mean and stdev for each jsons

aesthetic_scores = {}
corrupt_scores = {}
for subdir in subdirs:
    with open(aesthetic_jsons[subdir], 'r') as f:
        aesthetic_scores[subdir] = json.load(f)
    with open(corrupt_jsons[subdir], 'r') as f:
        corrupt_scores[subdir] = json.load(f)

aesthetic_mean = {k:mean(v.values()) for k,v in aesthetic_scores.items()}
aesthetic_stdev = {k:stdev(v.values()) for k,v in aesthetic_scores.items()}
corrupt_mean = {k:mean(v.values()) for k,v in corrupt_scores.items()}
corrupt_stdev = {k:stdev(v.values()) for k,v in corrupt_scores.items()}
# pretty print
for aesthetic_m, aesthetic_s, corrupt_m, corrupt_s in zip(aesthetic_mean, aesthetic_stdev, corrupt_mean, corrupt_stdev):
    print(f'{aesthetic_m} aesthetic: {aesthetic_mean[aesthetic_m]}+- {aesthetic_stdev[aesthetic_s]}')
    print(f'{corrupt_m} corrupt: {corrupt_mean[corrupt_m]}+- {corrupt_stdev[corrupt_s]}')
    print()


# plot histogram (bin=100) with transparency
plt.figure()
for subdir in subdirs:
    plt.hist(list(aesthetic_scores[subdir].values()), bins=30, alpha=0.5, label=subdir, density=True)
plt.legend()
plt.title('Aesthetic Scores')
plt.savefig('/ROOT/aesthetic_scores.png')

plt.figure()
for subdir in subdirs:
    plt.hist(list(corrupt_scores[subdir].values()), bins=60, alpha=0.5, label=subdir, density=True)
plt.legend()
plt.title('Corrupt Scores')
plt.savefig('/ROOT/corrupt_scores.png')

# flat dict
aesthetic_scores_flat = {k:list(v.values()) for k,v in aesthetic_scores.items()}
corrupt_scores_flat = {k:list(v.values()) for k,v in corrupt_scores.items()}

# 90percentile
percentile = 90
aesthetic_percentile = {k:np.percentile(v, percentile) for k,v in aesthetic_scores_flat.items()}
corrupt_percentile = {k:np.percentile(v, percentile) for k,v in corrupt_scores_flat.items()}

print(f'{percentile} percentile')
for aesthetic_name, corrupt_name in zip(aesthetic_percentile, corrupt_percentile):
    print(f'{aesthetic_name} aesthetic: {aesthetic_percentile[aesthetic_name]}')
    print(f'{corrupt_name} corrupt: {corrupt_percentile[corrupt_name]}')
    print()

