# Path: idtxl/estimators_multivariate_pid.py
# class SxPID(Estimator):
#     """Estimate partial information decomposition for multiple inputs.
# 
#     Implementation of the multivariate partial information decomposition (PID)
#     estimator for discrete data with (up to 4 inputs) and one output. The
#     estimator finds shared information, unique information and synergistic
#     information between the multiple inputs s1, s2, ..., sn with respect to the
#     output t for each realization (t, s1, ..., sn) and then average them
#     according to their distribution weights p(t, s1, ..., sn). Both the
#     pointwise (on the realization level) PID and the averaged PID are returned
#     (see the 'return' of 'estimate()').
# 
#     The algorithm uses recursion to compute the partial information
#     decomposition.
# 
#     References:
# 
#     - Makkeh, A. & Wibral, M. (2020). A differentiable pointwise partial
#       Information Decomposition estimator. https://github.com/Abzinger/SxPID.
# 
#     Args:
#         settings : dict
#             estimation parameters (with default parameters)
# 
#             - verbose : bool [optional] - print output to console
#               (default=False)
#     """
# 
#     def __init__(self, settings):
#         # get estimation parameters
#         self.settings = settings.copy()
#         self.settings.setdefault('verbose', False)
# 
#     def is_parallel():
#         return False
# 
#     def is_analytic_null_estimator(self):
#         return False
# 
#     def estimate(self, s, t):
#         """
#         Args:
#             s : list of numpy arrays
#                 1D arrays containing realizations of a discrete random variable
#             t : numpy array
#                 1D array containing realizations of a discrete random variable
# 
#         Returns:
#             dict of dict
#                 {
#                  'ptw' -> { realization -> {alpha -> [float, float, float]} }
# 
#                  'avg' -> {alpha -> [float, float, float]}
#                 }
#             where the list of floats is ordered
#             [informative, misinformative, informative - misinformative]
#             ptw stands for pointwise decomposition
#             avg stands for average decomposition
#         """
#         s, t, self.settings = _check_input(s, t, self.settings)
#         pdf = _get_pdf_dict(s, t)
# 
#         # Read lattices from a file
#         # Stored as {
#         #             n -> [{alpha -> children}, (alpha_1,...) ]
#         #           }
#         # children is a list of tuples
#         lattices = lt.lattices
#         num_source_vars = len(s)
#         retval_ptw, retval_avg = pid_goettingen.pid(
#             num_source_vars,
#             pdf_orig=pdf,
#             chld=lattices[num_source_vars][0],
#             achain=lattices[num_source_vars][1],
#             printing=self.settings['verbose'])
# 
#         # TODO AskM: Trivariate: does it make sense to name the alphas
#         #    for example shared_syn_s1_s2__syn_s1_s3 ?
#         results = {
#             'ptw': retval_ptw,
#             'avg': retval_avg,
#         }
#         return results


import time as tm
import numpy as np
import pytest
from idtxl.estimators_multivariate_pid import SxPID
    t_goettingen = tm.time() - tic

    print('\nInt addition - N = 1000')
    print('Estimator            SxPID\n')
    print('PID evaluation       {:.3f} s\n'.format(t_goettingen))
    print('Uni s1               {0:.8f}'.format(
                                         est_goettingen['avg'][((1,),)][2]))
    print('Uni s2               {0:.8f}'.format(
                                         est_goettingen['avg'][((2,),)][2]))
    print('Shared s1_s2         {0:.8f}'.format(
                                         est_goettingen['avg'][((1,),(2,),)][2]))
    print('Synergy s1_s2        {0:.8f}'.format(
                                         est_goettingen['avg'][((1,2,),)][2]))

    # assert np.isclose(est_tartu['syn_s1_s2'], est_sydney['syn_s1_s2'],
    #                   atol=1e-03), 'Synergies are not equal.'
    # assert np.isclose(est_tartu['shd_s1_s2'], est_sydney['shd_s1_s2'],
    #                   atol=1e-03), 'Shareds are not equal.'
    # assert np.isclose(est_tartu['unq_s1'], est_sydney['unq_s1'],
    #                   atol=1e-03), 'Unique1 is not equal.'
    # assert np.isclose(est_tartu['unq_s2'], est_sydney['unq_s2'],
    #                   atol=1e-03), 'Unique2 is not equal.'


def test_average_pid_three_hash():
    """Test Goettingen estimator on binary three hash."""
    s1 = np.asarray([0, 0, 0, 0, 1, 1, 1, 1])
    s2 = np.asarray([0, 0, 1, 1, 0, 0, 1, 1])
    s3 = np.asarray([0, 1, 0, 1, 0, 1, 0, 1])
    target = np.logical_xor(s3, np.logical_xor(s1, s2).astype(int)).astype(int)

# Path: idtxl/estimators_multivariate_pid.py
# class SxPID(Estimator):
#     """Estimate partial information decomposition for multiple inputs.
# 
#     Implementation of the multivariate partial information decomposition (PID)
#     estimator for discrete data with (up to 4 inputs) and one output. The
#     estimator finds shared information, unique information and synergistic
#     information between the multiple inputs s1, s2, ..., sn with respect to the
#     output t for each realization (t, s1, ..., sn) and then average them
#     according to their distribution weights p(t, s1, ..., sn). Both the
#     pointwise (on the realization level) PID and the averaged PID are returned
#     (see the 'return' of 'estimate()').
# 
#     The algorithm uses recursion to compute the partial information
#     decomposition.
# 
#     References:
# 
#     - Makkeh, A. & Wibral, M. (2020). A differentiable pointwise partial
#       Information Decomposition estimator. https://github.com/Abzinger/SxPID.
# 
#     Args:
#         settings : dict
#             estimation parameters (with default parameters)
# 
#             - verbose : bool [optional] - print output to console
#               (default=False)
#     """
# 
#     def __init__(self, settings):
#         # get estimation parameters
#         self.settings = settings.copy()
#         self.settings.setdefault('verbose', False)
# 
#     def is_parallel():
#         return False
# 
#     def is_analytic_null_estimator(self):
#         return False
# 
#     def estimate(self, s, t):
#         """
#         Args:
#             s : list of numpy arrays
#                 1D arrays containing realizations of a discrete random variable
#             t : numpy array
#                 1D array containing realizations of a discrete random variable
# 
#         Returns:
#             dict of dict
#                 {
#                  'ptw' -> { realization -> {alpha -> [float, float, float]} }
# 
#                  'avg' -> {alpha -> [float, float, float]}
#                 }
#             where the list of floats is ordered
#             [informative, misinformative, informative - misinformative]
#             ptw stands for pointwise decomposition
#             avg stands for average decomposition
#         """
#         s, t, self.settings = _check_input(s, t, self.settings)
#         pdf = _get_pdf_dict(s, t)
# 
#         # Read lattices from a file
#         # Stored as {
#         #             n -> [{alpha -> children}, (alpha_1,...) ]
#         #           }
#         # children is a list of tuples
#         lattices = lt.lattices
#         num_source_vars = len(s)
#         retval_ptw, retval_avg = pid_goettingen.pid(
#             num_source_vars,
#             pdf_orig=pdf,
#             chld=lattices[num_source_vars][0],
#             achain=lattices[num_source_vars][1],
#             printing=self.settings['verbose'])
# 
#         # TODO AskM: Trivariate: does it make sense to name the alphas
#         #    for example shared_syn_s1_s2__syn_s1_s3 ?
#         results = {
#             'ptw': retval_ptw,
#             'avg': retval_avg,
#         }
#         return results


# Path: test/test_estimators_multivariate_pid.py
import time as tm
import numpy as np
import pytest
from idtxl.estimators_multivariate_pid import SxPID

    t_goettingen = tm.time() - tic

    print('\nInt addition - N = 1000')
    print('Estimator            SxPID\n')
    print('PID evaluation       {:.3f} s\n'.format(t_goettingen))
    print('Uni s1               {0:.8f}'.format(
                                         est_goettingen['avg'][((1,),)][2]))
    print('Uni s2               {0:.8f}'.format(
                                         est_goettingen['avg'][((2,),)][2]))
    print('Shared s1_s2         {0:.8f}'.format(
                                         est_goettingen['avg'][((1,),(2,),)][2]))
    print('Synergy s1_s2        {0:.8f}'.format(
                                         est_goettingen['avg'][((1,2,),)][2]))

    # assert np.isclose(est_tartu['syn_s1_s2'], est_sydney['syn_s1_s2'],
    #                   atol=1e-03), 'Synergies are not equal.'
    # assert np.isclose(est_tartu['shd_s1_s2'], est_sydney['shd_s1_s2'],
    #                   atol=1e-03), 'Shareds are not equal.'
    # assert np.isclose(est_tartu['unq_s1'], est_sydney['unq_s1'],
    #                   atol=1e-03), 'Unique1 is not equal.'
    # assert np.isclose(est_tartu['unq_s2'], est_sydney['unq_s2'],
    #                   atol=1e-03), 'Unique2 is not equal.'


def test_average_pid_three_hash():
    """Test Goettingen estimator on binary three hash."""
    s1 = np.asarray([0, 0, 0, 0, 1, 1, 1, 1])
    s2 = np.asarray([0, 0, 1, 1, 0, 0, 1, 1])
    s3 = np.asarray([0, 1, 0, 1, 0, 1, 0, 1])
    target = np.logical_xor(s3, np.logical_xor(s1, s2).astype(int)).astype(int)

    pid_goettingen = SxPID(SETTINGS)