"""Takes the JSON file generated by the raw_to_json script and generates a
TF-IDF sparse matrix.
"""
from collections import OrderedDict
import codecs
import math
import json

import numpy as np
from sklearn.preprocessing import normalize

from dj_tiny_data import paths


def calculate_tf_idf_matrix(input_path=paths.OUTPUT_FILE,
                            output_path_matrix=paths.TF_IDF_MATRIX_PATH,
                            output_path_word_list=paths.WORD_LIST_PATH):
    with open(input_path, 'r') as finput:
        word_id_map = OrderedDict()
        track_list = list()
        for line in finput:
            track_info = json.loads(line)
            track_list.append(track_info)
            for word_id in track_info['wordcount']:
                if word_id not in word_id_map:
                    word_id_map[word_id] = {
                        'doc_freq': 0,
                        'index': len(word_id_map)
                    }
                word_id_map[word_id]['doc_freq'] += 1

    tf_idf_dense = np.zeros(shape=(len(track_list), len(word_id_map)))

    for track, row in zip(track_list, tf_idf_dense):
        for word_id in track['wordcount']:
            tf = math.sqrt(int(track['wordcount'][word_id]))
            idf = 1 + math.log(float(len(track_list)) /
                               (word_id_map[word_id]['doc_freq'] + 1))
            row[word_id_map[word_id]['index']] = tf*idf
    normalize(tf_idf_dense, norm='l2', axis=1, copy=False)
    np.savetxt(output_path_matrix, tf_idf_dense, delimiter=',')
    foutput = codecs.open(output_path_word_list, 'w', 'utf-8')
    for word_id in word_id_map:
        foutput.write('%s,%s\n' % (word_id, word_id_map[word_id]['doc_freq']))
    foutput.close()
