# test_coil_opt_rowvec_diff.py
import tempfile

import numpy as np
import pytest
import pynibs
import h5py
import os

def test_rowvec_diff_basic_case():
    # 4 candidate coil configurations, symmetric distance matrix
    efields_diff_mat = np.array(
        [
            [0.0, 1.0, 2.0, 3.0],
            [1.0, 0.0, 4.0, 1.5],
            [2.0, 4.0, 0.0, 0.5],
            [3.0, 1.5, 0.5, 0.0],
        ]
    )

    selected_coil_idcs = np.array([0])         # already chosen: coil 0
    candidate_coil_idcs = np.array([1, 2, 3])  # remaining candidates

    idx_in_candidates, max_min_diff = pynibs.optimization.rowvec_diff(
        candidate_coil_idcs, selected_coil_idcs, efields_diff_mat
    )

    # For each candidate, min distance to selected {0} is just efields_diff_mat[0, cand]
    # -> [1.0, 2.0, 3.0], so best is candidate index 2 (value 3.0)
    assert idx_in_candidates == 2
    assert max_min_diff == pytest.approx(3.0)


def test_rowvec_diff_multiple_selected():
    efields_diff_mat = np.array(
        [
            [0.0, 1.0, 2.0, 1.0],
            [1.0, 0.0, 3.0, 2.0],
            [2.0, 3.0, 0.0, 4.0],
            [1.0, 2.0, 4.0, 0.0],
        ]
    )

    selected_coil_idcs = np.array([0, 1])
    candidate_coil_idcs = np.array([2, 3])

    idx_in_candidates, max_min_diff = pynibs.optimization.rowvec_diff(
        candidate_coil_idcs, selected_coil_idcs, efields_diff_mat
    )

    # Distances from selected {0,1} to candidate 2: [2.0, 3.0] -> min = 2.0
    # Distances from selected {0,1} to candidate 3: [1.0, 2.0] -> min = 1.0
    # So candidate 2 should be chosen (index 0 in candidate_coil_idcs)
    assert idx_in_candidates == 0
    assert max_min_diff == pytest.approx(2.0)


def test_rowvec_diff_single_candidate():
    efields_diff_mat = np.array(
        [
            [0.0, 1.0],
            [1.0, 0.0],
        ]
    )

    selected_coil_idcs = np.array([0])
    candidate_coil_idcs = np.array([1])

    idx_in_candidates, max_min_diff = pynibs.rowvec_diff(
        candidate_coil_idcs, selected_coil_idcs, efields_diff_mat
    )

    assert idx_in_candidates == 0
    assert max_min_diff == pytest.approx(1.0)

def test_optimal_region_dot_basic():
    # 3 coil candidates, 4 voxels
    # Shape e: (n_coils, n_voxels)
    e = np.array(
        [
            [1.0, 0.0, 0.0, 0.0],  # coil 0: aligned with voxel 0
            [0.0, 1.0, 1.0, 0.0],  # coil 1: aligned with voxels 1 & 2
            [0.0, 0.0, 0.0, 1.0],  # coil 2: aligned with voxel 3
        ]
    )
    # fmri peak in voxels 1 & 2 -> coil 1 should be best
    fmri_stats = np.array([0.0, 2.0, 2.0, 0.0])

    best = pynibs.optimal_coilplacement_region(
        e=e.copy(),
        fmri_stats=fmri_stats.copy(),
        best_n=1,
        metric="dot",
        non_negative=False,
        pos_matrices=None,
        res_folder=None,
    )

    assert best.shape == (1,)
    assert best[0] == 1  # coil index 1 should be optimal


def test_optimal_region_cor_basic():
    # 2 coils, 4 voxels
    e = np.array(
        [
            [1.0, 2.0, 3.0, 4.0],  # coil 0 positively correlated with fmri
            [4.0, 3.0, 2.0, 1.0],  # coil 1 negatively correlated with fmri
        ]
    )
    fmri_stats = np.array([1.0, 2.0, 3.0, 4.0])

    best = pynibs.optimal_coilplacement_region(
        e=e.copy(),
        fmri_stats=fmri_stats.copy(),
        best_n=1,
        metric="cor",
        non_negative=False,
        pos_matrices=None,
        res_folder=None,
    )

    # coil 0 correlation > coil 1 correlation, so 0 should be chosen
    assert best.shape == (1,)
    assert best[0] == 0


def test_optimal_region_non_negative_behavior():
    # Include negative fmri values and ensure they are zeroed when non_negative=False
    e = np.eye(3)
    fmri_stats = np.array([1.0, -5.0, 2.0])

    best = pynibs.optimal_coilplacement_region(
        e=e.copy(),
        fmri_stats=fmri_stats.copy(),
        best_n=1,
        metric="dot",
        non_negative=False,
        pos_matrices=None,
        res_folder=None,
    )

    # After zeroing negatives, voxel 2 (value 2) should dominate -> coil index 2
    assert best[0] == 2


def test_optimal_region_invalid_metric_raises():
    e = np.ones((2, 3))
    fmri_stats = np.ones(3)

    with pytest.raises(ValueError):
        pynibs.optimal_coilplacement_region(
            e=e,
            fmri_stats=fmri_stats,
            best_n=1,
            metric="invalid",
            non_negative=False,
            pos_matrices=None,
            res_folder=None,
        )


# def test_optimal_region_io_branch():
#     """
#     Exercise the branch that writes HDF5 and XML and calls pynibs + localite,
#     without actually writing real navigator files.
#     """
#     # small synthetic data
#     n_coils = 2
#     n_vox = 3
#     e = np.ones((n_coils, n_vox))
#     fmri_stats = np.array([1.0, 0.5, 0.2])
#
#     # pos_matrices shape: (3, 4, n_coils) as used inside the function
#     pos_matrices = np.zeros((4, 4, n_coils))
#
#     # Create a temporary directory for all output files
#     with tempfile.TemporaryDirectory() as tmpdirname:
#         res_folder = tmpdirname
#
#         # Dummy localite object living under pynibs.simnibs.utils.nnav.localite
#         # class DummyLocalite:
#         #     def __call__(self):
#         #         return self
#         #
#         #     def write(self, mats, fn, names, overwrite):
#         #         # Just assert shapes and parameters are sensible
#         #         assert mats.shape[2] == n_coils
#         #         assert overwrite is True
#         #         assert fn.endswith(".xml")
#         #
#         # def fake_create_stimsite_from_matsimnibs(h5_fn, mats, overwrite, data, datanames):
#         #     # Check that file path is inside res_folder and best_n matches
#         #     assert os.path.dirname(h5_fn) == res_folder
#         #     assert mats.shape[2] == len(data)
#         #     assert overwrite is True
#         #     assert datanames == "Best coil id"
#
#         # try:
#             # Monkeypatch the functions directly on pynibs
#             # pynibs.simnibs.utils.nnav.localite = fake_localite
#             # pynibs.create_stimsite_from_matsimnibs = fake_create_stimsite_from_matsimnibs
#
#         best = pynibs.optimal_coilplacement_region(
#             e=e.copy(),
#             fmri_stats=fmri_stats.copy(),
#             best_n=2,
#             metric="dot",
#             non_negative=False,
#             pos_matrices=pos_matrices,
#             res_folder=res_folder,
#         )
#         # finally:
#         #     # Restore original functions
#         #     pynibs.simnibs.utils.nnav.localite = original_localite
#         #     pynibs.create_stimsite_from_matsimnibs = original_create_stimsite
#
#         # Check best indices length and range
#         assert best.shape == (2,)
#         assert np.all(best < n_coils)
#
#         # Check that the HDF5 file was created
#         h5_path = os.path.join(res_folder, "opt_coil_pos_dot.hdf5")
#         assert os.path.isfile(h5_path)
#
#         with h5py.File(h5_path, "r") as f:
#             assert "matsimnibs" in f
#             data = f["matsimnibs"][...]
#             # Should have shape (4, 3) after transpose in the function
#             assert data.shape == (4, 4)
