from __future__ import print_function

import sys

from pyspark.sql import SparkSession

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Użycie: mnmcount <file>", file=sys.stderr)
        sys.exit(-1)

    spark = (SparkSession
        .builder
        .appName("PythonMnMCount")
        .getOrCreate())
    # Pobranie z argumentów powłoki nazwy pliku danych dotyczących cukierków M&M’s.
    mnm_file = sys.argv[1]
    # Wczytanie pliku do egzemplarza Spark DataFrame.
    mnm_df = (spark.read.format("csv")
        .option("header", "true")
        .option("inferSchema", "true")
        .load(mnm_file))
    mnm_df.show(n=5, truncate=False)

    # Agregowanie liczby wystąpień kolorów przez wywołanie groupBy() dla pól "State" i "Color".
    # Wywołanie orderBy() i ustawienie wartości w kolejności malejącej.
    count_mnm_df = (mnm_df.select("State", "Color", "Count")
                    .groupBy("State", "Color")
                    .sum("Count")
                    .orderBy("sum(Count)", ascending=False))

    # Wyświetlenie zagregowanych wyników dla wszystkich stanów i kolorów.
    count_mnm_df.show(n=60, truncate=False)
    print("Total Rows = %d" % (count_mnm_df.count()))

    # Zagregowane wyniki dla Kalifornii odszukaj za pomocą filtrowania.
    ca_count_mnm_df = (mnm_df.select("*")
                       .where(mnm_df.State == 'CA')
                       .groupBy("State", "Color")
                       .sum("Count")
                       .orderBy("sum(Count)", ascending=False))

    # Wyświetlenie zagregowanych wyników dla Kalifornii.
    ca_count_mnm_df.show(n=10, truncate=False)

