from pyspark.sql import SparkSession
from create_mock_contracts import contracts
from operator import add

print("Created mock contracts map.")

import time

store_id_col = 2
country_id_col = 3
transaction_type_id_col = 4
date_col = 1
vendor_id_col = 0
total_col = 7

spark = SparkSession.builder.appName('poc').config('spark.driver.host', '127.0.0.1').getOrCreate()

transactions = spark.read.parquet('data/transactions/')

contracts_broadcast = spark.sparkContext.broadcast(contracts)

print("Starting...")

t0 = time.time()


def match_transaction(transaction):
    vendor_id = transaction[vendor_id_col]
    return ((contract.payee, contract.find_matching_clause(transaction).split * transaction[total_col])
            for contract in contracts_broadcast.value[vendor_id] if contract.matches(transaction))


matched_transactions = transactions.rdd.flatMap(match_transaction)

results = matched_transactions.reduceByKey(add).collect()

print(time.time() - t0)

for result in results:
    print(result)