#!/bin/env python

# Run pyspark shell with refinery jars.
# Note that refinery-job is not required, but since it's shaded it solves a bunch of dep imports.
# $ pyspark --jars ./refinery-spark/target/refinery-spark-0.2.33-SNAPSHOT.jar,./refinery-core/target/refinery-core-0.2.33-SNAPSHOT.jar,./refinery-core/target/refinery-core-0.2.33-SNAPSHOT.jar,.//refinery-job/target/refinery-job-0.2.33-SNAPSHOT-shaded.jar 
import os
from py4j.java_gateway import JavaClass, JavaObject

# We need to wrap refinery's HivePartition.
# This approach inefficient, but 
# but deailing  Java conversions (dict -> scalaMap -> ListMap) with varags has been
# fiddly. 
# In general, py4j does not like Scala immutable data structures - and Lists in particular.
# An alternative approach could be wrapping Scala code paths in Java interfaces.

HivePartition: JavaClass = sc._gateway.jvm.org.wikimedia.analytics.refinery.core.HivePartition
Some: JavaClass = sc._gateway.jvm.scala.Some
Tuple2: JavaClass = sc._gateway.jvm.scala.Tuple2

list_map: JavaObject = sc._gateway.jvm.scala.collection.immutable.ListMap()
pmap: dict = {"source_key": "testDataset", "year": "2023", "month": "11", "day": "7", "hour": "0"}
for key, value in pmap.items():
    tpl = Tuple2(key, Some(value))
    list_map = getattr(list_map, "$plus")(tpl)

partition: JavaObject = HivePartition("wmf", "webrequest", None, list_map)

# Initialize pydeequ. This code is taken from the library example
# at https://pypi.org/project/pydeequ/.
# We'll setup some analyzers, and pass their results to refinery's 
# DeequAnalyzersToDataQualityMetrics scala SerDe.
os.environ['SPARK_VERSION'] = '3.3'
import pydeequ

from pydeequ.analyzers import *
from pydeequ.repository import *
from pyspark.sql import Row

df = spark.sparkContext.parallelize([
            Row(a="foo", b=1, c=5),
            Row(a="bar", b=2, c=6),
            Row(a="baz", b=3, c=None)]).toDF()

repository: InMemoryMetricsRepository = InMemoryMetricsRepository(spark)
resultKey: ResultKey = ResultKey(spark, ResultKey.current_milli_time(), {'tag': 'sometag'})

AnalysisRunner(spark) \
    .onData(df) \
    .addAnalyzer(Size()) \
    .addAnalyzer(Completeness("b")) \
    .useRepository(repository) \
    .saveOrAppendResult(resultKey) \
    .run()

# DeequAnalyzersToDataQualityMetrics takes an instance of AnalyzersResult as input.
# This object is not available from the Python api, so we drop into jvm with
# a call to repository.repository instead. This should also ensure that compute metrics are not passed over the 
# python <-> jvm wire (we access jvm objects directly).
analyzers_result: JavaObject = repository.repository.load().get()

DeequAnalyzersToDataQualityMetrics: JavaClass = sc._jvm.org.wikimedia.analytics.refinery.spark.deequ.DeequAnalyzersToDataQualityMetrics
metrics: JavaObject = DeequAnalyzersToDataQualityMetrics(spark._jsparkSession, analyzers_result, partition, "123")
metrics.getAsDataFrame().show()

# Output: 
#+-------------+----------------+-------+--------+-----+------------+------------------+---------------+--------------------+-------------------+
#| dataset_date|            tags| entity|instance|value|        name|      source_table|pipeline_run_id|        partition_id|       partition_ts|
#+-------------+----------------+-------+--------+-----+------------+------------------+---------------+--------------------+-------------------+
#|1709805513911|{tag -> sometag}|Dataset|       *|  3.0|        Size|`wmf`.`webrequest`|            123|source_key=testDa...|2023-11-07 01:00:00|
#|1709805513911|{tag -> sometag}| Column|       b|  1.0|Completeness|`wmf`.`webrequest`|            123|source_key=testDa...|2023-11-07 01:00:00|
#+-------------+----------------+-------+--------+-----+------------+------------------+---------------+--------------------+-------------------+

from refinery import HivePartition
from refinery.dq import DeequAnalyzersToDataQualityMetrics

...
partition: HivePartition = HivePartition("wmf", "webrequest", None, list_map)
DeequAnalyzersToDataQualityMetrics(spark, analyzers_result, partition, "123") 
...