import os

import pytest

import tempfile

from unittest import mock

import numpy as np

import pandas as pd

from maui import Maui

from maui.maui_warnings import MauiWarning

samples = [f"Sample_{i}" for i in range(10)]

n_features_1 = 20

df1 = pd.DataFrame(

    np.random.randn(n_features_1, len(samples)),

    columns=samples,

    index=[f"data1_feature_{i}" for i in range(n_features_1)],

)

n_features_2 = 6

df2 = pd.DataFrame(

    np.random.randn(n_features_2, len(samples)),

    columns=samples,

    index=[f"data2_feature_{i}" for i in range(n_features_2)],

)

df_empty = pd.DataFrame(

    np.random.randn(0, len(samples)),

    columns=samples,

    index=[f"data0_feature_{i}" for i in range(0)],

)

def test_validate_X_fails_if_not_dict():

    maui_model = Maui()

    with pytest.raises(ValueError):

        maui_model._validate_X([1, 2, 3])

def test_validate_X_fails_if_samples_mismatch():

    maui_model = Maui()

    with pytest.raises(ValueError):

        df2_bad = df2.iloc[:, :2]

        data_with_mismatching_samples = {"a": df1, "b": df2_bad}

        maui_model._validate_X(data_with_mismatching_samples)

def test_validate_X_fails_if_some_data_empty():

    maui_model = Maui()

    with pytest.raises(ValueError):

        maui_model._validate_X({"a": df1, "e": df_empty})

def test_validate_X_returns_true_on_valid_data():

    maui_model = Maui()

    valid_data = {"a": df1, "b": df2}

    assert maui_model._validate_X(valid_data)

def test_dict2array():

    maui_model = Maui()

    arr = maui_model._dict2array({"data1": df1, "data2": df2})

    assert arr.shape[0] == len(df1.columns)

    assert arr.shape[1] == len(df1.index) + len(df2.index)

def test_maui_saves_feature_correlations():

    maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1)

    z = maui_model.fit_transform({"d1": df1, "d2": df2})

    r = maui_model.get_feature_correlations()

    assert r is not None

    assert hasattr(maui_model, "feature_correlations_")

def test_maui_saves_w():

    maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1)

    z = maui_model.fit_transform({"d1": df1, "d2": df2})

    w = maui_model.get_linear_weights()

    assert w is not None

    assert hasattr(maui_model, "w_")

def test_maui_saves_neural_weight_product():

    maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1)

    z = maui_model.fit_transform({"d1": df1, "d2": df2})

    nwp = maui_model.get_neural_weight_product()

    assert nwp is not None

    assert hasattr(maui_model, "nwp_")

    print(maui_model.encoder.summary())

    w1 = maui_model.encoder.layers[2].get_weights()[0]

    w2 = maui_model.encoder.layers[3].get_weights()[0]

    nwp_11 = np.dot(w1[0, :], w2[:, 0])

    assert np.allclose(nwp_11, nwp.iloc[0, 0])

def test_maui_updates_neural_weight_product_when_training():

    maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1)

    z_before = maui_model.fit_transform({"d1": df1, "d2": df2})

    nwp_before_fine_tuning = maui_model.get_neural_weight_product()

    maui_model.fine_tune({"d1": df1, "d2": df2})

    z_after = maui_model.transform({"d1": df1, "d2": df2})

    nwp_after_fine_tuning = maui_model.get_neural_weight_product()

    assert not np.allclose(z_before, z_after)

    assert not np.allclose(nwp_before_fine_tuning, nwp_after_fine_tuning)

def test_maui_clusters_with_single_k():

    maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1)

    maui_model.z_ = pd.DataFrame(

        np.random.randn(10, 2),

        index=[f"sample {i}" for i in range(10)],

        columns=["LF1", "LF2"],

    )

    maui_model.x_ = pd.DataFrame(

        np.random.randn(20, 10),

        index=[f"feature {i}" for i in range(20)],

        columns=[f"sample {i}" for i in range(10)],

    )

    yhat = maui_model.cluster(5)

    assert yhat.shape == (10,)

def test_maui_clusters_picks_optimal_k_by_ami():

    ami_mock = mock.Mock()

    ami_mock.side_effect = [

        2,

        3,

        1,

    ]  # the optimal AMI will be given at the second trial

    with mock.patch("sklearn.metrics.adjusted_mutual_info_score", ami_mock):

        maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1)

        maui_model.z_ = pd.DataFrame(

            np.random.randn(10, 2),

            index=[f"sample {i}" for i in range(10)],

            columns=["LF1", "LF2"],

        )

        maui_model.x_ = pd.DataFrame(

            np.random.randn(20, 10),

            index=[f"feature {i}" for i in range(20)],

            columns=[f"sample {i}" for i in range(10)],

        )

        the_y = pd.Series(np.arange(10), index=maui_model.z_.index)

        maui_model.cluster(

            ami_y=the_y, optimal_k_range=[1, 2, 3]

        )  # the second trial is k=2

        print(maui_model.kmeans_scores)

        assert maui_model.optimal_k_ == 2

def test_maui_clusters_picks_optimal_k_by_silhouette():

    silhouette_mock = mock.Mock()

    silhouette_mock.side_effect = [

        2,

        3,

        1,

    ]  # the optimal silhouette will be given at the second trial

    with mock.patch("sklearn.metrics.silhouette_score", silhouette_mock):

        maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1)

        maui_model.z_ = pd.DataFrame(

            np.random.randn(10, 2),

            index=[f"sample {i}" for i in range(10)],

            columns=["LF1", "LF2"],

        )

        maui_model.x_ = pd.DataFrame(

            np.random.randn(20, 10),

            index=[f"feature {i}" for i in range(20)],

            columns=[f"sample {i}" for i in range(10)],

        )

        maui_model.cluster(

            optimal_k_method="silhouette", optimal_k_range=[1, 2, 3]

        )  # the second trial is k=2

        assert maui_model.optimal_k_ == 2

def test_maui_clusters_picks_optimal_k_with_custom_scoring():

    scorer = mock.Mock()

    scorer.side_effect = [2, 3, 1]  # the optimal AMI will be given at the second trial

    scorer.__name__ = "mock_scorer"

    maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1)

    maui_model.z_ = pd.DataFrame(

        np.random.randn(10, 2),

        index=[f"sample {i}" for i in range(10)],

        columns=["LF1", "LF2"],

    )

    maui_model.x_ = pd.DataFrame(

        np.random.randn(20, 10),

        index=[f"feature {i}" for i in range(20)],

        columns=[f"sample {i}" for i in range(10)],

    )

    maui_model.cluster(

        optimal_k_method=scorer, optimal_k_range=[1, 2, 3]

    )  # the second trial is k=2

    assert maui_model.optimal_k_ == 2

def test_maui_computes_roc_and_auc():

    maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1)

    maui_model.z_ = pd.DataFrame(

        [

            [0, 1, 1, 1, 0, 1, 1, 0, 0],

            [1, 0, 0, 0, 0, 0, 1, 1, 0],

            [1, 0, 1, 0, 0, 0, 1, 1, 0],

            [1, 0, 0, 1, 0, 0, 1, 1, 0],

            [1, 0, 0, 0, 1, 1, 1, 1, 0],

            [1, 1, 1, 0, 0, 0, 1, 1, 1],

        ],

        index=[f"sample {i}" for i in range(6)],

        columns=[f"LF{i}" for i in range(9)],

    )

    y = pd.Series(["a", "b", "a", "c", "b", "c"], index=maui_model.z_.index)

    rocs = maui_model.compute_roc(y, cv_folds=2)

    assert rocs == maui_model.roc_curves_

    assert "a" in rocs

    assert "b" in rocs

    assert "c" in rocs

    assert "mean" in rocs

    aucs = maui_model.compute_auc(y, cv_folds=2)

    assert aucs == maui_model.aucs_

def test_maui_clusters_only_samples_in_y_index_when_optimizing():

    np.random.seed(0)

    maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1)

    maui_model.z_ = pd.DataFrame(

        np.random.randn(10, 2),

        index=[f"sample {i}" for i in range(10)],

        columns=["LF1", "LF2"],

    )

    maui_model.x_ = pd.DataFrame(

        np.random.randn(20, 10),

        index=[f"feature {i}" for i in range(20)],

        columns=[f"sample {i}" for i in range(10)],

    )

    y = pd.Series(

        ["a", "a", "a", "b", "b", "b"], index=[f"sample {i}" for i in range(6)]

    )

    yhat = maui_model.cluster(ami_y=y, optimal_k_range=[1, 2, 3])

    assert set(yhat.index) == set(y.index)

def test_select_clinical_factors():

    maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1)

    maui_model.z_ = pd.DataFrame(

        [

            [1, 1, 1, 0, 0, 0, 1, 0, 1],

            [1, 1, 1, 1, 0, 1, 1, 1, 0],

            [1, 1, 1, 1, 0, 1, 1, 1, 0],

            [1, 1, 1, 1, 0, 1, 1, 1, 0],

            [1, 1, 1, 1, 0, 1, 1, 1, 0],

            [1, 1, 1, 1, 1, 0, 0, 1, 0],

            [0, 0, 0, 1, 0, 0, 1, 1, 0],

            [0, 0, 0, 1, 0, 0, 1, 1, 0],

            [0, 0, 0, 1, 0, 0, 1, 1, 0],

            [0, 0, 0, 1, 0, 0, 1, 1, 0],

            [0, 0, 0, 1, 0, 1, 1, 1, 1],

        ],

        index=[f"sample {i}" for i in range(11)],

        columns=[f"LF{i}" for i in range(9)],

    )  # here the first 3 factors separate the groups and the last 6 do not

    durations = [

        1,

        2,

        3,

        4,

        5,

        6,

        1000,

        2000,

        3000,

        4000,

        5000,

    ]  # here the first 3 have short durations, the last 3 longer ones

    observed = [True] * 11  # all events observed

    survival = pd.DataFrame(

        dict(duration=durations, observed=observed),

        index=[f"sample {i}" for i in range(11)],

    )

    z_clin = maui_model.select_clinical_factors(survival, cox_penalizer=1, alpha=0.1)

    assert "LF0" in z_clin.columns

    assert "LF5" not in z_clin.columns

def test_maui_computes_harrells_c():

    maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1)

    maui_model.z_ = pd.DataFrame(

        [

            [1, 1, 1, 0, 0, 0, 1, 0, 1],

            [1, 1, 1, 1, 0, 1, 1, 1, 0],

            [1, 1, 1, 1, 0, 1, 1, 1, 0],

            [1, 1, 1, 1, 0, 1, 1, 1, 0],

            [1, 1, 1, 1, 0, 1, 1, 1, 0],

            [1, 1, 1, 1, 1, 0, 0, 1, 0],

            [0, 0, 0, 1, 0, 0, 1, 1, 0],

            [0, 0, 0, 1, 0, 0, 1, 1, 0],

            [0, 0, 0, 1, 0, 0, 1, 1, 0],

            [0, 0, 0, 1, 0, 0, 1, 1, 0],

            [0, 0, 0, 1, 0, 1, 1, 1, 1],

        ],

        index=[f"sample {i}" for i in range(11)],

        columns=[f"LF{i}" for i in range(9)],

    )  # here the first 3 factors separate the groups and the last 6 do not

    durations = [

        1,

        2,

        3,

        4,

        5,

        6,

        1000,

        2000,

        3000,

        4000,

        5000,

    ]  # here the first 3 have short durations, the last 3 longer ones

    observed = [True] * 11  # all events observed

    survival = pd.DataFrame(

        dict(duration=durations, observed=observed),

        index=[f"sample {i}" for i in range(11)],

    )

    cs = maui_model.c_index(

        survival,

        clinical_only=True,

        duration_column="duration",

        observed_column="observed",

        cox_penalties=[0.1, 1, 10, 100, 1000, 10000],

        cv_folds=3,

        sel_clin_alpha=0.1,

        sel_clin_penalty=1,

    )

    print(cs)

    assert np.allclose(cs, [0.5, 0.8, 0.5], atol=0.05)

def test_maui_produces_same_prediction_when_run_twice():

    """This is to show the maui encoder model picks the mean of

    the distribution, not a sample."""

    maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1)

    maui_model = maui_model.fit({"d1": df1, "d2": df2})

    z1 = maui_model.transform({"d1": df1, "d2": df2})

    z2 = maui_model.transform({"d1": df1, "d2": df2})

    assert np.allclose(z1, z2)

def test_maui_produces_different_prediction_when_run_twice_with_sampling():

    """This is to show the maui encoder model picks the mean of

    the distribution, not a sample."""

    maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1)

    maui_model = maui_model.fit({"d1": df1, "d2": df2})

    z1 = maui_model.transform({"d1": df1, "d2": df2}, encoder="sample")

    z2 = maui_model.transform({"d1": df1, "d2": df2}, encoder="sample")

    assert not np.allclose(z1, z2)

def test_maui_produces_nonnegative_zs_if_relu_embedding_true():

    maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1, relu_embedding=True)

    maui_model = maui_model.fit({"d1": df1, "d2": df2})

    z1 = maui_model.transform({"d1": df1, "d2": df2})

    assert np.all(z1 >= 0)

def test_maui_produces_pos_and_neg_zs_if_relu_embedding_false():

    maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1, relu_embedding=False)

    maui_model = maui_model.fit({"d1": df1, "d2": df2})

    z1 = maui_model.transform({"d1": df1, "d2": df2})

    assert not np.all(z1 >= 0)

def test_maui_runs_with_deep_not_stacked_vae():

    maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1, architecture="deep")

    z = maui_model.fit_transform({"d1": df1, "d2": df2})

def test_maui_complains_if_wrong_architecture():

    with pytest.raises(ValueError):

        maui_model = Maui(

            n_hidden=[10], n_latent=2, epochs=1, architecture="wrong value"

        )

def test_maui_supports_single_layer_vae():

    maui_model = Maui(n_hidden=None, n_latent=2, epochs=1)

    maui_model = maui_model.fit({"d1": df1, "d2": df2})

    z1 = maui_model.transform({"d1": df1, "d2": df2})

def test_maui_supports_not_deep_deep_vae():

    maui_model = Maui(n_hidden=None, n_latent=2, epochs=1, architecture="deep")

    z = maui_model.fit_transform({"d1": df1, "d2": df2})

def test_maui_drops_unexplanatody_factors_by_r2():

    maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1)

    maui_model.z_ = pd.DataFrame(

        [

            [1, 1, 1, 0, 0, 0, 1, 0, 0],

            [1, 1, 1, 1, 0, 1, 1, 1, 0],

            [1, 1, 1, 1, 0, 1, 1, 1, 0],

            [1, 1, 1, 1, 0, 1, 1, 1, 0],

            [1, 1, 1, 1, 0, 1, 1, 1, 0],

            [1, 1, 1, 1, 1, 0, 0, 1, 0],

            [0, 0, 0, 1, 0, 0, 1, 1, 0],

            [0, 0, 0, 1, 0, 0, 1, 1, 0],

            [0, 0, 0, 1, 0, 0, 1, 1, 0],

            [0, 0, 0, 1, 0, 0, 1, 1, 0],

            [0, 0, 0, 1, 0, 1, 1, 1, 0],

        ],

        index=[f"sample {i}" for i in range(11)],

        columns=[f"LF{i}" for i in range(9)],

        dtype=float,

    )  # here the first 8 latent factors have R2 above threshold, the last does not

    maui_model.x_ = pd.DataFrame(

        [[1], [1], [1], [1], [1], [1], [0], [0], [0], [0], [0]],

        index=[f"sample {i}" for i in range(11)],

        columns=["Feature 1"],

        dtype=float,

    )

    z_filt = maui_model.drop_unexplanatory_factors()

    assert z_filt.shape[1] == 8

def test_maui_merges_latent_factors():

    maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1)

    maui_model.z_ = pd.DataFrame(

        [

            [1, 1, 1, 0, 0, 0, 1, 0, 0],

            [1, 1, 1, 1, 0, 1, 1, 1, 0],

            [1, 1, 1, 1, 0, 1, 1, 1, 0],

            [1, 1, 1, 1, 0, 1, 1, 1, 0],

            [1, 1, 1, 1, 0, 1, 1, 1, 0],

            [1, 1, 1, 1, 1, 0, 0, 1, 0],

            [0, 0, 0, 1, 0, 0, 1, 1, 0],

            [0, 0, 0, 1, 0, 0, 1, 1, 0],

            [0, 0, 0, 1, 0, 0, 1, 1, 0],

            [0, 0, 0, 1, 0, 0, 1, 1, 0],

            [0, 0, 0, 1, 0, 1, 1, 1, 0],

        ],

        index=[f"sample {i}" for i in range(11)],

        columns=[f"LF{i}" for i in range(9)],

        dtype=float,

    )  # expect 0,1,2 to be merged, and 3,7 to be merged

    z_merged = maui_model.merge_similar_latent_factors(distance_metric="euclidean")

    assert z_merged.shape[1] == 6

    assert "0_1_2" in z_merged.columns

    assert "3_7" in z_merged.columns

def test_maui_merges_latent_factors_by_w():

    maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1)

    maui_model.z_ = pd.DataFrame(

        [

            [1, 1, 1, 0, 0, 0, 1, 0, 0],

            [1, 1, 1, 1, 0, 1, 1, 1, 0],

            [1, 1, 1, 1, 0, 1, 1, 1, 0],

            [1, 1, 1, 1, 0, 1, 1, 1, 0],

            [1, 1, 1, 1, 0, 1, 1, 1, 0],

            [1, 1, 1, 1, 1, 0, 0, 1, 0],

            [0, 0, 0, 1, 0, 0, 1, 1, 0],

            [0, 0, 0, 1, 0, 0, 1, 1, 0],

            [0, 0, 0, 1, 0, 0, 1, 1, 0],

            [0, 0, 0, 1, 0, 0, 1, 1, 0],

            [0, 0, 0, 1, 0, 1, 1, 1, 0],

        ],

        index=[f"sample {i}" for i in range(11)],

        columns=[f"LF{i}" for i in range(9)],

        dtype=float,

    )

    maui_model.x_ = pd.DataFrame(

        [[1], [1], [1], [1], [1], [1], [0], [0], [0], [0], [0]],

        index=[f"sample {i}" for i in range(11)],

        columns=["Feature 1"],

        dtype=float,

    )

    # with these z and x, expect 0,1,2 and 4,5 and 3,6,7

    z_merged = maui_model.merge_similar_latent_factors(

        distance_in="w", distance_metric="euclidean"

    )

    assert z_merged.shape[1] == 4

    assert "0_1_2" in z_merged.columns

    assert "3_6_7" in z_merged.columns

    assert "4_5" in z_merged.columns

def test_maui_merge_latent_factors_complains_if_unknown_merge_by():

    maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1)

    maui_model.z_ = pd.DataFrame(

        [

            [1, 1, 1, 0, 0, 0, 1, 0, 0],

            [1, 1, 1, 1, 0, 1, 1, 1, 0],

            [1, 1, 1, 1, 0, 1, 1, 1, 0],

            [1, 1, 1, 1, 0, 1, 1, 1, 0],

            [1, 1, 1, 1, 0, 1, 1, 1, 0],

            [1, 1, 1, 1, 1, 0, 0, 1, 0],

            [0, 0, 0, 1, 0, 0, 1, 1, 0],

            [0, 0, 0, 1, 0, 0, 1, 1, 0],

            [0, 0, 0, 1, 0, 0, 1, 1, 0],

            [0, 0, 0, 1, 0, 0, 1, 1, 0],

            [0, 0, 0, 1, 0, 1, 1, 1, 0],

        ],

        index=[f"sample {i}" for i in range(11)],

        columns=[f"LF{i}" for i in range(9)],

        dtype=float,

    )  # expect 0,1,2 to be merged, and 3,7 to be merged

    with pytest.raises(Exception):

        z_merged = maui_model.merge_similar_latent_factors(

            distance_in="xxx", distance_metric="euclidean"

        )

def test_maui_can_save_to_folder():

    maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1)

    maui_model = maui_model.fit({"d1": df1, "d2": df2})

    with tempfile.TemporaryDirectory() as tmpdirname:

        maui_model.save(tmpdirname)

        assert os.path.isfile(os.path.join(tmpdirname, "maui_weights.h5"))

        assert os.path.isfile(os.path.join(tmpdirname, "maui_args.json"))

def test_maui_can_load_from_folder():

    maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1)

    maui_model = maui_model.fit({"d1": df1, "d2": df2})

    with tempfile.TemporaryDirectory() as tmpdirname:

        maui_model.save(tmpdirname)

        maui_model_from_disk = Maui.load(tmpdirname)

    assert maui_model_from_disk.n_latent == maui_model.n_latent

    assert np.allclose(

        maui_model.vae.get_weights()[0], maui_model_from_disk.vae.get_weights()[0]

    )

    assert np.allclose(

        maui_model.transform({"d1": df1, "d2": df2}),

        maui_model_from_disk.transform({"d1": df1, "d2": df2}),

    )

def test_maui_can_print_verbose_training(capsys):

    maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1)

    maui_model = maui_model.fit({"d1": df1, "d2": df2})

    stdout, stderr = capsys.readouterr()

    assert stdout == ""

    maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1, verbose=1)

    maui_model = maui_model.fit({"d1": df1, "d2": df2})

    stdout, stderr = capsys.readouterr()

    assert "Epoch" in stdout

def test_maui_model_makes_2_layer_vae():

    maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1, input_dim=10)

    layers_names = [l.name for l in maui_model.vae.layers]

    assert "hidden_dim_0_mean" in layers_names

    assert "latent_mean" in layers_names

    assert "decode_hidden_0" in layers_names

    assert "reconstruction" in layers_names

    assert "decode_hidden_1" not in layers_names

def test_maui_model_makes_one_layer_vae():

    maui_model = Maui(n_hidden=[], n_latent=2, epochs=1, input_dim=10)

    layers_names = [l.name for l in maui_model.vae.layers]

    print(layers_names)

    assert layers_names[-1] == "reconstruction"

    assert not any(

        "decode_hidden" in name for name in layers_names

    ), "Has a decode hidden..."

    assert not any("hidden_dim" in name for name in layers_names), "Has a hidden dim..."

def test_maui_model_validates_feature_names_on_predict_after_fit():

    maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1)

    maui_model.fit({"d1": df1, "d2": df2})

    z = maui_model.transform({"d1": df1, "d2": df2})

    df1_wrong_features = df1.reindex(df1.index[: len(df1.index) - 1])

    with pytest.raises(ValueError):

        z = maui_model.transform({"df1": df1_wrong_features, "df2": df2})

def test_maui_model_saves_feature_names_to_disk():

    maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1)

    maui_model = maui_model.fit({"d1": df1, "d2": df2})

    with tempfile.TemporaryDirectory() as tmpdirname:

        maui_model.save(tmpdirname)

        maui_model_from_disk = Maui.load(tmpdirname)

    assert maui_model.feature_names == maui_model_from_disk.feature_names

def test_maui_model_loads_model_without_feature_names_from_disk_and_warns():

    maui_model = Maui(n_hidden=[10], n_latent=2, epochs=1)

    maui_model = maui_model.fit({"d1": df1, "d2": df2})

    with tempfile.TemporaryDirectory() as tmpdirname:

        maui_model.save(tmpdirname)

        os.remove(os.path.join(tmpdirname, "maui_feature_names.txt"))

        with pytest.warns(MauiWarning):

            maui_model_from_disk = Maui.load(tmpdirname)

        assert maui_model_from_disk.feature_names is None

def test_maui_can_fine_tune():

    maui_model = Maui(n_hidden=[], n_latent=2, epochs=1)

    maui_model = maui_model.fit({"d1": df1, "d2": df2})

    maui_model.fine_tune({"d1": df1, "d2": df2}, epochs=1)

def test_maui_complains_if_fine_tune_with_wrong_features():

    maui_model = Maui(n_hidden=[], n_latent=2, epochs=1)

    maui_model.fit({"d1": df1, "d2": df2})

    df1_wrong_features = df1.reindex(df1.index[: len(df1.index) - 1])

    with pytest.raises(ValueError):

        z = maui_model.fine_tune({"df1": df1_wrong_features, "df2": df2})