# Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache 2.0 License. 
# You may not use this file except in compliance with the License. A copy of
# the License is located at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.


"""
This script calculates the entropy for a given set of MDNS annotations. 
The mean and median entropy values are printed to the console.
"""

import pandas as pd
import numpy as np
import json
import random
import itertools
import sys
import os
import scipy
import numpy as np
import argparse

from scipy.stats import beta


def calculate_entropy(annotator_ratings):
    """Calculates entropy of the passed `pd.Series`
    """
    frequencies = pd.Series(annotator_ratings).value_counts()           # counts occurrence of each value
    entropy = scipy.stats.entropy(frequencies)  # get entropy from counts
    return entropy

def main(annotation_file):
    # Load raw annotator ratings to Dataframe
    df = pd.read_json(annotation_file, orient='index') 
    e_names = df.columns

    print('\nCalculating entropy per expression dimension for annotation file: ' + annotation_file)
    print('Num. expression dimensions (D): ' + str(len(e_names)))
    print('Num. images (N): '+ str(df.shape[0]) + '\n')
    print('Printing statistics below:')

    entropies = {}
    col=0
    for index, row in df.iterrows():
        entropies[row.name] = {}
        for e in e_names:
            entropy = calculate_entropy(df.at[row.name,e])
            entropies[row.name][e] = entropy


    entropy_df = pd.DataFrame.from_dict(entropies, orient='index')
    print('\nMean Etropy per expression dimension:')
    for e in e_names:
        print(e + ': '+ str(round(entropy_df[e].mean(),2)) + ' STD: '+ str(round(entropy_df[e].std(),2)))

    print()
    print('\nMedian Etropy per expression dimension:')
    for e in e_names:
        print(e + ': '+ str(round(entropy_df[e].median(),2)))


    print('\nDone.')
   
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-p", 
                        "--annotations_file_path", 
                        type = str, 
                        help = "Path to JSON file containing raw annotation data.",
                        required = True)

   
    args = parser.parse_args()
    main(args.annotations_file_path)




