# Copyright (c) 2023 Predibase, Inc., 2019 Uber Technologies, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import os
import shutil
from copy import deepcopy

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

try:
    import torchtext
except ImportError:
    torchtext = None

from ludwig.api import LudwigModel
from ludwig.backend import RAY
from ludwig.constants import BATCH_SIZE, COMBINER, EVAL_BATCH_SIZE, LOGITS, NAME, PREDICTIONS, PROBABILITIES, TRAINER
from ludwig.data.preprocessing import preprocess_for_prediction
from ludwig.features.number_feature import numeric_transformation_registry
from ludwig.globals import TRAIN_SET_METADATA_FILE_NAME
from ludwig.models.inference import to_inference_module_input_from_dataframe
from ludwig.utils import output_feature_utils
from ludwig.utils.tokenizers import TORCHSCRIPT_COMPATIBLE_TOKENIZERS
from tests.integration_tests import utils
from tests.integration_tests.utils import (
    audio_feature,
    bag_feature,
    binary_feature,
    category_feature,
    date_feature,
    generate_data,
    h3_feature,
    image_feature,
    LocalTestBackend,
    number_feature,
    sequence_feature,
    set_feature,
    text_feature,
    timeseries_feature,
    vector_feature,
)


@pytest.mark.integration_tests_e
@pytest.mark.parametrize("should_load_model", [True, False])
def test_torchscript(tmpdir, csv_filename, should_load_model):
    #######
    # Setup
    #######
    dir_path = tmpdir
    data_csv_path = os.path.join(tmpdir, csv_filename)

    # Single sequence input, single category output
    image_dest_folder = os.path.join(tmpdir, "generated_images")
    audio_dest_folder = os.path.join(tmpdir, "generated_audio")
    input_features = [
        binary_feature(),
        number_feature(),
        category_feature(encoder={"type": "passthrough", "vocab_size": 3}),
        category_feature(encoder={"type": "onehot", "vocab_size": 3}),
        category_feature(encoder={"type": "dense", "vocab_size": 3}),
        sequence_feature(encoder={"vocab_size": 3}),
        text_feature(encoder={"vocab_size": 3}),
        vector_feature(),
        image_feature(image_dest_folder),
        audio_feature(audio_dest_folder),
        timeseries_feature(),
        date_feature(),
        date_feature(),
        h3_feature(),
        set_feature(encoder={"vocab_size": 3}),
        bag_feature(encoder={"vocab_size": 3}),
    ]

    output_features = [
        category_feature(decoder={"vocab_size": 3}),
        binary_feature(),
        number_feature(),
        set_feature(decoder={"vocab_size": 3}),
        vector_feature(),
        sequence_feature(decoder={"vocab_size": 3}),
        text_feature(decoder={"vocab_size": 3}),
    ]

    predictions_column_name = "{}_predictions".format(output_features[0]["name"])

    # Generate test data
    data_csv_path = generate_data(input_features, output_features, data_csv_path)

    #############
    # Train model
    #############
    backend = LocalTestBackend()
    config = {
        "model_type": "ecd",
        "input_features": input_features,
        "output_features": output_features,
        TRAINER: {"train_steps": 1},
    }
    ludwig_model = LudwigModel(config, backend=backend)
    ludwig_model.train(
        dataset=data_csv_path,
        skip_save_training_description=True,
        skip_save_training_statistics=True,
        skip_save_model=True,
        skip_save_progress=True,
        skip_save_log=True,
        skip_save_processed_input=True,
    )

    ###################
    # save Ludwig model
    ###################
    ludwigmodel_path = os.path.join(dir_path, "ludwigmodel")
    shutil.rmtree(ludwigmodel_path, ignore_errors=True)
    ludwig_model.save(ludwigmodel_path)

    ###################
    # load Ludwig model
    ###################
    if should_load_model:
        ludwig_model = LudwigModel.load(ludwigmodel_path, backend=backend)

    ##############################
    # collect weight tensors names
    ##############################
    original_predictions_df, _ = ludwig_model.predict(dataset=data_csv_path)
    original_weights = deepcopy(list(ludwig_model.model.parameters()))
    original_weights = [t.cpu() for t in original_weights]

    #################
    # save torchscript
    #################
    torchscript_path = os.path.join(dir_path, "torchscript")
    shutil.rmtree(torchscript_path, ignore_errors=True)
    ludwig_model.model.save_torchscript(torchscript_path, device="cpu")

    ###################################################
    # load Ludwig model, obtain predictions and weights
    ###################################################
    ludwig_model = LudwigModel.load(ludwigmodel_path, backend=backend)
    loaded_prediction_df, _ = ludwig_model.predict(dataset=data_csv_path)
    loaded_weights = deepcopy(list(ludwig_model.model.parameters()))
    loaded_weights = [t.cpu() for t in loaded_weights]

    #####################################################
    # restore torchscript, obtain predictions and weights
    #####################################################
    training_set_metadata_json_fp = os.path.join(ludwigmodel_path, TRAIN_SET_METADATA_FILE_NAME)

    dataset, training_set_metadata = preprocess_for_prediction(
        ludwig_model.config_obj.to_dict(),
        dataset=data_csv_path,
        training_set_metadata=training_set_metadata_json_fp,
        include_outputs=False,
        backend=backend,
    )

    restored_model = torch.jit.load(torchscript_path)

    # Check the outputs for one of the features for correctness
    # Here we choose the first output feature (categorical)
    of_name = list(ludwig_model.model.output_features.keys())[0]

    data_to_predict = {
        name: torch.from_numpy(dataset.dataset[feature.proc_column])
        for name, feature in ludwig_model.model.input_features.items()
    }

    # Get predictions from restored torchscript.
    logits = restored_model(data_to_predict)
    restored_predictions = torch.argmax(output_feature_utils.get_output_feature_tensor(logits, of_name, "logits"), -1)

    restored_predictions = [training_set_metadata[of_name]["idx2str"][idx] for idx in restored_predictions]

    restored_weights = deepcopy(list(restored_model.parameters()))
    restored_weights = [t.cpu() for t in restored_weights]

    ###############################################
    # Check if weights and predictions are the same
    ###############################################

    # Check to weight values match the original model.
    assert utils.is_all_close(original_weights, loaded_weights)
    assert utils.is_all_close(original_weights, restored_weights)

    # Check that predictions are identical to the original model.
    assert np.all(original_predictions_df[predictions_column_name] == loaded_prediction_df[predictions_column_name])

    assert np.all(original_predictions_df[predictions_column_name] == restored_predictions)


@pytest.mark.integration_tests_e
def test_torchscript_e2e_tabular(csv_filename, tmpdir):
    data_csv_path = os.path.join(tmpdir, csv_filename)
    # Configure features to be tested:
    bin_str_feature_input_feature = binary_feature()
    bin_str_feature_output_feature = binary_feature(output_feature=True)
    transformed_number_features = [
        number_feature(preprocessing={"normalization": numeric_transformer})
        for numeric_transformer in numeric_transformation_registry.keys()
    ]
    input_features = [
        bin_str_feature_input_feature,
        binary_feature(),
        *transformed_number_features,
        number_feature(preprocessing={"outlier_strategy": "fill_with_mean"}),
        category_feature(encoder={"vocab_size": 3}),
        bag_feature(encoder={"vocab_size": 3}),
        set_feature(encoder={"vocab_size": 3}),
        vector_feature(),
        # TODO: future support
        # date_feature(),
        # h3_feature(),
    ]
    output_features = [
        bin_str_feature_output_feature,
        binary_feature(output_feature=True),
        number_feature(),
        category_feature(decoder={"vocab_size": 3}),
        set_feature(decoder={"vocab_size": 3}),
        vector_feature(),
        sequence_feature(decoder={"vocab_size": 3}),
        text_feature(decoder={"vocab_size": 3}),
    ]
    backend = LocalTestBackend()
    config = {
        "input_features": input_features,
        "output_features": output_features,
        TRAINER: {"train_steps": 1, BATCH_SIZE: 128},
    }

    # Generate training data
    training_data_csv_path = generate_data(input_features, output_features, data_csv_path)

    # Convert bool values to strings, e.g., {'Yes', 'No'}
    df = pd.read_csv(training_data_csv_path)
    false_value, true_value = "No", "Yes"
    df[bin_str_feature_input_feature[NAME]] = df[bin_str_feature_input_feature[NAME]].map(
        lambda x: true_value if x else false_value
    )
    df[bin_str_feature_output_feature[NAME]] = df[bin_str_feature_output_feature[NAME]].map(
        lambda x: true_value if x else false_value
    )
    df.to_csv(training_data_csv_path)

    validate_torchscript_outputs(tmpdir, config, backend, training_data_csv_path)


@pytest.mark.integration_tests_e
def test_torchscript_e2e_binary_only(csv_filename, tmpdir):
    data_csv_path = os.path.join(tmpdir, csv_filename)

    input_features = [
        binary_feature(),
    ]
    output_features = [
        binary_feature(),
    ]
    backend = LocalTestBackend()
    config = {
        "input_features": input_features,
        "output_features": output_features,
        TRAINER: {"train_steps": 1, BATCH_SIZE: 128},
    }

    # Generate training data
    training_data_csv_path = generate_data(input_features, output_features, data_csv_path)

    validate_torchscript_outputs(tmpdir, config, backend, training_data_csv_path)


@pytest.mark.integration_tests_e
def test_torchscript_e2e_tabnet_combiner(csv_filename, tmpdir):
    data_csv_path = os.path.join(tmpdir, csv_filename)
    # Configure features to be tested:
    input_features = [
        binary_feature(),
        number_feature(),
        category_feature(encoder={"vocab_size": 3}),
        bag_feature(encoder={"vocab_size": 3}),
        set_feature(encoder={"vocab_size": 3}),
    ]
    output_features = [
        binary_feature(),
        number_feature(),
        category_feature(decoder={"vocab_size": 3}),
    ]
    backend = LocalTestBackend()
    config = {
        "input_features": input_features,
        "output_features": output_features,
        COMBINER: {
            "type": "tabnet",
            "num_total_blocks": 2,
            "num_shared_blocks": 2,
        },
        TRAINER: {"train_steps": 1, BATCH_SIZE: 128},
    }

    # Generate training data
    training_data_csv_path = generate_data(input_features, output_features, data_csv_path)

    validate_torchscript_outputs(tmpdir, config, backend, training_data_csv_path)


@pytest.mark.integration_tests_e
@pytest.mark.xfail(reason="torchaudio 2.x: DifferentiableFIR not TorchScript-compatible (upstream)")
def test_torchscript_e2e_audio(csv_filename, tmpdir):
    data_csv_path = os.path.join(tmpdir, csv_filename)
    audio_dest_folder = os.path.join(tmpdir, "generated_audio")

    input_features = [
        audio_feature(audio_dest_folder),
    ]
    output_features = [
        binary_feature(),
    ]
    backend = LocalTestBackend()
    config = {
        "input_features": input_features,
        "output_features": output_features,
        TRAINER: {"train_steps": 1, BATCH_SIZE: 128},
    }
    training_data_csv_path = generate_data(input_features, output_features, data_csv_path)

    # NOTE: audio preprocessing mismatches by very small margins ~O(1e-6) but causes flakiness in e2e test.
    # Increasing tolerance is a workaround to reduce flakiness for now.
    # TODO: remove this workaround when audio preprocessing is fixed.
    validate_torchscript_outputs(tmpdir, config, backend, training_data_csv_path, tolerance=1e-6)


@pytest.mark.integration_tests_e
@pytest.mark.parametrize(
    "kwargs",
    [
        {"encoder": {"type": "stacked_cnn"}},  # Ludwig custom encoder
        {"encoder": {"type": "alexnet", "use_pretrained": False}},  # TorchVision pretrained model encoder
    ],
)
def test_torchscript_e2e_image(tmpdir, csv_filename, kwargs):
    data_csv_path = os.path.join(tmpdir, csv_filename)
    image_dest_folder = os.path.join(tmpdir, "generated_images")
    input_features = [
        image_feature(image_dest_folder, **kwargs),
    ]
    output_features = [
        binary_feature(),
    ]
    backend = LocalTestBackend()
    config = {
        "input_features": input_features,
        "output_features": output_features,
        TRAINER: {"train_steps": 1, BATCH_SIZE: 128},
    }
    training_data_csv_path = generate_data(input_features, output_features, data_csv_path)

    validate_torchscript_outputs(tmpdir, config, backend, training_data_csv_path)


@pytest.mark.integration_tests_e
def test_torchscript_e2e_text(tmpdir, csv_filename):
    data_csv_path = os.path.join(tmpdir, csv_filename)
    input_features = [
        text_feature(encoder={"vocab_size": 3}, preprocessing={"tokenizer": tokenizer})
        for tokenizer in TORCHSCRIPT_COMPATIBLE_TOKENIZERS
    ]
    output_features = [
        text_feature(decoder={"vocab_size": 3}),
    ]
    backend = LocalTestBackend()
    config = {
        "input_features": input_features,
        "output_features": output_features,
        TRAINER: {"train_steps": 1, BATCH_SIZE: 128},
    }
    training_data_csv_path = generate_data(input_features, output_features, data_csv_path)

    validate_torchscript_outputs(tmpdir, config, backend, training_data_csv_path)


@pytest.mark.skipif(
    torchtext is None or torch.torch_version.TorchVersion(torchtext.__version__) < (0, 14, 0),
    reason="requires torchtext 0.14.0 or higher",
)
@pytest.mark.integration_tests_e
def test_torchscript_e2e_text_hf_tokenizer(tmpdir, csv_filename):
    data_csv_path = os.path.join(tmpdir, csv_filename)
    input_features = [text_feature(encoder={"vocab_size": 3, "type": "bert"})]
    output_features = [
        category_feature(),
    ]
    backend = LocalTestBackend()
    config = {
        "input_features": input_features,
        "output_features": output_features,
        TRAINER: {"train_steps": 1, BATCH_SIZE: 128, EVAL_BATCH_SIZE: 128},
    }
    training_data_csv_path = generate_data(input_features, output_features, data_csv_path)

    validate_torchscript_outputs(tmpdir, config, backend, training_data_csv_path)


@pytest.mark.skipif(
    torchtext is None or torch.torch_version.TorchVersion(torchtext.__version__) < (0, 14, 0),
    reason="requires torchtext 0.14.0 or higher",
)
@pytest.mark.integration_tests_e
def test_torchscript_e2e_text_hf_tokenizer_truncated_sequence(tmpdir, csv_filename):
    data_csv_path = os.path.join(tmpdir, csv_filename)
    input_features = [text_feature(encoder={"vocab_size": 3, "type": "bert"}, preprocessing={"max_sequence_length": 3})]
    output_features = [
        category_feature(),
    ]
    backend = LocalTestBackend()
    config = {
        "input_features": input_features,
        "output_features": output_features,
        TRAINER: {"train_steps": 1, BATCH_SIZE: 128},
    }
    training_data_csv_path = generate_data(input_features, output_features, data_csv_path)

    validate_torchscript_outputs(tmpdir, config, backend, training_data_csv_path)


@pytest.mark.integration_tests_e
def test_torchscript_e2e_sequence(tmpdir, csv_filename):
    data_csv_path = os.path.join(tmpdir, csv_filename)
    input_features = [
        sequence_feature(encoder={"vocab_size": 3}, preprocessing={"tokenizer": "space"}),
    ]
    output_features = [
        sequence_feature(decoder={"vocab_size": 3}),
    ]
    backend = LocalTestBackend()
    config = {
        "input_features": input_features,
        "output_features": output_features,
        TRAINER: {"train_steps": 1, BATCH_SIZE: 128},
    }
    training_data_csv_path = generate_data(input_features, output_features, data_csv_path)

    validate_torchscript_outputs(tmpdir, config, backend, training_data_csv_path)


@pytest.mark.integration_tests_e
def test_torchscript_e2e_timeseries(tmpdir, csv_filename):
    data_csv_path = os.path.join(tmpdir, csv_filename)
    input_features = [
        timeseries_feature(),
    ]
    output_features = [
        binary_feature(),
    ]
    backend = LocalTestBackend()
    config = {
        "input_features": input_features,
        "output_features": output_features,
        TRAINER: {"train_steps": 1, BATCH_SIZE: 128},
    }
    training_data_csv_path = generate_data(input_features, output_features, data_csv_path)

    validate_torchscript_outputs(tmpdir, config, backend, training_data_csv_path)


@pytest.mark.integration_tests_e
def test_torchscript_e2e_h3(tmpdir, csv_filename):
    data_csv_path = os.path.join(tmpdir, csv_filename)
    input_features = [
        h3_feature(),
    ]
    output_features = [
        binary_feature(),
    ]
    backend = LocalTestBackend()
    config = {
        "input_features": input_features,
        "output_features": output_features,
        TRAINER: {"train_steps": 1, BATCH_SIZE: 128},
    }
    training_data_csv_path = generate_data(input_features, output_features, data_csv_path)

    validate_torchscript_outputs(tmpdir, config, backend, training_data_csv_path)


@pytest.mark.integration_tests_e
def test_torchscript_e2e_date(tmpdir, csv_filename):
    data_csv_path = os.path.join(tmpdir, csv_filename)
    input_features = [
        date_feature(),
    ]
    output_features = [
        binary_feature(),
    ]
    backend = LocalTestBackend()
    config = {
        "input_features": input_features,
        "output_features": output_features,
        TRAINER: {"train_steps": 1, BATCH_SIZE: 128},
    }
    training_data_csv_path = generate_data(input_features, output_features, data_csv_path)

    validate_torchscript_outputs(tmpdir, config, backend, training_data_csv_path)


@pytest.mark.integration_tests_e
@pytest.mark.parametrize("vector_type", [torch.Tensor, list[torch.Tensor]])
def test_torchscript_preproc_vector_alternative_type(tmpdir, csv_filename, vector_type):
    data_csv_path = os.path.join(tmpdir, csv_filename)
    feature = vector_feature()
    input_features = [
        feature,
    ]
    output_features = [
        binary_feature(),
    ]
    backend = LocalTestBackend()
    config = {
        "input_features": input_features,
        "output_features": output_features,
        TRAINER: {"train_steps": 1, BATCH_SIZE: 128},
    }
    training_data_csv_path = generate_data(input_features, output_features, data_csv_path)

    # Initialize Ludwig model
    ludwig_model, script_module = initialize_torchscript_module(tmpdir, config, backend, training_data_csv_path)

    # Obtain preprocessed inputs from Python model
    preproc_inputs_expected, _ = preprocess_for_prediction(
        ludwig_model.config_obj.to_dict(),
        training_data_csv_path,
        ludwig_model.training_set_metadata,
        backend=backend,
        include_outputs=False,
    )

    df = pd.read_csv(training_data_csv_path)
    inputs = to_inference_module_input_from_dataframe(df, config, load_paths=True)

    def transform_vector_list(vector_list, vector_type):
        vectors = []
        for vector_str in vector_list:
            vectors.append(torch.tensor([float(x) for x in vector_str.split()]))

        if vector_type == torch.Tensor:
            vectors = torch.stack(vectors)
        return vectors

    inputs[feature[NAME]] = transform_vector_list(inputs[feature[NAME]], vector_type)

    preproc_inputs = script_module.preprocessor_forward(inputs)

    # Check that preproc_inputs is the same as preproc_inputs_expected.
    for feature_name_expected, feature_values_expected in preproc_inputs_expected.dataset.items():
        feature_name = feature_name_expected[: feature_name_expected.rfind("_")]  # remove proc suffix
        if feature_name not in preproc_inputs.keys():
            continue

        feature_values = preproc_inputs[feature_name]
        assert utils.is_all_close(feature_values, feature_values_expected), f"feature: {feature_name}"


@pytest.mark.integration_tests_e
@pytest.mark.parametrize("padding", ["left", "right"])
@pytest.mark.parametrize("fill_value", ["", "1.0"])
def test_torchscript_preproc_timeseries_alternative_type(tmpdir, csv_filename, padding, fill_value):
    data_csv_path = os.path.join(tmpdir, csv_filename)
    feature = timeseries_feature(
        preprocessing={
            "padding": padding,
            "timeseries_length_limit": 4,
            "fill_value": "1.0",
        },
        encoder={"max_len": 7},
    )
    input_features = [
        feature,
    ]
    output_features = [
        binary_feature(),
    ]
    backend = LocalTestBackend()
    config = {
        "input_features": input_features,
        "output_features": output_features,
        TRAINER: {"train_steps": 1, BATCH_SIZE: 128},
    }
    training_data_csv_path = generate_data(input_features, output_features, data_csv_path, nan_percent=0.2)

    # Initialize Ludwig model
    ludwig_model, script_module = initialize_torchscript_module(tmpdir, config, backend, training_data_csv_path)

    # Obtain preprocessed inputs from Python model
    preproc_inputs_expected, _ = preprocess_for_prediction(
        ludwig_model.config_obj.to_dict(),
        training_data_csv_path,
        ludwig_model.training_set_metadata,
        backend=backend,
        include_outputs=False,
    )

    df = pd.read_csv(training_data_csv_path)
    inputs = to_inference_module_input_from_dataframe(df, config, load_paths=True)

    def transform_timeseries_from_str_list_to_tensor_list(timeseries_list):
        timeseries = []
        for timeseries_str in timeseries_list:
            timeseries.append(torch.tensor([float(x) for x in timeseries_str.split()]))
        return timeseries

    inputs[feature[NAME]] = transform_timeseries_from_str_list_to_tensor_list(inputs[feature[NAME]])

    preproc_inputs = script_module.preprocessor_forward(inputs)

    # Check that preproc_inputs is the same as preproc_inputs_expected.
    for feature_name_expected, feature_values_expected in preproc_inputs_expected.dataset.items():
        feature_name = feature_name_expected[: feature_name_expected.rfind("_")]  # remove proc suffix
        assert feature_name in preproc_inputs.keys(), f'feature "{feature_name}" not found.'

        feature_values = preproc_inputs[feature_name]
        assert utils.is_all_close(feature_values, feature_values_expected), f'feature "{feature_name}" value mismatch.'


@pytest.mark.integration_tests_e
@pytest.mark.parametrize(
    "feature",
    [
        number_feature(),
        binary_feature(),
        category_feature(encoder={"vocab_size": 3}),
        bag_feature(encoder={"vocab_size": 3}),
        set_feature(encoder={"vocab_size": 3}),
        text_feature(encoder={"vocab_size": 3}),
        sequence_feature(encoder={"vocab_size": 3}),
        timeseries_feature(),
        h3_feature(),
        # TODO: future support
        # audio_feature(),  # default BFILL strategy is unintuitive at inference time
        # image_feature(),  # default BFILL strategy is unintuitive at inference time
        # vector_feature(), # does not have a missing_value_strategy
        # date_feature(),   # default fill with datetime.now() strategy is not scriptable
    ],
)
def test_torchscript_preproc_with_nans(tmpdir, csv_filename, feature):
    data_csv_path = os.path.join(tmpdir, csv_filename)
    input_features = [
        feature,
    ]
    output_features = [
        binary_feature(),
    ]
    backend = LocalTestBackend()
    config = {
        "input_features": input_features,
        "output_features": output_features,
        TRAINER: {"train_steps": 1, BATCH_SIZE: 128},
    }
    training_data_csv_path = generate_data(input_features, output_features, data_csv_path, nan_percent=0.2)

    # Initialize Ludwig model
    ludwig_model, script_module = initialize_torchscript_module(tmpdir, config, backend, training_data_csv_path)

    # Obtain preprocessed inputs from Python model
    preproc_inputs_expected, _ = preprocess_for_prediction(
        ludwig_model.config_obj.to_dict(),
        training_data_csv_path,
        ludwig_model.training_set_metadata,
        backend=backend,
        include_outputs=False,
    )

    df = pd.read_csv(training_data_csv_path)
    inputs = to_inference_module_input_from_dataframe(df, config, load_paths=True)
    preproc_inputs = script_module.preprocessor_forward(inputs)

    # Check that preproc_inputs is the same as preproc_inputs_expected.
    for feature_name_expected, feature_values_expected in preproc_inputs_expected.dataset.items():
        feature_name = feature_name_expected[: feature_name_expected.rfind("_")]  # remove proc suffix
        if feature_name not in preproc_inputs.keys():
            continue

        feature_values = preproc_inputs[feature_name]
        assert utils.is_all_close(feature_values, feature_values_expected), f"feature: {feature_name}"


@pytest.mark.skipif(torch.cuda.device_count() == 0, reason="test requires at least 1 gpu")
@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires gpu support")
@pytest.mark.integration_tests_e
@pytest.mark.distributed
@pytest.mark.parametrize(
    "feature_fn",
    [
        number_feature,
        image_feature,
        audio_feature,
        h3_feature,
        date_feature,
        # TODO: future support
        # binary_feature(),                # Torchscript takes List[str] as input, so currently CPU only
        # category_feature(encoder={"vocab_size": 3}),  # Torchscript takes List[str] as input, so currently CPU only
        # set_feature(encoder={"vocab_size": 3}),       # Torchscript takes List[str] as input, so currently CPU only
        # sequence_feature(encoder={"vocab_size": 3}),  # Torchscript takes List[str] as input, so currently CPU only
        # text_feature(encoder={"vocab_size": 3}),      # Torchscript takes List[str] as input, so currently CPU only
        # vector_feature(),                # Torchscript takes List[str] as input, so currently CPU only
        # bag_feature(encoder={"vocab_size": 3}),       # Torchscript takes List[str] as input, so currently CPU only
        # timeseries_feature(),            # Torchscript takes List[str] as input, so currently CPU only
    ],
)
def test_torchscript_preproc_gpu(tmpdir, csv_filename, feature_fn):
    data_csv_path = os.path.join(tmpdir, csv_filename)

    feature_kwargs = {}
    if feature_fn in {image_feature, audio_feature}:
        dest_folder = os.path.join(tmpdir, "generated_samples")
        feature_kwargs["folder"] = dest_folder

    input_features = [
        feature_fn(**feature_kwargs),
    ]
    output_features = [
        binary_feature(),
    ]

    config = {
        "input_features": input_features,
        "output_features": output_features,
        TRAINER: {"train_steps": 1, BATCH_SIZE: 128},
    }
    backend = RAY
    training_data_csv_path = generate_data(input_features, output_features, data_csv_path)
    _, script_module = initialize_torchscript_module(
        tmpdir,
        config,
        backend,
        training_data_csv_path,
        device=torch.device("cuda"),
    )

    df = pd.read_csv(training_data_csv_path)
    inputs = to_inference_module_input_from_dataframe(
        df,
        config,
        load_paths=True,
        device=torch.device("cuda"),
    )
    preproc_inputs = script_module.preprocessor_forward(inputs)

    for name, values in preproc_inputs.items():
        assert values.is_cuda, f'feature "{name}" tensors are not on GPU'


@pytest.mark.skipif(torch.cuda.device_count() == 0, reason="test requires at least 1 gpu")
@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires gpu support")
@pytest.mark.integration_tests_e
@pytest.mark.distributed
@pytest.mark.parametrize(
    "feature_fn",
    [
        number_feature,
        category_feature,
        binary_feature,
        set_feature,
        vector_feature,
        sequence_feature,
        text_feature,
    ],
)
def test_torchscript_postproc_gpu(tmpdir, csv_filename, feature_fn):
    data_csv_path = os.path.join(tmpdir, csv_filename)

    feature_kwargs = {}
    if feature_fn in {category_feature, set_feature, sequence_feature, text_feature}:
        feature_kwargs["vocab_size"] = 3

    input_features = [
        number_feature(),
    ]
    output_features = [
        feature_fn(**feature_kwargs),
    ]

    config = {
        "input_features": input_features,
        "output_features": output_features,
        TRAINER: {"train_steps": 1, BATCH_SIZE: 128},
    }
    backend = RAY
    training_data_csv_path = generate_data(input_features, output_features, data_csv_path)
    _, script_module = initialize_torchscript_module(
        tmpdir,
        config,
        backend,
        training_data_csv_path,
        device=torch.device("cuda"),
    )

    df = pd.read_csv(training_data_csv_path)
    inputs = to_inference_module_input_from_dataframe(
        df,
        config,
        load_paths=True,
        device=torch.device("cuda"),
    )
    postproc_outputs = script_module(inputs)

    for feature_name, feature_outputs in postproc_outputs.items():
        for output_name, output_values in feature_outputs.items():
            assert utils.is_all_tensors_cuda(output_values), f"{feature_name}.{output_name} tensors are not on GPU"


def validate_torchscript_outputs(tmpdir, config, backend, training_data_csv_path, tolerance=1e-8):
    # Train Ludwig (Pythonic) model:
    ludwig_model, script_module = initialize_torchscript_module(
        tmpdir,
        config,
        backend,
        training_data_csv_path,
    )

    # Obtain predictions from Python model
    preds_dict, _ = ludwig_model.predict(dataset=training_data_csv_path, return_type=dict)

    df = pd.read_csv(training_data_csv_path)
    inputs = to_inference_module_input_from_dataframe(df, config, load_paths=True)
    outputs = script_module(inputs)

    # TODO: these are the only outputs we provide from Torchscript for now
    ts_outputs = {PREDICTIONS, PROBABILITIES, LOGITS}

    # Compare results from Python trained model against Torchscript
    for feature_name, feature_outputs_expected in preds_dict.items():
        assert feature_name in outputs

        feature_outputs = outputs[feature_name]
        for output_name, output_values_expected in feature_outputs_expected.items():
            if output_name not in ts_outputs:
                continue

            assert output_name in feature_outputs
            output_values = feature_outputs[output_name]
            assert utils.has_no_grad(output_values), f'"{feature_name}.{output_name}" tensors have gradients'
            assert utils.is_all_close(
                output_values, output_values_expected
            ), f'"{feature_name}.{output_name}" tensors are not close to ludwig model'


def initialize_torchscript_module(tmpdir, config, backend, training_data_csv_path, device=None):
    # Initialize Ludwig model
    ludwig_model = LudwigModel(config, backend=backend)
    ludwig_model.train(
        dataset=training_data_csv_path,
        skip_save_training_description=True,
        skip_save_training_statistics=True,
        skip_save_model=True,
        skip_save_progress=True,
        skip_save_log=True,
        skip_save_processed_input=True,
    )

    # Always use CPU for torchscript since inference inputs are CPU tensors
    if device is None:
        device = torch.device("cpu")

    # Create graph inference model (Torchscript) from trained Ludwig model.
    script_module = ludwig_model.to_torchscript(device=device)
    # Ensure torchscript saving/loading does not affect final predictions.
    script_module_path = os.path.join(tmpdir, "inference_module.pt")
    torch.jit.save(script_module, script_module_path)
    script_module = torch.jit.load(script_module_path)
    return ludwig_model, script_module
