#!/usr/bin/env python

"""
Export track metadata from the SQLite database
"""
import logging
import sqlite3
import argparse
import numpy as np
import pandas as pd


def main():
    logger = logging.getLogger(__name__)
    logging.basicConfig(level=logging.INFO)

    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('path_to_metadata', type=str,
                        help='Path to SQLite database')
    parser.add_argument('path_to_genre', type=str,
                        help='Path to genres file')
    parser.add_argument('path_to_location', type=str,
                        help='Path to location file')
    parser.add_argument('path_to_mxm_data', type=str,
                        help='Path to musixmatch dataset')
    parser.add_argument('path_to_language_data', type=str,
                        help='Path to langauge data')
    parser.add_argument('path_to_output', type=str,
                        help='Output path for the feather file')

    args = parser.parse_args()
    path_to_metadata = args.path_to_metadata
    path_to_genre = args.path_to_genre
    path_to_location = args.path_to_location
    path_to_mxm_data = args.path_to_mxm_data
    path_to_language_data = args.path_to_language_data
    path_to_output = args.path_to_output

    logger.info('Loading data...')

    # load tracks in the mxm dataset
    mxm = pd.read_feather(path_to_mxm_data)
    mxm_tracks = mxm[['track_id']]

    # load songs table, it has some columns with metadata
    conn = sqlite3.connect(path_to_metadata)
    songs = pd.read_sql("SELECT * FROM songs", conn)

    logger.info('Replacing 0 with NAs in year column...')
    songs['year'] = songs.year.replace(0, np.nan)

    # rename year column to avoid clashing with "year" word in bag of words
    songs.rename(columns={'year': 'release_year'}, inplace=True)

    logger.info('Subsetting songs to match the musixmatch dataset...')
    df = mxm_tracks.merge(songs, how='left', on='track_id')

    # add genre
    genre = pd.read_csv(path_to_genre, sep='\t', comment='#',
                        names=['track_id', 'genre'])
    df = df.merge(genre, how='left', on='track_id')

    # add location
    location = pd.read_csv(path_to_location, sep='<SEP>',
                           names=['artist_id', 'latitude', 'longitude',
                                  'artist_name', 'location'],
                           engine='python')
    location.drop('artist_name', axis=1, inplace=True)
    df = df.merge(location, how='left', on='artist_id')

    # merge language
    language = pd.read_feather(path_to_language_data)
    df = df.merge(language, how='left', on='track_id')

    nas = df.apply(lambda c: np.sum(c.isna()), axis=0) / df.shape[0] * 100
    logger.info('NA stats: {}'.format(nas))

    # save this
    df = df.reset_index()
    df.drop('index', axis=1, inplace=True)

    logger.info('Saving file in {}'.format(path_to_output))
    df.to_feather(path_to_output)


if __name__ == "__main__":
    main()
