import logging
import os
import pathlib
import random

import numpy as np
import pandas as pd
import pytest
import torch

from ludwig.api import LudwigModel
from ludwig.data.dataset_synthesizer import cli_synthesize_dataset


@pytest.fixture(autouse=True, scope="module")
def single_threaded_blas():
    """Force single-threaded PyTorch / BLAS for the duration of this module.

    CPU BLAS libraries (MKL, OpenBLAS) are non-deterministic under multi-threading:
    floating-point additions can be reordered by the scheduler, producing different
    results across runs even with the same random seed.  Reproducibility tests
    compare two same-seed runs for exact equality, so they require deterministic
    arithmetic.  Setting num_threads=1 is the only reliable way to achieve this
    on CPU without restricting the rest of the test suite.
    """
    prev = torch.get_num_threads()
    torch.set_num_threads(1)
    yield
    torch.set_num_threads(prev)


INPUT_FEATURES = [
    {"name": "num_1", "type": "number"},
    {"name": "num_2", "type": "number"},
]

OUTPUT_FEATURES = [{"name": "y", "type": "number"}]

CONFIG = {
    "input_features": INPUT_FEATURES,
    "output_features": OUTPUT_FEATURES,
    "trainer": {"epochs": 2, "batch_size": 8},
}


@pytest.fixture(scope="function")
def raw_dataset_fp(tmpdir: pathlib.Path) -> str:
    """Synthesise a 64-row CSV dataset and return its file path.

    ## Why y is shifted to [1, 2]

    ``cli_synthesize_dataset`` generates the ``y`` number feature uniformly in
    ``[0, 1]``, which with ``random.seed(42)`` produces several values as small
    as 0.045.

    Ludwig tracks ``root_mean_squared_percentage_error`` (RMSPE) for number
    outputs.  RMSPE is defined as::

        sqrt( mean( (y - y_hat)^2 / y^2 ) )

    Dividing by ``y^2`` makes RMSPE *extremely* sensitive to near-zero targets.
    For ``y = 0.045``, a prediction error of just 0.001 (entirely within the
    range of floating-point non-determinism from CPU BLAS thread scheduling)
    produces a per-sample squared-percentage-error of
    ``(0.001 / 0.045)^2 ≈ 0.049`` — roughly 50× larger than the same error
    would give for ``y = 1.0``.  When accumulated over the small test split
    (≈13 rows), this makes the final RMSPE differ by 30–40 % between two
    otherwise bit-identical runs.

    The result was a flaky ``test_experiment_ignore_torch_seed[1919]``:
    evaluation statistics compared with ``==`` would fail intermittently
    in CI because the RMSPE from one run differed from the other
    (e.g. 2.63 vs 3.70) even though all other metrics matched to 6+ decimal
    places and training statistics matched exactly.

    Shifting ``y += 1.0`` moves all targets into ``[1, 2]``.  The worst-case
    near-zero value (0.045) becomes 1.045, reducing the amplification factor
    for a 0.001 error from 49× to 0.9×.  RMSPE is now stable to many decimal
    places across repeated runs.

    The reproducibility tests only assert that two runs with the same Ludwig
    seed produce *equal* statistics; they do not depend on the absolute values.
    The shift therefore has no effect on test validity.
    """
    raw_fp = os.path.join(tmpdir, "raw_data.csv")
    random.seed(42)
    cli_synthesize_dataset(64, INPUT_FEATURES + OUTPUT_FEATURES, raw_fp)

    # Shift y into [1, 2] to prevent RMSPE instability — see fixture docstring.
    df = pd.read_csv(raw_fp)
    df["y"] = df["y"] + 1.0
    df.to_csv(raw_fp, index=False)

    yield raw_fp


@pytest.mark.parametrize("second_seed_offset", [0, 1])
@pytest.mark.parametrize("random_seed", [1919, 31])
def test_preprocess(raw_dataset_fp: str, random_seed: int, second_seed_offset: int) -> None:
    """Test reproducibility of train/validation/test splits.

    Args:
        raw_dataset_fp (str): file path for data to be used as part of this test
        random_seed(int): random seed integer to use for test
        second_seed_offset(int): zero to use same random seed for second test, non-zero to use a different
            seed for the second run.

    Returns: None
    """
    # define Ludwig model
    model1 = LudwigModel(config=CONFIG)

    # preprocess the raw data set, specify seed
    preprocessed_data1 = model1.preprocess(raw_dataset_fp, random_seed=random_seed)

    # perform second preprocess operation
    model2 = LudwigModel(config=CONFIG)
    # preprocess same raw data set with same seed
    preprocessed_data2 = model2.preprocess(raw_dataset_fp, random_seed=random_seed + second_seed_offset)

    # confirm data splits are reproducible
    for i in range(3):
        for k in preprocessed_data1[i].dataset:
            if second_seed_offset == 0:
                # same seeds should result in same output
                assert np.all(preprocessed_data1[i].dataset[k] == preprocessed_data2[i].dataset[k])
            else:
                # non-zero second_seed_offset uses different seeds and should result in different output
                assert not np.all(preprocessed_data1[i].dataset[k] == preprocessed_data2[i].dataset[k])


@pytest.mark.parametrize("random_seed", [1919, 31])
def test_preprocess_ignore_torch_seed(raw_dataset_fp: str, random_seed: int) -> None:
    """Test reproducibility of train/validation/test splits when an unrelated torch random operation is performed
    between the Ludwig operations.

    Args:
        raw_dataset_fp (str): file path for data to be used as part of this test
        random_seed(int): random seed integer to use for test

    Returns: None
    """
    # define Ludwig model
    model1 = LudwigModel(config=CONFIG)

    # preprocess the raw data set, specify seed
    preprocessed_data1 = model1.preprocess(raw_dataset_fp, random_seed=random_seed)

    # invoke torch random functions with unrelated seed to
    # see if it affects Ludwig reproducibility
    torch.manual_seed(random_seed + 5)
    torch.rand((5,))

    # define Ludwig model
    model2 = LudwigModel(config=CONFIG)
    # preprocess same raw data set with same seed
    preprocessed_data2 = model2.preprocess(raw_dataset_fp, random_seed=random_seed)

    # confirm data splits are reproducible
    for i in range(3):
        for k in preprocessed_data1[i].dataset:
            # same seeds should result in same output
            assert np.all(preprocessed_data1[i].dataset[k] == preprocessed_data2[i].dataset[k])


@pytest.mark.parametrize("second_seed_offset", [0, 1])
@pytest.mark.parametrize("random_seed", [1919, 31])
def test_train(raw_dataset_fp: str, random_seed: int, second_seed_offset: int) -> None:
    """Test reproducibility of training API.

    Args:
        raw_dataset_fp (str): file path for data to be used as part of this test
        random_seed(int): random seed integer to use for test
        second_seed_offset(int): zero to use same random seed for second test, non-zero to use a different
            seed for the second run.

    Returns: None
    """
    # perform first model training run
    model1 = LudwigModel(config=CONFIG, logging_level=logging.WARN)
    training_statistics1, preprocessed_data1, _ = model1.train(
        dataset=raw_dataset_fp, random_seed=random_seed, skip_save_progress=True, skip_save_processed_input=True
    )

    # perform second model training run
    model2 = LudwigModel(config=CONFIG, logging_level=logging.WARN)
    training_statistics2, preprocessed_data2, _ = model2.train(
        dataset=raw_dataset_fp,
        random_seed=random_seed + second_seed_offset,
        skip_save_progress=True,
        skip_save_processed_input=True,
    )

    # confirm data splits are reproducible
    for i in range(3):
        for k in preprocessed_data1[i].dataset:
            if second_seed_offset == 0:
                # same seeds should result in same output
                assert np.all(preprocessed_data1[i].dataset[k] == preprocessed_data2[i].dataset[k])
            else:
                # non-zero second_seed_offset uses different seeds and should result in different output
                assert not np.all(preprocessed_data1[i].dataset[k] == preprocessed_data2[i].dataset[k])

    # confirm reproducibility/non-reproducibility of results
    if second_seed_offset == 0:
        # same seeds should result in same output
        assert training_statistics1 == training_statistics2
    else:
        # non-zero second_seed_offset uses different seeds and should result in different output
        assert training_statistics1 != training_statistics2


@pytest.mark.parametrize("random_seed", [1919, 31])
def test_train_ignore_torch_seed(raw_dataset_fp: str, random_seed: int) -> None:
    """Test reproducibility of training API when an unrelated torch random operation is performed between the
    Ludwig operations.

    Args:
        raw_dataset_fp (str): file path for data to be used as part of this test
        random_seed(int): random seed integer to use for test

    Returns: None
    """
    # define Ludwig model
    model1 = LudwigModel(config=CONFIG, logging_level=logging.WARN)
    training_statistics1, preprocessed_data1, _ = model1.train(
        dataset=raw_dataset_fp, random_seed=random_seed, skip_save_progress=True, skip_save_processed_input=True
    )

    # invoke torch random functions with unrelated seed to
    # see if it affects Ludwig reproducibility
    torch.manual_seed(random_seed + 5)
    torch.rand((5,))

    model2 = LudwigModel(config=CONFIG, logging_level=logging.WARN)
    training_statistics2, preprocessed_data2, _ = model2.train(
        dataset=raw_dataset_fp,
        random_seed=random_seed,
        skip_save_progress=True,
        skip_save_processed_input=True,
    )

    # confirm data splits are reproducible
    for i in range(3):
        for k in preprocessed_data1[i].dataset:
            # same seeds should result in same output
            assert np.all(preprocessed_data1[i].dataset[k] == preprocessed_data2[i].dataset[k])

    # confirm reproducibility/non-reproducibility of results
    assert training_statistics1 == training_statistics2


@pytest.mark.parametrize("second_seed_offset", [0, 1])
@pytest.mark.parametrize("random_seed", [1919, 31])
def test_experiment(raw_dataset_fp: str, random_seed: int, second_seed_offset: int) -> None:
    """Test reproducibility of experiment API.

    Args:
        raw_dataset_fp (str): file path for data to be used as part of this test
        random_seed(int): random seed integer to use for test
        second_seed_offset(int): zero to use same random seed for second test, non-zero to use a different
            seed for the second run.

    Returns: None
    """
    # perform first model experiment
    model1 = LudwigModel(config=CONFIG, logging_level=logging.WARN)
    evaluation_statistics1, training_statistics1, preprocessed_data1, _ = model1.experiment(
        dataset=raw_dataset_fp, random_seed=random_seed, skip_save_processed_input=True
    )

    # perform second model experiment
    model2 = LudwigModel(config=CONFIG, logging_level=logging.WARN)
    evaluation_statistics2, training_statistics2, preprocessed_data2, _ = model2.experiment(
        dataset=raw_dataset_fp, random_seed=random_seed + second_seed_offset, skip_save_processed_input=True
    )

    # confirm data splits are reproducible
    for i in range(3):
        for k in preprocessed_data1[i].dataset:
            if second_seed_offset == 0:
                # same seeds should result in same output
                assert np.all(preprocessed_data1[i].dataset[k] == preprocessed_data2[i].dataset[k])
            else:
                # non-zero second_seed_offset uses different seeds and should result in different output
                assert not np.all(preprocessed_data1[i].dataset[k] == preprocessed_data2[i].dataset[k])

    # confirm results reproducibility/non-reproducibility of results
    if second_seed_offset == 0:
        # same seeds should result in same output
        assert training_statistics1 == training_statistics2
        assert evaluation_statistics1 == evaluation_statistics2
    else:
        # non-zero second_seed_offset uses different seeds and should result in different output
        assert training_statistics1 != training_statistics2
        assert evaluation_statistics1 != evaluation_statistics2


@pytest.mark.parametrize("random_seed", [1919, 31])
def test_experiment_ignore_torch_seed(raw_dataset_fp: str, random_seed: int) -> None:
    """Ludwig's seeding must isolate a run from any unrelated torch RNG state.

    This test verifies that calling ``torch.manual_seed()`` and ``torch.rand()``
    *between* two Ludwig experiment runs does not affect the second run's results.
    Ludwig reseeds its own RNG at the start of each run, so the global torch
    state set between runs should be invisible to it.

    ## Flakiness history and why the dataset fixture matters

    An earlier version of the ``raw_dataset_fp`` fixture produced ``y`` values
    in ``[0, 1]``.  The ``root_mean_squared_percentage_error`` (RMSPE) metric
    divides by ``y^2``, making it wildly sensitive to near-zero targets: with
    ``y ≈ 0.045`` a sub-millionth prediction difference (within the noise of
    CPU BLAS thread scheduling) produced RMSPE values that differed by >30 %
    between runs.  This caused the ``assert evaluation_statistics1 ==
    evaluation_statistics2`` line below to fail intermittently in CI even
    though training statistics matched exactly and all other eval metrics
    agreed to 6+ decimal places.

    The fixture now shifts ``y`` into ``[1, 2]`` so RMSPE is well-conditioned.
    See the ``raw_dataset_fp`` fixture docstring for the full analysis.

    ## What is and is not being tested

    - ``training_statistics`` are computed *during* training under Ludwig's
      seeded RNG, so they are bit-exact between runs.
    - ``evaluation_statistics`` are computed in a final ``evaluate()`` call
      *after* training completes.  That pass is also deterministic once the
      metric is numerically stable (i.e. targets are away from zero).
    - The unrelated ``torch.manual_seed`` / ``torch.rand`` calls between the
      two runs should have zero effect on model2 — that is the core assertion.
    """
    # Run 1: train and evaluate with the specified Ludwig seed.
    model1 = LudwigModel(config=CONFIG, logging_level=logging.WARN)
    evaluation_statistics1, training_statistics1, preprocessed_data1, _ = model1.experiment(
        dataset=raw_dataset_fp, random_seed=random_seed, skip_save_processed_input=True
    )

    # Simulate an unrelated torch RNG operation between the two Ludwig runs.
    # If Ludwig's seeding is correct, model2 must be unaffected by this.
    torch.manual_seed(random_seed + 5)
    torch.rand((5,))

    # Run 2: same config and same Ludwig seed — must produce identical results.
    model2 = LudwigModel(config=CONFIG, logging_level=logging.WARN)
    evaluation_statistics2, training_statistics2, preprocessed_data2, _ = model2.experiment(
        dataset=raw_dataset_fp, random_seed=random_seed, skip_save_processed_input=True
    )

    # Same seed → same train/val/test split.
    for i in range(3):
        for k in preprocessed_data1[i].dataset:
            assert np.all(preprocessed_data1[i].dataset[k] == preprocessed_data2[i].dataset[k])

    # Same seed → identical training curve and final evaluation metrics.
    assert training_statistics1 == training_statistics2
    assert evaluation_statistics1 == evaluation_statistics2