# Copyright 2025 TOYOTA MOTOR CORPORATION.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations
from typing import Any, Mapping
import uuid
import copy
import pytest
import pandas as pd
import numpy as np
import privjail as pj
import privjail.pandas as ppd
import privjail.numpy as pnp

def load_dataframe() -> tuple[ppd.PrivDataFrame, pd.DataFrame]:
    data = {
        "a": [1, 2, 3, 4, 5],
        "b": [2, 4, 4, 4, 3],
    }
    domains = {
        "a": ppd.RealDomain(dtype="int64", range=(None, None)),
        "b": ppd.RealDomain(dtype="int64", range=(None, None)),
    }
    accountant = pj.PureDPAccountant()
    accountant.set_as_root(name=str(uuid.uuid4()))
    pdf = ppd.PrivDataFrame(data, domains=domains, distance=pj.RealExpr(1), accountant=accountant)
    df = pd.DataFrame(data)
    assert (pdf.columns == df.columns).all()
    assert (pdf._value == df).all().all()
    return pdf, df

def assert_budget_spent(pdf: ppd.PrivDataFrame, expected_budget: float) -> None:
    root_name = pdf.accountant._root_name
    assert pj.budgets_spent()[root_name] == pytest.approx(("pure", expected_budget))

def assert_equal_sensitive_series(sensitive_ser: ppd.SensitiveSeries[Any], expected_ser: pd.Series[Any]) -> None:
    assert isinstance(sensitive_ser, ppd.SensitiveSeries)
    assert (sensitive_ser.index == expected_ser.index).all()
    assert (sensitive_ser._value == expected_ser).all()

def assert_equal_sensitive_dataframes(sensitive_df: ppd.SensitiveDataFrame, expected_df: pd.DataFrame) -> None:
    assert isinstance(sensitive_df, ppd.SensitiveDataFrame)
    assert (sensitive_df.index == expected_df.index).all()
    assert (sensitive_df.columns == expected_df.columns).all()
    assert (sensitive_df._value == expected_df).all().all()

def merge_domains(domains: Mapping[str, ppd.Domain], updates: Mapping[str, ppd.Domain]) -> dict[str, ppd.Domain]:
    merged: dict[str, ppd.Domain] = dict(domains)
    merged.update(updates)
    return merged

def test_priv_dataframe_size() -> None:
    pdf, df = load_dataframe()

    # The `shape` member should be a pair of a sensitive value (row) and a non-sensitive value (column)
    assert isinstance(pdf.shape, tuple)
    assert len(pdf.shape) == len(df.shape) == 2
    assert isinstance(pdf.shape[0], pj.Prisoner)
    assert pdf.shape[0]._value == df.shape[0]
    assert pdf.shape[0].max_distance == 1
    assert pdf.shape[1] == len(pdf.columns) == len(df.columns)

    # The `size` member should be a sensitive value
    assert isinstance(pdf.size, pj.Prisoner)
    assert pdf.size._value == df.size
    assert pdf.size.max_distance == len(pdf.columns) == len(df.columns)

    # Builtin `len()` function should raise an error because it must be an integer value
    with pytest.raises(pj.DPError):
        len(pdf)

def test_priv_dataframe_comp() -> None:
    pdf, df = load_dataframe()

    # A non-sensitive value should be successfully compared against a private dataframe
    assert isinstance(pdf == 3, ppd.PrivDataFrame)
    assert isinstance(pdf != 3, ppd.PrivDataFrame)
    assert isinstance(pdf <  3, ppd.PrivDataFrame)
    assert isinstance(pdf <= 3, ppd.PrivDataFrame)
    assert isinstance(pdf >  3, ppd.PrivDataFrame)
    assert isinstance(pdf >= 3, ppd.PrivDataFrame)
    assert ((pdf == 3)._value == (df == 3)).all().all()
    assert ((pdf != 3)._value == (df != 3)).all().all()
    assert ((pdf <  3)._value == (df <  3)).all().all()
    assert ((pdf <= 3)._value == (df <= 3)).all().all()
    assert ((pdf >  3)._value == (df >  3)).all().all()
    assert ((pdf >= 3)._value == (df >= 3)).all().all()

    # A non-sensitive value should be successfully compared against a private series
    assert isinstance(pdf["a"] == 3, ppd.PrivSeries)
    assert isinstance(pdf["a"] != 3, ppd.PrivSeries)
    assert isinstance(pdf["a"] <  3, ppd.PrivSeries)
    assert isinstance(pdf["a"] <= 3, ppd.PrivSeries)
    assert isinstance(pdf["a"] >  3, ppd.PrivSeries)
    assert isinstance(pdf["a"] >= 3, ppd.PrivSeries)
    assert ((pdf["a"] == 3)._value == (df["a"] == 3)).all()
    assert ((pdf["a"] != 3)._value == (df["a"] != 3)).all()
    assert ((pdf["a"] <  3)._value == (df["a"] <  3)).all()
    assert ((pdf["a"] <= 3)._value == (df["a"] <= 3)).all()
    assert ((pdf["a"] >  3)._value == (df["a"] >  3)).all()
    assert ((pdf["a"] >= 3)._value == (df["a"] >= 3)).all()

    # An irrelevant, non-private dataframe should not be compared against a private dataframe
    with pytest.raises(TypeError): pdf == df
    with pytest.raises(TypeError): pdf != df
    with pytest.raises(TypeError): pdf <  df
    with pytest.raises(TypeError): pdf <= df
    with pytest.raises(TypeError): pdf >  df
    with pytest.raises(TypeError): pdf >= df

    # An irrelevant, non-private 2d array should not be compared against a private dataframe
    with pytest.raises(TypeError): pdf == [[list(range(len(pdf.columns)))] for x in range(5)]
    with pytest.raises(TypeError): pdf != [[list(range(len(pdf.columns)))] for x in range(5)]
    with pytest.raises(TypeError): pdf <  [[list(range(len(pdf.columns)))] for x in range(5)]
    with pytest.raises(TypeError): pdf <= [[list(range(len(pdf.columns)))] for x in range(5)]
    with pytest.raises(TypeError): pdf >  [[list(range(len(pdf.columns)))] for x in range(5)]
    with pytest.raises(TypeError): pdf >= [[list(range(len(pdf.columns)))] for x in range(5)]

    # An irrelevant, non-private array should not be compared against a private series
    with pytest.raises(TypeError): pdf["a"] == [0, 1, 2, 3, 4]
    with pytest.raises(TypeError): pdf["a"] != [0, 1, 2, 3, 4]
    with pytest.raises(TypeError): pdf["a"] <  [0, 1, 2, 3, 4]
    with pytest.raises(TypeError): pdf["a"] <= [0, 1, 2, 3, 4]
    with pytest.raises(TypeError): pdf["a"] >  [0, 1, 2, 3, 4]
    with pytest.raises(TypeError): pdf["a"] >= [0, 1, 2, 3, 4]

    accountant = pj.PureDPAccountant()
    accountant.set_as_root(name=str(uuid.uuid4()))
    x = pj.Prisoner(value=0, distance=pj.RealExpr(1), accountant=accountant)

    # A sensitive value should not be compared against a private dataframe
    with pytest.raises(TypeError): pdf == x
    with pytest.raises(TypeError): pdf != x
    with pytest.raises(TypeError): pdf <  x
    with pytest.raises(TypeError): pdf <= x
    with pytest.raises(TypeError): pdf >  x
    with pytest.raises(TypeError): pdf >= x

    # A sensitive value should not be compared against a private series
    with pytest.raises(TypeError): pdf["a"] == x
    with pytest.raises(TypeError): pdf["a"] != x
    with pytest.raises(TypeError): pdf["a"] <  x
    with pytest.raises(TypeError): pdf["a"] <= x
    with pytest.raises(TypeError): pdf["a"] >  x
    with pytest.raises(TypeError): pdf["a"] >= x

    # Sensitive dataframes of potentially different size should not be compared
    pdf_ = pdf[pdf["a"] >= 0]
    with pytest.raises(pj.DPError): pdf == pdf_
    with pytest.raises(pj.DPError): pdf != pdf_
    with pytest.raises(pj.DPError): pdf <  pdf_
    with pytest.raises(pj.DPError): pdf <= pdf_
    with pytest.raises(pj.DPError): pdf >  pdf_
    with pytest.raises(pj.DPError): pdf >= pdf_

    # Sensitive series of potentially different size should not be compared
    with pytest.raises(pj.DPError): pdf["a"] == pdf_["a"]
    with pytest.raises(pj.DPError): pdf["a"] != pdf_["a"]
    with pytest.raises(pj.DPError): pdf["a"] <  pdf_["a"]
    with pytest.raises(pj.DPError): pdf["a"] <= pdf_["a"]
    with pytest.raises(pj.DPError): pdf["a"] >  pdf_["a"]
    with pytest.raises(pj.DPError): pdf["a"] >= pdf_["a"]

def test_priv_dataframe_getitem() -> None:
    pdf, df = load_dataframe()

    # A single-column view should be successfully retrieved from a private dataframe
    assert isinstance(pdf["a"], ppd.PrivSeries)
    assert (pdf["a"]._value == df["a"]).all()

    # A multi-column view should be successfully retrieved from a private dataframe
    assert isinstance(pdf[["a", "b"]], ppd.PrivDataFrame)
    assert (pdf[["a", "b"]]._value == df[["a", "b"]]).all().all()

    # An irrelevant, non-sensitve bool vector should not be accepted for filtering (dataframe)
    with pytest.raises(TypeError):
        pdf[[True, True, False, False, True]]

    # An irrelevant, non-sensitve bool vector should not be accepted for filtering (series)
    with pytest.raises(TypeError):
        pdf["a"][[True, True, False, False, True]]

    # A bool-filtered view should be successfully retrieved from a private dataframe
    assert isinstance(pdf[pdf["a"] > 3], ppd.PrivDataFrame)
    assert (pdf[pdf["a"] > 3]._value == df[df["a"] > 3]).all().all()

    # A bool-filtered view should be successfully retrieved from a private series
    assert isinstance(pdf["a"][pdf["a"] > 3], ppd.PrivSeries)
    assert (pdf["a"][pdf["a"] > 3]._value == df["a"][df["a"] > 3]).all()

    # A sensitve bool vector of potentially different size should not be accepted for filtering (dataframe)
    pdf_ = pdf[pdf["a"] >= 0]
    with pytest.raises(pj.DPError):
        pdf[pdf_["a"] > 3]

    # A sensitve bool vector of potentially different size should not be accepted for filtering (series)
    with pytest.raises(pj.DPError):
        pdf["a"][pdf_["a"] > 3]

    accountant = pj.PureDPAccountant()
    accountant.set_as_root(name=str(uuid.uuid4()))
    x = pj.Prisoner(value=0, distance=pj.RealExpr(1), accountant=accountant)

    # A sensitive value should not be used as a column name
    with pytest.raises(TypeError):
        pdf[x]

    # A slice should not be used for selecting rows of a dataframe
    # TODO: it might be legal to accept a slice
    with pytest.raises(TypeError):
        pdf[2:5]

    # A slice should not be used for selecting rows of a series
    # TODO: it might be legal to accept a slice
    with pytest.raises(TypeError):
        pdf["a"][2:5]

def test_priv_dataframe_setitem() -> None:
    pdf, df = load_dataframe()

    # A non-sensitive value should be successfully assigned to a single-column view
    pdf["c"] = 10
    df["c"] = 10
    assert (pdf.columns == df.columns).all()
    assert (pdf._value == df).all().all()

    # A private series should be successfully assigned to a single-column view
    pdf["d"] = pdf["a"]
    df["d"] = df["a"]
    assert (pdf.columns == df.columns).all()
    assert (pdf._value == df).all().all()

    # A non-sensitive value should be successfully assigned to a multi-column view
    pdf[["e", "f"]] = 100
    df[["e", "f"]] = 100
    assert (pdf.columns == df.columns).all()
    assert (pdf._value == df).all().all()

    # A private dataframe should be successfully assigned to a multi-column view
    pdf[["g", "h"]] = pdf[["a", "b"]]
    df[["g", "h"]] = df[["a", "b"]]
    assert (pdf.columns == df.columns).all()
    assert (pdf._value == df).all().all()

    # A non-sensitive value should be successfully assigned to a bool-filtered view
    pdf[pdf["a"] < 3] = 8
    df[df["a"] < 3] = 8
    assert (pdf.columns == df.columns).all()
    assert (pdf._value == df).all().all()

    # A non-sensitive array should be successfully assigned to a bool-filtered view
    pdf[pdf["a"] == 8] = list(range(len(pdf.columns)))
    df[df["a"] == 8] = list(range(len(df.columns)))
    assert (pdf.columns == df.columns).all()
    assert (pdf._value == df).all().all()

    # A private dataframe should be successfully assigned to a bool-filtered view
    pdf[pdf["a"] < 3] = pdf[pdf["a"] < 4]
    df[df["a"] < 3] = df[df["a"] < 4]
    assert (pdf.columns == df.columns).all()
    assert (pdf._value == df).all().all()

    accountant = pj.PureDPAccountant()
    accountant.set_as_root(name=str(uuid.uuid4()))
    x = pj.Prisoner(value=0, distance=pj.RealExpr(1), accountant=accountant)

    # A sensitive value should not be assigned to a single-column view
    with pytest.raises(TypeError):
        pdf["x"] = x

    # A sensitive value should not be assigned to a multi-column view
    with pytest.raises(TypeError):
        pdf[["x", "y"]] = x

    # A sensitive value should not be assigned to a bool-filtered view
    with pytest.raises(TypeError):
        pdf[pdf["a"] > 3] = x

    # An irrelevant, non-private array should not be assigned to a column view
    with pytest.raises(TypeError):
        pdf["x"] = [1, 2, 3, 4, 5]

    # An irrelevant, non-private 2d array should not be assigned to a bool-filtered view
    with pytest.raises(TypeError):
        pdf[pdf["a"] > 3] = [[list(range(len(pdf.columns)))] for x in range(5)]

    # A sensitive value should not be used as a column name
    with pytest.raises(TypeError):
        pdf[x] = 10

    # A sensitve bool vector of potentially different size should not be accepted for filtering (dataframe)
    pdf_ = pdf[pdf["a"] >= 0]
    with pytest.raises(pj.DPError):
        pdf[pdf_["a"] > 3] = 10

    # A sensitve bool vector of potentially different size should not be accepted for filtering (series)
    with pytest.raises(pj.DPError):
        pdf["a"][pdf_["a"] > 3] = 10

    # A slice should not be used for selecting rows
    # TODO: it might be legal to accept a slice
    with pytest.raises(TypeError):
        pdf[2:5] = 0

def test_priv_dataframe_setitem_numpy() -> None:
    pdf, df = load_dataframe()

    # A numpy array should be successfully assigned to a column
    pdf["a_"] = pj.clip_norm(pdf["a"].values, bound=3.0, ord=1)
    assert pdf["a_"]._value.to_list() == pytest.approx([1., 2., 3., 3., 3.])
    assert pdf["a_"].domain.dtype == "float64"
    assert pdf["a_"].domain.range == pytest.approx((-3.0, 3.0))

def test_priv_dataframe_replace() -> None:
    pdf, df = load_dataframe()

    # Default behaviour
    assert (pdf.replace(4, 10)._value == df.replace(4, 10)).all().all()

def test_priv_series_replace() -> None:
    pdf, df = load_dataframe()

    # Default behaviour
    assert (pdf["b"].replace(4, 10)._value == df["b"].replace(4, 10)).all()

def test_priv_dataframe_dropna() -> None:
    pdf, df = load_dataframe()

    pdf = pdf.replace(3, np.nan)
    df = df.replace(3, np.nan)

    # Default behaviour
    assert (pdf.dropna()._value == df.dropna()).all().all()

    # Should return an error with ignore_index=True
    with pytest.raises(pj.DPError):
        pdf.dropna(ignore_index=True)

def test_priv_series_dropna() -> None:
    pdf, df = load_dataframe()

    pdf = pdf.replace(3, np.nan)
    df = df.replace(3, np.nan)

    # Default behaviour
    assert (pdf["b"].dropna()._value == df["b"].dropna()).all()

    # Should return an error with ignore_index=True
    with pytest.raises(pj.DPError):
        pdf["b"].dropna(ignore_index=True)

def test_priv_series_value_counts() -> None:
    pdf, df = load_dataframe()

    # Should return an error without arguments
    with pytest.raises(pj.DPError):
        pdf["b"].value_counts()

    # Should return an error for non-categorical series
    with pytest.raises(pj.DPError):
        pdf["b"].value_counts(sort=False)

    pdf._domains = merge_domains(pdf._domains, {"b": ppd.CategoryDomain(categories=[1, 2, 3, 4, 5])}) # FIXME: use public API

    # Should return an error with `sort=True`
    with pytest.raises(pj.DPError):
        pdf["b"].value_counts()

    # Should return correct counts when all possible values are provided
    pdf._domains = merge_domains(pdf._domains, {"b": ppd.CategoryDomain(categories=[2, 3, 4])}) # FIXME: use public API
    counts = pdf["b"].value_counts(sort=False)
    assert_equal_sensitive_series(counts, pd.Series({2: 1, 3: 1, 4: 3}))
    assert counts[2].max_distance == 1

    # Should return correct counts when only a part of possible values are provided
    pdf._domains = merge_domains(pdf._domains, {"b": ppd.CategoryDomain(categories=[3, 4])}) # FIXME: use public API
    counts = pdf["b"].value_counts(sort=False)
    assert_equal_sensitive_series(counts, pd.Series({3: 1, 4: 3}))
    assert counts[3].max_distance == 1

    # Should return correct counts when non-existent values are provided
    pdf._domains = merge_domains(pdf._domains, {"b": ppd.CategoryDomain(categories=[1, 3, 4, 5])}) # FIXME: use public API
    counts = pdf["b"].value_counts(sort=False)
    assert_equal_sensitive_series(counts, pd.Series({1: 0, 3: 1, 4: 3, 5: 0}))
    assert counts[1].max_distance == 1

    # Should be able to get a sensitive value from a sensitive series
    c4 = counts[4]
    assert isinstance(c4, pj.Prisoner)
    assert c4.max_distance == 1
    assert c4._value == 3

    # FIXME
    # # Should be able to get a sensitive view from a sensitive series
    # c3 = counts[1:3][3]
    # c4 = counts[1:3][4]
    # assert isinstance(c3, pj.Prisoner)
    # assert isinstance(c4, pj.Prisoner)
    # assert c3.max_distance == c4.max_distance == 1
    # assert c3._value == 1
    # assert c4._value == 3

def test_crosstab() -> None:
    pdf, df = load_dataframe()

    original_domain = copy.copy(pdf._domains)

    rowvalues = [1, 2, 3, 4, 5]
    colvalues = [1, 2, 3, 4, 5]

    # Should raise an error without rowvalues/colvalues
    with pytest.raises(pj.DPError): ppd.crosstab(pdf["a"], pdf["b"])
    pdf._domains = merge_domains(original_domain, {"a": ppd.CategoryDomain(categories=rowvalues)}) # FIXME: use public API
    with pytest.raises(pj.DPError): ppd.crosstab(pdf["a"], pdf["b"])
    pdf._domains = merge_domains(original_domain, {"b": ppd.CategoryDomain(categories=colvalues)}) # FIXME: use public API
    with pytest.raises(pj.DPError): ppd.crosstab(pdf["a"], pdf["b"])

    pdf._domains = merge_domains(original_domain,
                                 {"a": ppd.CategoryDomain(categories=rowvalues),
                                  "b": ppd.CategoryDomain(categories=colvalues)}) # FIXME: use public API

    # Should raise an error with margins=True
    with pytest.raises(pj.DPError):
        ppd.crosstab(pdf["a"], pdf["b"], margins=True)

    # Should raise an error with series of potentially different size
    pdf_ = pdf[pdf["a"] >= 0]
    with pytest.raises(pj.DPError):
        ppd.crosstab(pdf["a"], pdf_["b"])

    # Should return correct counts when all possible values are provided
    counts = ppd.crosstab(pdf["a"], pdf["b"])
    ans = pd.DataFrame([[0, 1, 0, 0, 0],
                        [0, 0, 0, 1, 0],
                        [0, 0, 0, 1, 0],
                        [0, 0, 0, 1, 0],
                        [0, 0, 1, 0, 0]],
                       index=rowvalues, columns=colvalues)
    assert_equal_sensitive_dataframes(counts, ans)

    pj.laplace_mechanism(counts, eps=1.0)

def test_cut() -> None:
    pdf, df = load_dataframe()

    # Default behaviour
    assert isinstance(ppd.cut(pdf["a"], bins=[0, 3, 6]), ppd.PrivSeries)
    assert (ppd.cut(pdf["a"], bins=[0, 3, 6])._value == pd.cut(df["a"], bins=[0, 3, 6])).all()

    # Should raise an error with a scalar bins
    with pytest.raises(TypeError):
        ppd.cut(pdf["a"], bins=2)

def test_dataframe_groupby() -> None:
    pdf, df = load_dataframe()

    # Should raise an error without `keys`
    with pytest.raises(pj.DPError):
        pdf.groupby("b")

    # Should be able to group by a single column
    keys = [1, 2, 4]
    # FIXME: use public API
    pdf._domains = merge_domains(pdf._domains, {"b": ppd.CategoryDomain(categories=list(keys))})
    pdf._value["b"] = pdf._value["b"].astype(pd.CategoricalDtype(categories=keys))
    groups = pdf.groupby("b")
    assert len(groups) == len(keys)

    # Should be able to loop over groups
    for i, (key, pdf_) in enumerate(groups):
        assert key == keys[i]
        assert isinstance(pdf_, ppd.PrivDataFrame)

    # A group with key=1 should be empty but its columns and dtypes should match the original
    assert isinstance(groups.get_group(1), ppd.PrivDataFrame)
    assert len(groups.get_group(1)._value) == 0
    assert (groups.get_group(1).columns == pdf.columns).all()
    assert (groups.get_group(1).dtypes == pdf.dtypes).all()

    # Check for a group with key=2
    assert isinstance(groups.get_group(2), ppd.PrivDataFrame)
    print(groups.get_group(2)._value)
    assert (groups.get_group(2)._value == pd.DataFrame({"a": [1], "b": [2]}, index=[0])).all().all()
    assert (groups.get_group(2).columns == pdf.columns).all()
    assert (groups.get_group(2).dtypes == pdf.dtypes).all()

    # Check for a group with key=4
    assert isinstance(groups.get_group(4), ppd.PrivDataFrame)
    assert (groups.get_group(4)._value == pd.DataFrame({"a": [2, 3, 4], "b": [4, 4, 4]}, index=[1, 2, 3])).all().all()
    assert (groups.get_group(4).columns == pdf.columns).all()
    assert (groups.get_group(4).dtypes == pdf.dtypes).all()

def test_groupby_accounting_group_isolation() -> None:
    pdf, df = load_dataframe()

    keys = [2, 4]
    pdf._domains = merge_domains(pdf._domains, {"b": ppd.CategoryDomain(categories=list(keys))})
    groups = pdf.groupby("b")

    child1 = groups.get_group(2)
    child2 = groups.get_group(4)

    # Each child should have a different accounting_group
    assert child1._accounting_group is not child2._accounting_group

    # Create a new accountant for child1
    new_acc = pj.create_accountant("approx", parent=child1.accountant)

    # Setting the new accountant on child1 should work
    child1.set_accountant(new_acc)
    assert child1.accountant is new_acc

    # Setting the new accountant on child2 should fail (different accounting_group)
    with pytest.raises(pj.DPError):
        child2.set_accountant(new_acc)

    # child -> parent direction should fail
    child_acc = pj.create_accountant("approx", parent=child2.accountant)
    with pytest.raises(pj.DPError):
        pdf.set_accountant(child_acc)

    # parent -> child direction should work (accounting_group unchanged)
    parent_acc = pj.create_accountant("approx", parent=pdf.accountant)
    original_group = child2._accounting_group
    child2.set_accountant(parent_acc)
    assert child2.accountant is parent_acc
    assert child2._accounting_group is original_group

def test_dataframe_values_returns_ndarray() -> None:
    pdf, df = load_dataframe()

    arr = pdf.values

    assert isinstance(arr, pnp.PrivNDArray)
    assert arr._value.tolist() == df.to_numpy(dtype=float).tolist()
    assert arr.max_distance == pytest.approx(pdf.max_distance)
    assert arr.domain.norm_bound is None

    pdf["a"] = pdf["a"].clip(0, 5)
    pdf["b"] = pdf["b"].clip(0, 5)
    clipped_arr = pdf.values

    assert isinstance(clipped_arr, pnp.PrivNDArray)
    assert clipped_arr._value.tolist() == df.to_numpy(dtype=float).tolist()
    assert clipped_arr.max_distance == pytest.approx(pdf.max_distance)
    assert clipped_arr.domain.norm_type == "l1"
    assert clipped_arr.domain.norm_bound == pytest.approx(10.0)

    pdf._domains = merge_domains(pdf._domains, {"a": ppd.CategoryDomain(categories=[1, 2, 3, 4, 5])}) # FIXME: use public API
    pdf._domains = merge_domains(pdf._domains, {"b": ppd.CategoryDomain(categories=[1, 2, 3, 4, 5])}) # FIXME: use public API

    counts = ppd.crosstab(pdf["a"], pdf["b"])
    counts_arr = counts.values

    assert isinstance(counts_arr, pnp.SensitiveNDArray)
    assert counts_arr.shape == (5, 5)
    assert counts_arr.max_distance == pytest.approx(1.0)
    assert counts_arr.norm_type == "l1"

    expected = [[0, 1, 0, 0, 0],
                [0, 0, 0, 1, 0],
                [0, 0, 0, 1, 0],
                [0, 0, 0, 1, 0],
                [0, 0, 1, 0, 0]]
    assert counts_arr._value.tolist() == expected

def test_series_values_returns_ndarray() -> None:
    pdf, df = load_dataframe()

    arr = pdf["a"].values

    assert isinstance(arr, pnp.PrivNDArray)
    assert arr._value.tolist() == df["a"].to_numpy(dtype=float).tolist()
    assert arr.max_distance == pytest.approx(pdf["a"].max_distance)
    assert arr.domain.norm_bound is None

    clipped_arr = pdf["a"].clip(0, 5).values

    assert isinstance(clipped_arr, pnp.PrivNDArray)
    assert clipped_arr._value.tolist() == df["a"].to_numpy(dtype=float).tolist()
    assert clipped_arr.max_distance == pytest.approx(pdf.max_distance)
    assert clipped_arr.domain.norm_type == "l1"
    assert clipped_arr.domain.norm_bound == pytest.approx(5.0)

    pdf._domains = merge_domains(pdf._domains, {"a": ppd.CategoryDomain(categories=[1, 2, 3, 4, 5])}) # FIXME: use public API

    counts = pdf["a"].value_counts(sort=False)
    counts_arr = counts.values

    assert isinstance(counts_arr, pnp.SensitiveNDArray)
    assert counts_arr.shape == (5,)
    assert counts_arr.max_distance == pytest.approx(1.0)
    assert counts_arr.norm_type == "l1"

def test_privacy_budget() -> None:
    pdf, df = load_dataframe()

    eps = 0.1
    pdf1 = pdf[pdf["b"] >= 3]
    pdf1._domains = merge_domains(pdf1._domains, {"b": ppd.CategoryDomain(categories=[3, 4, 5])}) # FIXME: use public API
    counts = pdf1["b"].value_counts(sort=False)

    pj.laplace_mechanism(counts, eps=eps)
    assert_budget_spent(pdf, eps)

    pj.laplace_mechanism(counts[4], eps=eps)
    assert_budget_spent(pdf, eps * 2)

    pdf2 = pdf[pdf["a"] >= 3]

    pj.laplace_mechanism(pdf2.shape[0], eps=eps)
    assert_budget_spent(pdf, eps * 3)

    # Privacy budget for different data sources should be managed independently
    pdf_, df_ = load_dataframe()

    pj.laplace_mechanism(pdf_.shape[0], eps=eps)
    assert_budget_spent(pdf, eps * 3)
    assert_budget_spent(pdf_, eps)

def test_privacy_budget_parallel_composition() -> None:
    pdf, df = load_dataframe()

    eps = 0.1

    # value_counts()
    pdf._domains = merge_domains(pdf._domains, {"b": ppd.CategoryDomain(categories=[2, 3, 4])}) # FIXME: use public API
    counts = pdf["b"].value_counts(sort=False)

    pj.laplace_mechanism(counts, eps=eps)
    assert_budget_spent(pdf, eps)

    pj.laplace_mechanism(counts[2], eps=eps)
    assert_budget_spent(pdf, eps * 2)

    pj.laplace_mechanism(counts[3], eps=eps)
    assert_budget_spent(pdf, eps * 2)

    # crosstab()
    pdf._domains = merge_domains(pdf._domains, {"a": ppd.CategoryDomain(categories=[1, 2, 3, 4, 5])}) # FIXME: use public API
    pdf._domains = merge_domains(pdf._domains, {"b": ppd.CategoryDomain(categories=[1, 2, 3, 4, 5])}) # FIXME: use public API
    crosstab = ppd.crosstab(pdf["a"], pdf["b"])

    for idx in crosstab.index:
        for col in crosstab.columns:
            pj.laplace_mechanism(crosstab[col][idx], eps=eps)

    assert_budget_spent(pdf, eps * 3)

    s = pj.SensitiveInt(0)
    for idx in crosstab.index:
        for col in crosstab.columns:
            s += crosstab[col][idx]
    assert s.max_distance == 1
    pj.laplace_mechanism(s, eps=eps)
    assert_budget_spent(pdf, eps * 4)

    assert crosstab[5][1].max_distance == 1
    pj.laplace_mechanism(crosstab[5][1], eps=eps)
    assert_budget_spent(pdf, eps * 5)

    # FIXME: use public API
    pdf._domains = merge_domains(pdf._domains, {"b": ppd.CategoryDomain(categories=[1, 2, 3, 4, 5])})
    pdf._value["b"] = pdf._value["b"].astype(pd.CategoricalDtype(categories=[1, 2, 3, 4, 5]))
    print(pdf._value)

    # groupby()
    s = pj.SensitiveInt(0)
    for key, pdf_ in pdf.groupby("b"):
        s += pdf_.shape[0]
        print(s._value)
    assert s.max_distance == 1
    assert s._value == len(pdf._value)
    pj.laplace_mechanism(s, eps=eps)
    assert_budget_spent(pdf, eps * 6)

    s = pj.SensitiveInt(0)
    for key, pdf_ in pdf.groupby("b"):
        s += key * pdf_.shape[0]
    assert s.max_distance == 5
    pj.laplace_mechanism(s, eps=eps)
    assert_budget_spent(pdf, eps * 7)
