source: main/trunk/model-sites-dev/mars/collect/deam/04-MERGE-CSV-FEATURES-AND-AV-FILES.py@ 34442

#!/usr/bin/env python

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import with_statement


import csv
import re
from collections import OrderedDict

def csv_load_features_to_dict(csv_filename):
    csv_features_file = open(csv_filename,"r")

    try:
        csv_features_reader = csv.reader(csv_features_file, delimiter=',', quotechar='"')

        csv_features_dict = {}
        
        line_count = 0
        header_row = None
        
        for row in csv_features_reader:
            if line_count == 0:
                row.pop(0)
                header_row = row
            else:
                full_file_id = row[0]
                
                pat = re.compile('^import/(\\d+).json$')
                mat = pat.match(full_file_id)
                song_id = int(mat.group(1))

                row.pop(0)
                row_ordered_dict = OrderedDict()
                
                for i in range(0, len(header_row)):
                    field = header_row[i];
                    value = row[i];
                    
                    row_ordered_dict[field] = value
                
                csv_features_dict[song_id] = row_ordered_dict

            line_count += 1

    finally:
        csv_features_file.close()
        return csv_features_dict

def csv_load_groundtruth_to_dict(csv_filename):

    # Written in the 'with' style that automatically closes the open()
    with open(csv_filename,"r") as csv_file:
        csv_reader = csv.reader(csv_file, delimiter=',')

        csv_gt_dict = {}

        line_count = 0
        header_row = None

        for row in csv_reader:
            if line_count == 0:
                row.pop(0)
                header_row = row
            else:
                song_id = int(row[0])
                row.pop(0)

                row_ordered_dict = OrderedDict()

                for i in range(0, len(row)):
                    field = header_row[i];
                    value = row[i];

                    row_ordered_dict[field] = value
                
                csv_gt_dict[song_id] = row_ordered_dict

            line_count += 1

    return csv_gt_dict


def print_songkey_dict(message,songkey_dict,limit=3):
    print()
    print("****")
    print("* " + message + " (Sample limit = " + str(limit) + "):")
    print("****")

    i = 1
    for song_id_key, ordered_vals in sorted(songkey_dict.items()): 
        print("Song id " + str(song_id_key) + ": \n" + str(ordered_vals))
        if i >= limit:
            break
        i = i + 1

def add_groundtruth_to_features(songkey_features_dict,gt_label,songkey_gt_dict,gt_field):
    new_features_label = gt_label + "_" + gt_field

    # Doesn't need to be sorted, but tidier this way 
    for song_id_key, ordered_feature_vals in sorted(songkey_features_dict.items()): 
        songkey_features_dict[song_id_key][new_features_label] = songkey_gt_dict[song_id_key][gt_field]
    

def ordered_dict_features_to_list(songkey_combined_dict):
    
    features_list = []

    header = []

    # Set up the header labels
    first_key = next(iter(songkey_combined_dict))
    first_ordered_dict = songkey_combined_dict[first_key]

    for feature_key in first_ordered_dict:
        header.append(feature_key)

    features_list.append(header)

    # Now move on to processing each of the song_key_id entries
    #
    # Doesn't need to be sorted, but tidier this way 
    
    for song_id_key, ordered_combined_vals in sorted(songkey_combined_dict.items()): 

        row = []

        for feature_key in ordered_combined_vals:
            feature_val = ordered_combined_vals[feature_key]

            row.append(feature_val)


        features_list.append(row)

    return features_list
        

def csv_save_combined_features(csv_filename,combined_features_list):
    csv_features_file = open(csv_filename,"w",newline="")
    csv_features_writer = csv.writer(csv_features_file, delimiter=',', quotechar='"')

    csv_features_writer.writerows(combined_features_list);



csv_input_filename  = 'etc/deam-essentia-features-collated.csv'
csv_output_filename = 'etc/deam-essentia-features-arousal-valence.csv'

groundtruth_dir = 'prepare/annotations/annotations averaged per song/dynamic (per second annotations)'
arousal_csv_filename = groundtruth_dir + '/arousal.csv'
valence_csv_filename = groundtruth_dir + '/valence.csv'

# Essentia Features extracted: 21 seconds in, for 6 second block
# This equates to DEAM ground-trugh date values:
#   [sample_21000ms,sample_27000ms)
#
# => Take the ground-trugh value at 'sample_26500' as the culmination of
#    the affect the music had on users participating in the experiment

gt_field='sample_26500ms'


csv_features_dict = csv_load_features_to_dict(csv_input_filename)

arousal_groundtruth_dict = csv_load_groundtruth_to_dict(arousal_csv_filename)
valence_groundtruth_dict = csv_load_groundtruth_to_dict(valence_csv_filename)

print_songkey_dict("Arousal Ground-truth",arousal_groundtruth_dict)
print_songkey_dict("Valence Ground-truth",valence_groundtruth_dict)

add_groundtruth_to_features(csv_features_dict,"arousal",arousal_groundtruth_dict,gt_field)
add_groundtruth_to_features(csv_features_dict,"valence",valence_groundtruth_dict,gt_field)

print_songkey_dict("Essentia Features",csv_features_dict,1)

combined_features_list = ordered_dict_features_to_list(csv_features_dict)


print()
print("****")
print("* Saving combined CSV data to: " + csv_output_filename)
print("****")
csv_save_combined_features(csv_output_filename,combined_features_list)