#!/usr/bin/env python3
"""Zliczanie słów w „Panu Tadeuszu” w Sparku: wersja RDD i wersja DataFrame/SQL.

Tokenizacja jak w pozostałych wersjach demo (Java, Streaming, Hive): małe litery,
podział na ciągi znaków niebędących literami Unicode, puste elementy odrzucone.

Użycie:
    spark-submit wordcount_spark.py {rdd|df|both} {small|large} [--pause SECONDS] [--label yarn]

--pause zostawia aplikację przy życiu po policzeniu, żeby można było obejrzeć Spark UI (port 4040).
--label dopisuje przedrostek do katalogu wynikowego (np. dla uruchomienia na YARN).
Wersja RDD zapisuje wynik przez saveAsTextFile, które nie nadpisuje istniejącego katalogu:
przed ponownym uruchomieniem usuń go (hdfs dfs -rm -r /user/hadoop/output/spark_rdd_large).
"""
import argparse
import re
import sys
import time

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

BASE_PATH = 'hdfs://127.0.0.1:9000/user/hadoop'
INPUT_PATHS = {'small': BASE_PATH + '/pan_tadeusz/small', 'large': BASE_PATH + '/pan_tadeusz/large'}

# [^\W\d_] to litera Unicode (znak „słowa” bez cyfr i podkreślenia)
LETTERS = re.compile(r'[^\W\d_]+')


def count_rdd(spark, input_path, output_path):
    """Klasyczne API RDD: flatMap -> map -> reduceByKey."""
    lines = spark.sparkContext.textFile(input_path)
    counts = (lines
              .flatMap(lambda line: LETTERS.findall(line.lower()))   # każde słowo osobno
              .map(lambda word: (word, 1))                           # para (słowo, 1)
              .reduceByKey(lambda a, b: a + b)                       # suma po słowie
              .cache())                                              # wynik jest mały, więc trzymamy go w pamięci
    counts.map(lambda pair: '%s\t%d' % pair).saveAsTextFile(output_path)
    top = counts.takeOrdered(20, key=lambda pair: (-pair[1], pair[0]))
    return top, counts.count(), counts.values().sum()


def count_dataframe(spark, input_path, output_path):
    """DataFrame i SQL: ten sam wynik, ale optymalizuje go Catalyst, a nie programista."""
    lines = spark.read.text(input_path)                               # jedna kolumna: value
    lines.createOrReplaceTempView('pan_tadeusz')
    counts = spark.sql(r"""
        SELECT word, COUNT(*) AS cnt
        FROM (SELECT explode(split(lower(value), '[^\\p{L}]+')) AS word FROM pan_tadeusz)
        WHERE word <> ''
        GROUP BY word
    """).cache()
    counts.write.mode('overwrite').option('sep', '\t').csv(output_path)
    top = [tuple(row) for row in counts.orderBy(F.desc('cnt'), F.asc('word')).limit(20).collect()]
    return top, counts.count(), counts.agg(F.sum('cnt')).first()[0]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('mode', choices=['rdd', 'df', 'both'])
    parser.add_argument('dataset', choices=['small', 'large'])
    parser.add_argument('--pause', type=int, default=0, help='ile sekund utrzymać aplikację po policzeniu')
    parser.add_argument('--label', default='', help='przedrostek katalogu wynikowego, np. yarn')
    args = parser.parse_args()

    spark = SparkSession.builder.appName('word count Pan Tadeusz (Spark)').getOrCreate()
    spark.sparkContext.setLogLevel('WARN')
    print('Spark %s, Python %s, master=%s' % (spark.version, sys.version.split()[0], spark.sparkContext.master))

    for mode, count_fn in (('rdd', count_rdd), ('df', count_dataframe)):
        if args.mode not in (mode, 'both'):
            continue
        prefix = args.label + '_' if args.label else ''
        output_path = '%s/output/spark_%s%s_%s' % (BASE_PATH, prefix, mode, args.dataset)
        spark.sparkContext.setJobDescription('word count (%s, %s)' % (mode.upper(), args.dataset))
        started = time.time()
        top, distinct_words, total_words = count_fn(spark, INPUT_PATHS[args.dataset], output_path)
        print('\n== wersja %s, zbiór %s: %d różnych słów, %d wszystkich, czas %.1f s' %
              (mode.upper(), args.dataset, distinct_words, total_words, time.time() - started))
        print('== wynik zapisany w %s' % output_path)
        print('== 20 najczęstszych słów:')
        for word, count in top:
            print('%s\t%d' % (word, count))

    if args.pause:
        print('\n== Spark UI: %s (pauza %d s)' % (spark.sparkContext.uiWebUrl, args.pause))
        sys.stdout.flush()
        time.sleep(args.pause)
    spark.stop()


if __name__ == '__main__':
    main()