# Start the Spark Session
# This uses local mode for simplicity
# the use of findspark is optional

# install pyspark if needed
# ! pip install pyspark

# import findspark
# findspark.init("/home/luca/Spark/spark-3.3.0-bin-hadoop3")

from pyspark.sql import SparkSession
spark = (SparkSession.builder
         .appName("PySpark histograms")
         .master("local[*]")
         .getOrCreate()
        )

"""
Explanation: Weighted histograms with the Apache Spark DataFrame API
This provides and example of how to generate weighted histograms using the Spark DataFrame API.
A weighted histogram is a generalization of a frequency histogram. It takes into account weights for each value entry. When all weights are set to 1, the weighted histogram is a frequency/count histogram.
Disambiguation: we refer here to computing histograms of the DataFrame data, rather than histograms of the columns statistics used by the cost based optimizer.
End of explanation
"""


num_events = 100
scale = 100
seed1 = 4242
seed2 = 4243

df = spark.sql(f"""
select random({seed1}) * {scale} as random_value,
       random({seed2}) as weight 
from range({num_events})
""")

df.show(5)

"""
Explanation: Generate a DataFrame with toy data for demo purposes
End of explanation
"""


# import the computeWeightedHistogram function 
# see implementation details at: 
# https://github.com/LucaCanali/Miscellaneous/blob/master/Spark_Notes/Spark_Histograms/python/sparkhistogram/histogram.py

# requires the package sparkhistogram
! pip install sparkhistogram

from sparkhistogram import computeWeightedHistogram


# Compute the histogram using the computeWeightedHistogram function

histogram = computeWeightedHistogram(df, "random_value", "weight", -20, 90, 11)

# Alternative syntax: compute the histogram using transform on the DataFrame
# requires Spark 3.3.0 or higher
# histogram = df.transform(computeWeightedHistogram, "random_value", "weight", -20, 90, 11)

# this triggers the computation as show() is an action
histogram.show()

# Fetch the histogram data into a Pandas DataFrame for visualization
# At this stage data is reduced to a small number of rows (one row per bin)
# so it can be easily handled by one machine

# toPandas() is an action and triggers the computation
hist_pandasDF = histogram.toPandas()

hist_pandasDF

"""
Explanation: Compute the weighted histogram
End of explanation
"""


import matplotlib.pyplot as plt 
plt.style.use('seaborn-darkgrid')
plt.rcParams.update({'font.size': 20, 'figure.figsize': [14,10]})

f, ax = plt.subplots()

# histogram data
x = hist_pandasDF["value"]
y = hist_pandasDF["weighted_sum"]

# bar plot
ax.bar(x, y, width = 3.0, color='red')

ax.set_xlabel('Bucket values')
ax.set_ylabel('Sum of the weights (weighted count)')
ax.set_title("Distribution of event weight sums")

# Label for the resonances spectrum peaks
txt_opts = {'horizontalalignment': 'center',
            'verticalalignment': 'center',
            'transform': ax.transAxes}

plt.show()

spark.stop()

"""
Explanation: Histogram plotting
This plots the weighted histogram: wheighted sum of the number of per bin.
End of explanation
"""


from pyspark.sql.functions import sum

def computeWeightedHistogram(df: "DataFrame", value_col: str, weight_col: str,
                             min_val: float, max_val: float, bins: int) -> "DataFrame":
    """ This is a dataframe function to compute the weighted histogram of a DataFrame column.
        A weighted histogram is a generalization of a frequency histogram.
        Parameters
        ----------
        df: the dataframe with the data to compute
        value_col: column name on which to compute the histogram
                   the column needs to be of numeric type
        weight_col: numeric-type column with the weights,
                    the bucket value is computed as sum of weights.
                    If all weight are set to 1, you get a frequency histogram
        min_val: minimum value in the histogram
        max_val: maximum value in the histogram
        bins: number of histogram buckets to compute
        Output DataFrame
        ----------------
        bucket: the bucket number, range from 1 to bins (included)
        value: midpoint value of the given bucket
        weighted_sum: weighted sum of the number of values in the bucket
    """
    step = (max_val - min_val) / bins
    # this will be used to fill in for missing buckets, i.e. buckets with no corresponding values
    df_buckets = spark.sql(f"select id+1 as bucket from range({bins})")
    histdf = (df
              .selectExpr(f"width_bucket({value_col}, {min_val}, {max_val}, {bins}) as bucket", f"{weight_col}")
              .groupBy("bucket")
              .agg(sum(f"{weight_col}").alias("weighted_sum"))  # sum the weights on the weight_col
              .join(df_buckets, "bucket", "right_outer") # add missing buckets and remove buckets out of range
              .selectExpr("bucket", f"{min_val} + (bucket - 0.5) * {step} as value",  # use center value of the buckets
                          "nvl(weighted_sum, 0) as weighted_sum") # buckets with no values will have a sum of 0
              .orderBy("bucket")
              )
    return histdf

"""
Explanation: Note added
Use this to define the computeWeightedHistogram function if you cannot pip install sparkhistogram
End of explanation
"""


