#
# USAGE: cat URLlist.csv | get_language_coverage.py <index_cover_directory>
# Takes a newspapermap.csv-formatted URL list on stdin
# Takes a pointer to a directory of common crawl index record files
# Prints language coverage statistics and sorts them into language-specific
# directories under the index directory. 
#
# Author : John Hewitt : johnhew@seas.upenn.edu
# Also sorts the data in to one-folder-per-language

import json
import csv
import sys
import os
import collections
import shutil

#Move to the directory with the index output
indexdir = sys.argv[1]
os.chdir(indexdir)

# Create a dictionary mapping URLs to languages and languages to the total webpage counts
URL_to_lang = {}
lang_page_counts = collections.defaultdict(int)

# Populate the URL-to-Lang dictionary with your URL source. 
for line in csv.DictReader(sys.stdin, delimiter='\t'):
    URL_to_lang[line['url'][7:].strip('/').replace('/','-')] = line['language']

# Clean out the language directory if it exists.
if 'langs' in os.listdir('.'):
    shutil.rmtree('./langs')

# Create the base folder for the language-specific directory
os.mkdir('langs/')

# For each file in the index output, (1)determine the language, (2)update file count for the 
# language, and (3)copy the file to its lang-specific folder. 
for count, indexfile in enumerate(os.listdir(sys.argv[1])):
    if count % 100 == 0: print >> sys.stderr, 'processing #',count
    language = None

    # (1)
    URL = indexfile[indexfile.find('--') + 2 : indexfile.rfind('-')]
    try:
        language = URL_to_lang[URL]
        if not language:
          continue
    except:
        print >> sys.stderr, 'no match', URL
        continue

    # (2)
    try:
        for line in open(indexfile):
            lang_page_counts[language] += 1
    except:
        print >> sys.stderr, 'error parsing file', indexfile
        continue

    # (3)
    if language not in os.listdir('./langs'):
        os.mkdir('./langs/'+language)
    try:
        shutil.copy('./'+indexfile, 'langs/'+language+'/'+indexfile)
    except IOError:
        print >> sys.stderr, 'error copying file', indexfile


# Write the number of pages statistics
for lang in sorted(lang_page_counts, key = lambda x : lang_page_counts[x]):
    print lang, lang_page_counts[lang]

