# test_multichannel_focality.py

import numpy as np
import pynibs

def test_calc_ll_shapes_and_values():
    # simple 2-coil, 3-element lead-field (flattened: 3*3 rows)
    l = np.array(
        [
            [1.0, 0.0],
            [0.0, 1.0],
            [1.0, 1.0],
            [2.0, 0.0],
            [0.0, 2.0],
            [1.0, -1.0],
        ]
    )
    ll = pynibs.optimization.calc_ll(l)

    # shape: [n_coils, n_coils]
    assert ll.shape == (2, 2)

    # ll = l.T @ l, check against explicit computation
    ll_manual = l.T @ l
    assert np.allclose(ll, ll_manual)


def test_calc_e_shape_and_linear_mapping():
    # 2 coils, 3 elements -> l has 3*3 rows
    l = np.array(
        [
            [1.0, 0.0],
            [0.0, 1.0],
            [1.0, 1.0],
            [2.0, 0.0],
            [0.0, 2.0],
            [1.0, -1.0],
            [0.5, 0.5],
            [1.0, 0.0],
            [0.0, 1.0],
        ]
    )
    j = np.array([1.0, -1.0])  # currents for 2 coils

    e = pynibs.optimization.calc_e(l, j)

    # There are n_k = len(l)/3 elements, each with 3 components
    n_k = int(len(l) / 3)
    assert e.shape == (n_k, 3)

    # Check that flatten(e) == l @ j
    assert np.allclose(e.reshape(-1), l @ j)


def test_get_focality_score_quadratic_form():
    # 3 coils, simple symmetric positive semidefinite ll
    j = np.array([1.0, 2.0, -1.0])
    ll = np.array(
        [
            [2.0, 0.5, 0.0],
            [0.5, 1.0, 0.3],
            [0.0, 0.3, 0.5],
        ]
    )

    eps = pynibs.optimization.get_focality_score(j, ll)

    # eps = j^T * ll * j
    eps_manual = j @ ll @ j.T
    assert np.isclose(eps, eps_manual)


def test_goal_uses_focality_score_and_scales_by_target_stim():
    # minimal setup: 3 coils so hyperellipsoid part is used directly
    n_coils = 3
    # lead-field l: 6 elements (just arbitrary; only ll is needed for goal)
    l = np.random.randn(6, n_coils)
    ll = pynibs.optimization.calc_ll(l)

    # fake SVD outputs (u,s) consistent with ll
    # use identity for u to test mapping j = u @ x more transparently
    u = np.eye(n_coils)
    s = np.array([1.0, 2.0, 3.0])

    phi = np.array([0.1, 0.2, 0.3])  # length >= 3

    target_stim = 50.0
    eps, j = pynibs.optimization.goal(phi, l, ll, target_stim=target_stim, u=u, s=s)

    # j has length n_coils and is scaled by target_stim
    assert j.shape == (n_coils,)
    assert np.allclose(j / target_stim, u @ (1.0 / np.sqrt(s[0:3]) * np.array(
        [
            np.cos(phi[0]),
            np.sin(phi[0]) * np.cos(phi[1]),
            np.sin(phi[0]) * np.sin(phi[1]),
        ]
    )))

    # eps * n_elm should match focality_score(j,ll) where n_elm = l.shape[0]/100
    n_elm = l.shape[0] / 100.0
    expected_eps = pynibs.optimization.get_focality_score(j, ll) / n_elm
    assert np.isclose(eps, expected_eps)

def test_goal_handles_small_singular_values():
    # construct s with very small entries to exercise thresholding
    n_coils = 3
    l = np.random.randn(6, n_coils)
    ll = pynibs.optimization.calc_ll(l)
    u = np.eye(n_coils)
    s = np.array([1e-12, 1e-11, 1e-9])  # all below threshold

    phi = np.array([0.0, 0.0, 0.0])

    eps, j = pynibs.optimization.goal(phi, l, ll, target_stim=1.0, u=u, s=s)

    # after threshold, s < 1e-10 set to 1e-10
    s_thresholded = np.array([1e-10, 1e-10, 1e-9])
    e = 1.0 / np.sqrt(s_thresholded[0:3])

    x_expected = np.zeros(n_coils)
    x_expected[0] = e[0] * np.cos(phi[0])
    x_expected[1] = e[1] * np.sin(phi[0]) * np.cos(phi[1])
    x_expected[2] = e[2] * np.sin(phi[0]) * np.sin(phi[1])

    # check mapping through u
    assert np.allclose(j, (u @ x_expected) * 1.0)
    
def test_calculate_triangle_area_simple():
    # Single triangle in xy-plane with right angle: area = 0.5 * 1 * 1 = 0.5
    p1 = np.array([[0.0, 0.0, 0.0]])
    p2 = np.array([[1.0, 0.0, 0.0]])
    p3 = np.array([[0.0, 1.0, 0.0]])

    area = pynibs.optimization.calculate_triangle_area(p1, p2, p3)

    assert area.shape == (1,)
    assert np.isclose(area[0], 0.5)


def test_constraint_min_current_satisfied_and_violated():
    n_stim = 4
    # channels 2 and 3 are constrained by default
    # Construct x so that all entries on constrained channels exceed min_current
    min_current = 0.7
    n_channel = 5
    x_good = np.zeros(n_stim * n_channel)
    # channel indices (2,3) -> slices [2,3] in channel axis
    x_good[n_stim * 2 : n_stim * 3] = min_current + 0.1
    x_good[n_stim * 3 : n_stim * 4] = min_current + 0.1

    res_good = pynibs.optimization.constraint_min_current(
        x_good,
        n_stim=n_stim,
        channels=(2, 3),
        min_current=min_current,
        frac=0.95,
        factor=1.0,
    )

    # At most tiny numerical deviations; must be >= 0 when satisfied
    assert res_good >= 0.0

    # Now make many entries below threshold so fraction is violated
    x_bad = np.zeros_like(x_good)
    x_bad[n_stim * 2 : n_stim * 3] = 0.1
    x_bad[n_stim * 3 : n_stim * 4] = 0.1

    res_bad = pynibs.optimization.constraint_min_current(
        x_bad,
        n_stim=n_stim,
        channels=(2, 3),
        min_current=min_current,
        frac=0.95,
        factor=1.0,
    )

    assert res_bad < 0.0


def test_get_current_init_shapes_and_bounds():
    n_channel = 3
    n_stim = 5
    bounds = [(0.0, 1.0), (0.0, 2.0), (-1.0, 1.0)]
    only_pos = [True, True, False]

    x0 = pynibs.optimization.get_current_init(n_channel=n_channel, n_stim=n_stim, bounds=bounds, only_pos=only_pos, seed=42)

    # Shape n_channel * n_stim
    assert x0.shape == (n_channel * n_stim,)

    # Check per-channel ranges and sign behavior
    for ch in range(n_channel):
        seg = x0[n_stim * ch : n_stim * (ch + 1)]
        low, high = min(bounds[ch]), max(bounds[ch])
        assert np.all(seg >= low)
        assert np.all(seg <= high)
        if only_pos[ch]:
            assert np.all(seg >= 0.0)


def test_get_score_calls_get_score_raw_and_normalizes(monkeypatch):
    # small synthetic setup
    n_ele = 2
    n_channel = 2
    n_stim = 2
    x = [1,-1,1,-1]
    e = np.array([[0,0,-1,0,0,-1],[0,0,1,0,0,1]]).T
    score = pynibs.optimization.get_score(
        x=x,
        e=e,
        n_stim=n_stim,
        n_ele=n_ele,
        n_channel=n_channel,
        x_opt=None,
        opt_target="stims",
        fun=None,
        fun_args=None,
    )
    assert score == 4000


def test_get_score_raw_elms_mode_and_fun_scaling(monkeypatch):
    n_ele = 3
    n_channel = 2
    n_stim = 2

    # e: (n_ele*3, n_channel)
    e = np.eye(n_ele * 3, n_channel)
    x = np.array([1.0, -1.0, 0.5, 0.5])  # shape (n_channel * n_stim,)
    x_opt = None

    # simple scaling function to check that fun is applied
    def scale_fun(e_mag, factor):
        return e_mag * factor

    fun_args = {"factor": 2.0}

    # force elms mode, and ensure no NaN branch is taken by using non-degenerate data
    score = pynibs.optimization.get_score_raw(
        x=x,
        e=e,
        n_stim=n_stim,
        n_ele=n_ele,
        n_channel=n_channel,
        x_opt=x_opt,
        opt_target="elms",
        fun=scale_fun,
        fun_args=fun_args,
        res_scale=1.0,
        verbose=False,
    )

    # Just check it returns a finite scalar and is non-negative
    assert np.isscalar(score)
    assert not np.isnan(score)
    assert score >= 0.0


def test_get_score_raw_stims_mode_and_nan_handling():
    n_ele = 3
    n_channel = 2
    n_stim = 2

    # e: (n_ele*3, n_channel)
    e = np.eye(n_ele * 3, n_channel)
    x = np.array([np.nan, 0.0, 0.0, 0.0])
    x_opt = None

    # This will generate zero e-fields, correlation can be NaN; function should replace NaN with n_stim**2
    score = pynibs.optimization.get_score_raw(
        x=x,
        e=e,
        n_stim=n_stim,
        n_ele=n_ele,
        n_channel=n_channel,
        x_opt=x_opt,
        opt_target="stims",
        fun=None,
        fun_args=None,
        res_scale=1.0,
        verbose=False,
    )

    # In NaN case, r_sum should become n_stim**2
    assert np.isclose(score, n_stim ** 2)


def test_get_score_raw_single_channel_basic_and_with_x_opt():
    # e: (n_ele*3, n_channels) but here used as (n_ele*3, n_channels)
    n_ele = 4
    n_channels = 3
    e = np.random.randn(n_ele * 3, n_channels)

    # x selects row indices; use small integers
    x = np.array([0, 2, 5], dtype=int)

    score1 = pynibs.optimization.get_score_raw_single_channel(x=x, e=e, x_opt=None)
    assert np.isscalar(score1)
    assert not np.isnan(score1)

    # With x_opt, it should prepend indices
    x_opt = np.array([1, 3], dtype=int)
    score2 = pynibs.optimization.get_score_raw_single_channel(x=x, e=e, x_opt=x_opt)
    assert np.isscalar(score2)
    assert not np.isnan(score2)

    # Different selection (x_opt) should generally change the score
    assert score1 != score2


def test_sphere_returns_points_and_triangle_indices():
    # Simple square made of two triangles in the xy-plane
    points = np.array(
        [
            [0.0, 0.0, 0.0],  # 0
            [1.0, 0.0, 0.0],  # 1
            [1.0, 1.0, 0.0],  # 2
            [0.0, 1.0, 0.0],  # 3
        ]
    )
    # Two triangles: (0,1,2) and (0,2,3)
    connections = np.array([[0, 1, 2], [0, 2, 3]])

    # Center near the middle, radius small: should hit all 4 points and both tris
    center = np.array([0.5, 0.5, 0.0])
    radius = 1.0

    valid_p, idx_c = pynibs.optimization.sphere(points, connections, radius, center)

    # All points within radius
    assert sorted(valid_p) == [0, 1, 2, 3]
    # Both triangles should be included
    assert sorted(idx_c) == [0, 1]


def test_sphere_with_smaller_radius_selects_subset():
    points = np.array(
        [
            [0.0, 0.0, 0.0],  # 0
            [2.0, 0.0, 0.0],  # 1
            [2.0, 2.0, 0.0],  # 2
            [0.0, 2.0, 0.0],  # 3
        ]
    )
    connections = np.array([[0, 1, 2], [1, 2, 3]])

    # Center near point 0, radius small: only point 0 should be inside
    center = np.array([0.0, 0.0, 0.0])
    radius = 0.9

    valid_p, idx_c = pynibs.optimization.sphere(points, connections, radius, center)

    assert valid_p == [0]
    # Only triangles that contain point 0 (triangle 0)
    assert idx_c == [0]


def test_target_to_max_ratio_scalar_and_vector_e():
    # 1D e-field, k-based target
    e = np.array([0.5, 1.0, 2.0, 1.5])
    k = 2  # target at 2.0
    t2m = pynibs.optimization.target_to_max_ratio(e, k=k, e_target=None)

    # e_target = 2.0, e.max() = 2.0 => ratio 1
    assert np.isclose(t2m, 1.0)

    # 2D vector e-field, with explicit e_target
    e_vec = np.array(
        [
            [1.0, 0.0],
            [0.0, 2.0],
            [0.0, 0.5],
        ]
    )
    # Use e_target explicitly (e.g., 2.0), max norm is 2.0 from [0,2]
    t2m_vec = pynibs.optimization.target_to_max_ratio(e_vec, k=None, e_target=2.0)
    assert np.isclose(t2m_vec, 1.0)

    # Lower e_target should yield ratio < 1
    t2m_lower = pynibs.optimization.target_to_max_ratio(e_vec, k=None, e_target=1.0)
    assert 0.0 < t2m_lower < 1.0


def test_target_to_max_ratio_assertion():
    e = np.array([1.0, 2.0, 3.0])
    # Both k and e_target set should raise an assertion
    try:
        pynibs.optimization.target_to_max_ratio(e, k=1, e_target=2.0)
    except AssertionError:
        pass
    else:
        assert False, "Expected AssertionError when both k and e_target are provided"


def test_get_maximality_currents_enumerates_sign_combinations():
    # Simple 2-channel case
    # e: (n_chans, 3) leadfield; choose vectors so that [1,1] direction gives larger norm than [1,-1]
    e = np.array(
        [
            [1.0, 0.0, 0.0],   # chan 0
            [0.5, 1.0, 0.0],   # chan 1
        ]
    )
    current_constraint = 10.0

    j_max, e_max = pynibs.optimization.get_maximality_currents(e, current_constraint=current_constraint)

    # j_max is ±1 per channel
    assert j_max.shape == (2,)
    assert set(np.unique(j_max)).issubset({-1, 1})

    # For 2 channels, combinations are [1,1] and [1,-1] (the function halves the set)
    # Check that the chosen combination yields the reported maximum
    combs = np.array([[1, 1], [1, -1]])
    e_max_candidates = []
    for comb in combs:
        a = comb[:, np.newaxis] * current_constraint * e
        e_norm = np.linalg.norm(np.sum(a, axis=0))
        e_max_candidates.append(e_norm)
    e_max_candidates = np.array(e_max_candidates)

    assert np.isclose(e_max, e_max_candidates.max())
