# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: LicenseRef-.amazon.com.-AmznSL-1.0
# Licensed under the Amazon Software License  http://aws.amazon.com/asl/

from smclarify.util import pdf, pdfs_aligned_nonzero
import numpy as np


def test_pdf():
    xs = [1, 2, 2, 3, 3, 3]
    pdf_res = dict(pdf(xs))
    assert pdf_res[1] == 1 / 6
    assert pdf_res[2] == 2 / 6
    assert pdf_res[3] == 3 / 6


def test_pdfs_aligned():
    (pa, pb) = pdfs_aligned_nonzero([1, 2, 3], [4, 5, 5])
    assert np.array_equal(pa, [])
    assert np.array_equal(pb, [])

    (pa, pb) = pdfs_aligned_nonzero([1, 2, 3], [1, 1, 4, 5, 5])
    assert np.array_equal(pa, [1 / 3])
    assert np.array_equal(pb, [2 / 5])

    (pa, pb, pc) = pdfs_aligned_nonzero([1, 2, 3], [1, 1, 4, 5, 5], [1, 1, 1, 2])
    assert np.array_equal(pa, [1 / 3])
    assert np.array_equal(pb, [2 / 5])
    assert np.array_equal(pc, [3 / 4])