from fractions import Fraction

import numpy as np
import pytest

from dynasor.modes.tools import (
    as_fraction,
    det,
    factor_analysis,
    group_eigvals,
    inv,
    make_table,
    symmetrize_eigenvectors,
    trace,
)


# ------------------------------------------------------------------ trace ---

def test_trace_identity():
    A = np.eye(3)
    assert trace(A) == pytest.approx(3.0)


def test_trace_arbitrary():
    A = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=float)
    assert trace(A) == pytest.approx(15.0)


# ------------------------------------------------------------------- det ----

def test_det_2x2():
    A = np.array([[3, 4], [2, 5]], dtype=float)
    assert det(A) == pytest.approx(np.linalg.det(A))


def test_det_3x3_identity():
    A = np.eye(3)
    assert det(A) == pytest.approx(1.0)


def test_det_3x3_arbitrary():
    np.random.seed(0)
    A = np.array([[2, 1, 0], [0, 3, 1], [1, 0, 4]], dtype=float)
    assert det(A) == pytest.approx(np.linalg.det(A))


def test_det_singular():
    A = np.zeros((3, 3))
    assert det(A) == pytest.approx(0.0)


# ------------------------------------------------------------------- inv ----

def test_inv_returns_fraction_by_default():
    A = np.array([[2, 0, 0], [0, 3, 0], [0, 0, 6]], dtype=int)
    A_inv = inv(A)
    assert A_inv.dtype == object
    assert isinstance(A_inv[0, 0], Fraction)


def test_inv_returns_float_when_requested():
    A = np.array([[2, 0, 0], [0, 4, 0], [0, 0, 1]], dtype=int)
    A_inv = inv(A, as_fraction=False)
    assert np.allclose(A_inv, np.linalg.inv(A))


def test_inv_correctness():
    A = np.array([[1, 2, 3], [0, 1, 4], [5, 6, 0]], dtype=int)
    A_inv = inv(A, as_fraction=False)
    product = A @ A_inv
    assert np.allclose(product, np.eye(3))


# -------------------------------------------------------------- group_eigvals

def test_group_eigvals_no_degeneracy():
    vals = [1.0, 2.0, 3.0]
    groups = group_eigvals(vals)
    assert groups == [[0], [1], [2]]


def test_group_eigvals_all_degenerate():
    vals = [1.0, 1.0, 1.0]
    groups = group_eigvals(vals)
    assert groups == [[0, 1, 2]]


def test_group_eigvals_partial_degeneracy():
    vals = [1.0, 1.0, 2.0, 3.0, 3.0]
    groups = group_eigvals(vals)
    assert groups == [[0, 1], [2], [3, 4]]


def test_group_eigvals_unsorted_raises():
    with pytest.raises(AssertionError):
        group_eigvals([3.0, 1.0, 2.0])


# ----------------------------------------------------------- factor_analysis -

def test_factor_analysis_returns_unitary():
    np.random.seed(7)
    L = np.random.normal(size=(6, 3))
    R = factor_analysis(L, iterations=50)
    assert R.shape == (3, 3)
    # R should be orthogonal (rotation)
    assert np.allclose(R @ R.T, np.eye(3), atol=1e-10)


def test_factor_analysis_quartimax():
    np.random.seed(7)
    L = np.random.normal(size=(6, 3))
    R = factor_analysis(L, iterations=50, method='quartimax')
    assert np.allclose(R @ R.T, np.eye(3), atol=1e-10)


def test_factor_analysis_numeric_gamma():
    np.random.seed(7)
    L = np.random.normal(size=(6, 3))
    R = factor_analysis(L, iterations=50, method=0.5)
    assert np.allclose(R @ R.T, np.eye(3), atol=1e-10)


# ------------------------------------------------------- symmetrize_eigvecs --

def test_symmetrize_eigenvectors_shape():
    np.random.seed(1)
    eigvecs = np.random.normal(size=(3, 4, 3))
    result = symmetrize_eigenvectors(eigvecs)
    assert result.shape == eigvecs.shape


def test_symmetrize_eigenvectors_with_cell():
    np.random.seed(2)
    eigvecs = np.random.normal(size=(2, 2, 3))
    cell = np.diag([3.0, 3.0, 3.0])
    result = symmetrize_eigenvectors(eigvecs, cell=cell)
    assert result.shape == eigvecs.shape


# ---------------------------------------------------------------- as_fraction -

def test_as_fraction_scalar():
    result = as_fraction(3)
    assert isinstance(result, Fraction)
    assert result == Fraction(3)


def test_as_fraction_array():
    arr = np.array([1, 2, 3])
    result = as_fraction(arr)
    assert result.shape == arr.shape
    assert all(isinstance(x, Fraction) for x in result)


def test_as_fraction_tuple():
    t = (1, 2, 3)
    result = list(as_fraction(t))
    assert all(isinstance(x, Fraction) for x in result)


def test_as_fraction_list():
    lst = [4, 5, 6]
    result = as_fraction(lst)
    assert all(isinstance(x, Fraction) for x in result)


# ------------------------------------------------------------------- inv (negative det branch) ---

def test_inv_negative_determinant():
    # Odd permutation matrix has det = -1, triggering the detx2 < 0 branch in inv()
    A = np.array([[0, 1, 0], [1, 0, 0], [0, 0, 1]], dtype=int)
    assert np.linalg.det(A) < 0
    A_inv = inv(A, as_fraction=False)
    assert np.allclose(A_inv, np.linalg.inv(A))
    assert np.allclose(A @ A_inv, np.eye(3))


# ----------------------------------------------------------------- make_table ---

def test_make_table():
    M = np.array([[1.0, 2.5], [3.0, 4.75]])
    result = make_table(M)
    assert '1.00' in result
    assert '4.75' in result
    # Each row on its own line
    lines = result.splitlines()
    assert len(lines) == 2
