"""Tests for classifier widget initialization"""

from pathlib import Path

import imageio
import numpy as np
import pandas as pd
import pytest

from napari_feature_classifier.classifier_widget import (
    ClassifierWidget,
    LoadClassifierContainer,
)
from napari_feature_classifier.feature_loader_widget import make_features


def run_and_wait(container, qtbot, timeout=30_000):
    """Trigger container.run() and block until the background worker finishes.

    Waits for the run button to be re-enabled, which is the last action in
    both _on_run_done and _on_run_error — guaranteeing those callbacks have
    fully completed regardless of which thread they ran in.
    """
    container.run()
    qtbot.waitUntil(lambda: container._run_button.enabled, timeout=timeout)


lbl_img_np = imageio.v2.imread(
    Path("src/napari_feature_classifier/sample_data/test_labels.tif")
)
clf_path = Path("src/napari_feature_classifier/sample_data/test_labels_classifier.clf")


# make_napari_viewer is a pytest fixture that returns a napari viewer object
# capsys is a pytest fixture that captures stdout and stderr output streams
def test_classifier_widgets_initialization_no_features_selected(make_napari_viewer):
    """
    Tests if the main widget launches
    """
    # make viewer and add an image layer using our fixture
    viewer = make_napari_viewer()
    _ = viewer.add_labels(lbl_img_np)

    # Start init widget
    classifier_widget = ClassifierWidget(viewer)

    classifier_widget.initialize_run_widget()

    # TODO: Catch that it doesn't actually initialize, but instead shows a
    # message to the user that now features were selected
    assert classifier_widget._run_container is None


def test_classifier_initializtion_without_label_image(make_napari_viewer):
    viewer = make_napari_viewer()
    classifier_widget = ClassifierWidget(viewer)
    assert classifier_widget._init_container is not None
    assert classifier_widget._init_container._last_selected_label_layer is None


features = make_features(np.unique(lbl_img_np)[1:], roi_id="ROI1", n_features=6)
features_no_roi_id = features.drop(columns=["roi_id"])


# make_napari_viewer is a pytest fixture that returns a napari viewer object
@pytest.mark.parametrize("features", [features, features_no_roi_id])
def test_running_classification_through_widget(
    features, make_napari_viewer, qtbot, tmp_path, monkeypatch
):
    """
    Tests if the main widget launches
    """
    monkeypatch.chdir(tmp_path)
    # make viewer and add an image layer using our fixture
    viewer = make_napari_viewer()
    label_layer = viewer.add_labels(lbl_img_np)
    label_layer.features = features

    # Start init widget
    classifier_widget = ClassifierWidget(viewer)

    # Select relevant features
    assert classifier_widget._init_container is not None
    classifier_widget._init_container._feature_combobox.value = [
        "feature_1",
        "feature_2",
        "feature_3",
    ]

    classifier_widget.initialize_run_widget()

    # Check that the run widget is initialized
    assert classifier_widget._run_container is not None

    # Add some annotations manually
    label_layer.features.loc[0, "annotations"] = 1.0
    label_layer.features.loc[1, "annotations"] = 1.0
    label_layer.features.loc[3, "annotations"] = 2.0

    # Run the classifier and wait for the background worker to finish
    run_and_wait(classifier_widget._run_container, qtbot)

    # Assert something that the layer is visible, predictions exist and are not NaN
    assert classifier_widget._run_container._prediction_manager.prediction_layer.visible
    assert "prediction" in label_layer.features.columns
    assert pd.notna(label_layer.features["prediction"]).all().all()

    # Check that the classifier file was saved
    assert (tmp_path / "lbl_img_np_classifier.clf").exists()


# Test classifier results export
def test_prediction_export(make_napari_viewer, qtbot, capsys, tmp_path, monkeypatch):
    """
    Tests if the main widget launches
    """
    monkeypatch.chdir(tmp_path)
    # make viewer and add an image layer using our fixture
    viewer = make_napari_viewer()
    label_layer = viewer.add_labels(lbl_img_np)
    label_layer.features = features

    # Start init widget
    classifier_widget = ClassifierWidget(viewer)

    # Select relevant features
    assert classifier_widget._init_container is not None
    classifier_widget._init_container._feature_combobox.value = [
        "feature_1",
        "feature_2",
        "feature_3",
    ]

    classifier_widget.initialize_run_widget()

    # Check that the run widget is initialized
    assert classifier_widget._run_container is not None

    # Add some annotations manually
    label_layer.features.loc[0, "annotations"] = 1.0
    label_layer.features.loc[1, "annotations"] = 1.0
    label_layer.features.loc[3, "annotations"] = 2.0

    # Run the classifier and wait for the background worker to finish
    run_and_wait(classifier_widget._run_container, qtbot)

    # Test result export
    classifier_widget._run_container._export_panel.export_results()
    df = pd.read_csv(classifier_widget._run_container._export_panel._export_path)
    assert df.shape == (16, 5)
    assert df["prediction"].isna().sum() == 0

    message = "INFO: Annotations were saved at lbl_img_np_predictions.csv"
    assert message in capsys.readouterr().out


# TODO: Add a test to check the overwrite confirmations working correctly
# For classifier files, for annotations and for exported predictions


features = make_features(np.unique(lbl_img_np)[1:], roi_id="ROI1", n_features=6)


# make_napari_viewer is a pytest fixture that returns a napari viewer object
def test_classifier_fails_running_without_annotation(
    make_napari_viewer, qtbot, capsys, tmp_path, monkeypatch
):
    """
    Tests if the main widget launches
    """
    monkeypatch.chdir(tmp_path)
    # make viewer and add an image layer using our fixture
    viewer = make_napari_viewer()
    label_layer = viewer.add_labels(lbl_img_np)
    label_layer.features = features

    # Start init widget
    classifier_widget = ClassifierWidget(viewer)

    # Select relevant features
    assert classifier_widget._init_container is not None
    classifier_widget._init_container._feature_combobox.value = [
        "feature_1",
        "feature_2",
        "feature_3",
    ]

    classifier_widget.initialize_run_widget()

    # Run the classifier and wait for the background worker to finish
    assert classifier_widget._run_container is not None
    run_and_wait(classifier_widget._run_container, qtbot)

    expected_message = (
        "INFO: Training failed. A typical reason are not "
        "having enough annotations. \nThe error message was: Found array "
        "with 0 sample(s) (shape=(0, 3)) while a minimum of 1 is required "
        "by RandomForestClassifier."
    )
    assert expected_message in capsys.readouterr().out


def test_layer_selection_changes(make_napari_viewer):
    """
    Tests if the main widget launches
    """
    # make viewer and add an image layer using our fixture
    viewer = make_napari_viewer()
    label_layer1 = viewer.add_labels(lbl_img_np, name="test_labels")
    label_layer2 = viewer.add_labels(lbl_img_np, name="test_labels_2")
    label_layer2.features = features
    label_layer1.features = features

    # Start init widget
    classifier_widget = ClassifierWidget(viewer)

    # Select relevant features
    assert classifier_widget._init_container is not None
    classifier_widget._init_container._feature_combobox.value = [
        "feature_1",
        "feature_2",
        "feature_3",
    ]

    classifier_widget.initialize_run_widget()

    assert classifier_widget._run_container is not None
    assert (
        classifier_widget._run_container._last_selected_label_layer.name
        == "test_labels_2"
    )
    assert (
        str(classifier_widget._run_container._export_panel._export_path)
        == "test_labels_2_predictions.csv"
    )
    viewer.layers.selection.active = label_layer1
    assert (
        classifier_widget._run_container._last_selected_label_layer.name
        == "test_labels"
    )
    assert (
        classifier_widget._run_container._export_panel._export_path.name
        == "test_labels_predictions.csv"
    )


# make_napari_viewer is a pytest fixture that returns a napari viewer object
# capsys is a pytest fixture that captures stdout and stderr output streams
def test_load_classifier_widget(make_napari_viewer, capsys):
    """
    Tests if the load classifier widget launches
    """
    # make viewer and add an image layer using our fixture
    viewer = make_napari_viewer()
    label_layer = viewer.add_labels(lbl_img_np)
    label_layer.features = features

    # Start init widget
    loading_widget = LoadClassifierContainer(viewer)

    loading_widget._clf_destination.value = clf_path
    loading_widget.load()

    # Some basic checks that loading looks to have worked
    assert loading_widget._run_container is not None
    assert loading_widget._run_container._prediction_manager.prediction_layer.visible
    assert "prediction" in label_layer.features.columns


def test_change_of_loader_filter(make_napari_viewer, capsys):
    viewer = make_napari_viewer()
    loading_widget = LoadClassifierContainer(viewer)
    assert loading_widget._filter.value == "*.clf"
    assert loading_widget._clf_destination.filter is None
    loading_widget._filter.value = "*.pkl"
    assert loading_widget._clf_destination.filter == "*.pkl"
    assert loading_widget._filter.value == "*.pkl"
