import pytest


@pytest.fixture()
def batch_size():
    return 8


@pytest.fixture()
def num_samples():
    return 100


@pytest.fixture()
def summary_network():
    return None


@pytest.fixture()
def adapter_without_sample_weight():
    from bayesflow import ContinuousApproximator

    return ContinuousApproximator.build_adapter(
        inference_variables=["mean", "std"],
        inference_conditions=["x"],
    )


@pytest.fixture()
def adapter_with_sample_weight():
    from bayesflow import ContinuousApproximator

    return ContinuousApproximator.build_adapter(
        inference_variables=["mean", "std"],
        inference_conditions=["x"],
        sample_weight="weight",
    )


@pytest.fixture()
def adapter_unconditional():
    from bayesflow import ContinuousApproximator

    return ContinuousApproximator.build_adapter(
        inference_variables=["mean", "std"],
    )


@pytest.fixture(params=["adapter_unconditional", "adapter_without_sample_weight", "adapter_with_sample_weight"])
def adapter(request):
    return request.getfixturevalue(request.param)


@pytest.fixture()
def normal_simulator():
    from tests.utils.normal_simulator import NormalSimulator

    return NormalSimulator()


@pytest.fixture()
def normal_simulator_with_sample_weight():
    from tests.utils.normal_simulator import NormalSimulator
    from bayesflow import make_simulator

    def weight(mean):
        return dict(weight=1.0)

    return make_simulator([NormalSimulator(), weight])


@pytest.fixture(params=["normal_simulator", "normal_simulator_with_sample_weight"])
def simulator(request):
    return request.getfixturevalue(request.param)


@pytest.fixture()
def train_dataset(batch_size, adapter, simulator):
    from bayesflow import OfflineDataset
    from tests.utils import check_combination_simulator_adapter

    check_combination_simulator_adapter(simulator, adapter)

    num_batches = 4
    data = simulator.sample((num_batches * batch_size,))
    return OfflineDataset(data=data, adapter=adapter, batch_size=batch_size, workers=4, max_queue_size=num_batches)


@pytest.fixture()
def validation_dataset(batch_size, adapter, simulator):
    from bayesflow import OfflineDataset
    from tests.utils import check_combination_simulator_adapter

    check_combination_simulator_adapter(simulator, adapter)

    num_batches = 2
    data = simulator.sample((num_batches * batch_size,))
    return OfflineDataset(data=data, adapter=adapter, batch_size=batch_size, workers=4, max_queue_size=num_batches)


@pytest.fixture()
def mean_std_summary_network():
    from tests.utils import MeanStdSummaryNetwork

    return MeanStdSummaryNetwork()


@pytest.fixture()
def continuous_approximator(adapter, summary_network):
    from bayesflow import ContinuousApproximator
    from bayesflow.networks import CouplingFlow

    return ContinuousApproximator(
        adapter=adapter,
        inference_network=CouplingFlow(subnet="mlp", depth=2, subnet_kwargs=dict(widths=(32, 32))),
        summary_network=summary_network,
    )


@pytest.fixture()
def point_approximator_without_parametric_score(adapter, summary_network):
    from bayesflow import ScoringRuleApproximator
    from bayesflow.networks import PointNetwork

    if "-> 'inference_conditions'" not in str(adapter) and "-> 'summary_conditions'" not in str(adapter):
        pytest.skip("Scoring rule approximator does not support unconditional estimation")

    return ScoringRuleApproximator(
        adapter=adapter,
        inference_network=PointNetwork(
            points="mean",
            subnet="mlp",
            subnet_kwargs=dict(widths=(8, 8)),
        ),
        summary_network=summary_network,
    )


@pytest.fixture()
def ensemble_approximator_continuous_and_point(continuous_approximator, point_approximator_without_parametric_score):
    from bayesflow import EnsembleApproximator

    return EnsembleApproximator(
        dict(cont_approx=continuous_approximator, point_approx=point_approximator_without_parametric_score)
    )


@pytest.fixture(
    params=["continuous_approximator", "ensemble_approximator_continuous_and_point"],
)
def approximator(request):
    return request.getfixturevalue(request.param)


@pytest.fixture(params=["continuous_approximator", "scoring_rule_approximator", "model_comparison_approximator"])
def approximator_with_summaries(request):
    from bayesflow.adapters import Adapter
    from bayesflow.networks import MLP

    adapter = Adapter()
    match request.param:
        case "continuous_approximator":
            from bayesflow.approximators import ContinuousApproximator
            from bayesflow.networks import CouplingFlow

            return ContinuousApproximator(adapter=adapter, inference_network=CouplingFlow(), summary_network=None)
        case "scoring_rule_approximator":
            from bayesflow.approximators import ScoringRuleApproximator
            from bayesflow.networks import PointNetwork

            return ScoringRuleApproximator(
                adapter=adapter,
                inference_network=PointNetwork(points=["mean"]),
                summary_network=None,
            )
        case "model_comparison_approximator":
            from bayesflow.approximators import ModelComparisonApproximator

            return ModelComparisonApproximator(
                num_models=2, classifier_network=MLP(widths=(8, 8)), adapter=adapter, summary_network=None
            )
        case _:
            raise ValueError("Invalid param for approximator class.")
