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

#!/usr/bin/env python

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

import csv
import re
from collections import OrderedDict

def csv_features_to_dict(csv_filename):
    csv_features_file = open(csv_filename,"r")
    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

    return csv_features_dict

def csv_groundtruth_to_dict(csv_filename):

    csv_file = open(csv_filename)
    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_array(songkey_combined_dict):
    
    features_array = []

    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_array.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_array.append(row)

    return features_array
        

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

    csv_features_writer.writerows(combined_features_array);



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_features_to_dict(csv_input_filename)

arousal_groundtruth_dict = csv_groundtruth_to_dict(arousal_csv_filename)
valence_groundtruth_dict = csv_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_array = ordered_dict_features_to_array(csv_features_dict)


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