{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"background-color:#1e0d15; padding:20px; border-bottom:3px solid #cc3d56; text-align:center;\"><a href=\"https://www.tmlt.io/\"><img src=\"https://s3.amazonaws.com/tumult.demos/data/images/header.svg\" width=200px></a><div style=\"background-color:#1e0d15; text-align:center; padding-top:15px;  color:white;\">copyright 2022 Tumult Labs</div></div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Execute Simple Private Quantile Calculation with Tumult Core\n",
    "\n",
    "### Import the necessary libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from pyspark.sql import SparkSession\n",
    "\n",
    "from tmlt.core.privacy_framework.components import (\n",
    "    PandasDataFrameAggregationByColumn,\n",
    "    SparkGroupByApplyUDF,\n",
    "    VectorQuantile,\n",
    ")\n",
    "from tmlt.core.privacy_framework.domains import (\n",
    "    NumpyIntegerDomain,\n",
    "    PandasDataFrameDomain,\n",
    "    PandasSeriesDomain,\n",
    "    SparkDataFrameDomain,\n",
    "    SparkIntegerColumnDescriptor,\n",
    "    SparkStringColumnDescriptor,\n",
    ")\n",
    "from tmlt.core.privacy_framework.metrics import SymmetricDifference\n",
    "\n",
    "# If Java 1.8 is not your default, switch to Java 1.8\n",
    "#%env JAVA_HOME=/Library/Java/JavaVirtualMachines/jdk1.8.0_281.jdk/Contents/Home"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Perform the quantile calculation and print the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Main function.\"\"\"\n",
    "spark = SparkSession.builder.config(\n",
    "    \"spark.ui.showConsoleProgress\", \"false\"\n",
    ").getOrCreate()\n",
    "spark.sparkContext.setLogLevel(\"ERROR\")\n",
    "\n",
    "sdf = spark.createDataFrame(  # pylint: disable=no-member\n",
    "    pd.DataFrame(\n",
    "        [\n",
    "            [\"F\", 28],\n",
    "            [\"F\", 26],\n",
    "            [\"F\", 27],\n",
    "            [\"M\", 23],\n",
    "            [\"F\", 29],\n",
    "            [\"M\", 22],\n",
    "            [\"M\", 24],\n",
    "            [\"M\", 25],\n",
    "        ],\n",
    "        columns=[\"Sex\", \"Age\"],\n",
    "    )\n",
    ")\n",
    "\n",
    "print(\"Dataframe:\")\n",
    "sdf.show()\n",
    "\n",
    "vector_quantile = VectorQuantile(\n",
    "    input_domain=PandasSeriesDomain(NumpyIntegerDomain()),\n",
    "    quantile=0.5,\n",
    "    lower=22,\n",
    "    upper=29,\n",
    "    epsilon=1,\n",
    ")\n",
    "\n",
    "df_aggregation_function = PandasDataFrameAggregationByColumn(\n",
    "    input_domain=PandasDataFrameDomain(\n",
    "        {\"Age\": PandasSeriesDomain(NumpyIntegerDomain())}\n",
    "    ),\n",
    "    input_metric=SymmetricDifference(),\n",
    "    column_to_aggregation={\"Age\": vector_quantile},\n",
    ")\n",
    "measure = SparkGroupByApplyUDF(\n",
    "    input_domain=SparkDataFrameDomain(\n",
    "        {\n",
    "            \"Sex\": SparkStringColumnDescriptor(),\n",
    "            \"Age\": SparkIntegerColumnDescriptor(),\n",
    "        }\n",
    "    ),\n",
    "    input_metric=SymmetricDifference(),\n",
    "    groupby_domains={\"Sex\": [\"M\", \"F\"]},\n",
    "    aggregation_function=df_aggregation_function,\n",
    ")\n",
    "\n",
    "print(\"Output:\")\n",
    "measure(sdf).withColumnRenamed(\"Age\", \"Noisy Median Age\").show()\n",
    "\n",
    "print(f\"Privacy relation with epsilon=1: {measure.privacy_relation(1, 1)}\")\n",
    "print(f\"Privacy relation with epsilon=.9: {measure.privacy_relation(1, '0.9')}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div style=\"background-color:#1e0d15; padding:20px; border-top:3px solid #cc3d56; text-align:center;\"><a href=\"https://www.tmlt.io/\"><img src=\"https://s3.amazonaws.com/tumult.demos/data/images/footer.svg\" width=200px></a></div><div style=\"background-color:#1e0d15; text-align:center; padding-bottom:20px;  color:white;\">Privacy protection, redefined</div>"
   ]
  }
 ],
 "metadata": {
  "code_folding": [
   1
  ],
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
