import random

from neuralpp.inference.graphical_model.learn.graphical_model_sgd_learner import (
    GraphicalModelSGDLearner,
)
from neuralpp.inference.graphical_model.representation.factor.fixed.fixed_pytorch_factor import (
    FixedPyTorchTableFactor,
)
from neuralpp.inference.graphical_model.representation.factor.neural.MLPFactor import (
    MLPFactor,
)
from neuralpp.inference.graphical_model.representation.model.model import (
    compute_query_probability,
    cross_entropy_for_dataset,
)
from neuralpp.inference.graphical_model.representation.random.random_dataset import (
    generate_dataset_given_observation_and_query_variables,
)
from neuralpp.inference.graphical_model.variable.integer_variable import IntegerVariable
from neuralpp.util import util
from neuralpp.util.util import try_noisy_test_up_to_n_times


def joint_learning():

    x = IntegerVariable("x", 2)
    y = IntegerVariable("y", 2)
    z = IntegerVariable("z", 2)

    y_is_negation_of_x = FixedPyTorchTableFactor.from_function(
        [x, y], lambda vx, vy: float(vy == 1 - vx)
    )
    noise = 0.9
    z_noisy_equals_y = FixedPyTorchTableFactor.from_function(
        [y, z], lambda vy, vz: noise if vz == vy else (1 - noise)
    )

    ground_truth_model = [y_is_negation_of_x, z_noisy_equals_y]

    dataset = generate_dataset_given_observation_and_query_variables(
        ground_truth_model,
        observed_variables=[x],
        query_variables=[y, z],
        number_of_observations=100,
        datapoints_per_observation=1,
    )

    print("Dataset sample:")
    print(util.join(random.sample(dataset, 20), "\n"))

    y_from_x = MLPFactor.make([x], y)
    learned_model = [y_from_x, z_noisy_equals_y]
    ground_truth_cross_entropy = cross_entropy_for_dataset(
        dataset, ground_truth_model
    ).item()

    print(f"Ground truth model cross entropy: {ground_truth_cross_entropy :.3f}")

    GraphicalModelSGDLearner(learned_model, dataset, loss_decrease_tol=1e-3).learn()

    cross_entropy = cross_entropy_for_dataset(dataset, learned_model).item()

    print(f"Ground truth model cross entropy: {ground_truth_cross_entropy :.3f}")
    print(f"Learned model cross entropy     : {cross_entropy :.3f}")

    for observation_dict, query_dict in random.sample(dataset, 10):
        likelihood = compute_query_probability(
            observation_dict, query_dict, learned_model
        )
        print(
            f"Given {observation_dict}, answer is {query_dict}, learned model assigns it probability {likelihood.item():.2f}"
        )

    expected_cross_entropy_upper_bound = (
        ground_truth_cross_entropy * 1.2
    )  # based on a few runs and observing that that's the usual approximation
    assert (
        cross_entropy <= expected_cross_entropy_upper_bound
    ), f"Learning performance was poorer than expected. Expected upper bound: {expected_cross_entropy_upper_bound:0.3f}, actual: {cross_entropy:.3f}"


def test_joint_learning():
    try_noisy_test_up_to_n_times(joint_learning, n=3)
