"""
Bloch eigenmode solver/operators

This module contains functions for generating and solving the
 3D Bloch eigenproblem. The approach is to transform the problem
 into the (spatial) fourier domain, transforming the equation

    1/mu * curl(1/eps * curl(H_eigenmode)) = (w/c)^2 H_eigenmode

 into

    conv(1/mu_k, ik x conv(1/eps_k, ik x H_k)) = (w/c)^2 H_k

 where:

  - the `_k` subscript denotes a 3D fourier transformed field
  - each component of `H_k` corresponds to a plane wave with wavevector `k`
  - `x` is the cross product
  - `conv()` denotes convolution

 Since `k` and `H` are orthogonal for each plane wave, we can use each
 `k` to create an orthogonal basis (k, m, n), with `k x m = n`, and
 `|m| = |n| = 1`. The cross products are then simplified as follows:

  - `h` is shorthand for `H_k`
  - `(...)_xyz` denotes the `(x, y, z)` basis
  - `(...)_kmn` denotes the `(k, m, n)` basis
  - `hm` is the component of `h` in the `m` direction, etc.

  We know

    k @ h = kx hx + ky hy + kz hz = 0 = hk
    h = hk + hm + hn = hm + hn
    k = kk + km + kn = kk = |k|

  We can write

    k x h = (ky hz - kz hy,
             kz hx - kx hz,
             kx hy - ky hx)_xyz
          = ((k x h) @ k, (k x h) @ m, (k x h) @ n)_kmn
          = (0, (m x k) @ h, (n x k) @ h)_kmn         # triple product ordering
          = (0, kk (-n @ h), kk (m @ h))_kmn          # (m x k) = -|k| n, etc.
          = |k| (0, -h @ n, h @ m)_kmn

  which gives us a straightforward way to perform the cross product
  while simultaneously transforming into the `_kmn` basis.
  We can also write

    k x h = (km hn - kn hm,
             kn hk - kk hn,
             kk hm - km hk)_kmn
          = (0, -kk hn, kk hm)_kmn
          = (-kk hn)(mx, my, mz)_xyz + (kk hm)(nx, ny, nz)_xyz
          = |k| (hm * (nx, ny, nz)_xyz
               - hn * (mx, my, mz)_xyz)

  which gives us a way to perform the cross product while simultaneously
  trasnforming back into the `_xyz` basis.

 We can also simplify `conv(X_k, Y_k)` as `fftn(X * ifftn(Y_k))`.

 Using these results and storing `H_k` as `h = (hm, hn)`, we have

    e_xyz = fftn(1/eps * ifftn(|k| (hm * n - hn * m)))
    b_mn = |k| (-e_xyz @ n, e_xyz @ m)
    h_mn = fftn(1/mu * ifftn(b_m * m + b_n * n))

 which forms the operator from the left side of the equation.

 We can then use a preconditioned block Rayleigh iteration algorithm, as in
  SG Johnson and JD Joannopoulos, Block-iterative frequency-domain methods
  for Maxwell's equations in a planewave basis, Optics Express 8, 3, 173-190 (2001)
 (similar to that used in MPB) to find the eigenvectors for this operator.

 ===

 Typically you will want to do something like

    recip_lattice = numpy.diag(1/numpy.array(epsilon[0].shape * dx))
    n, v = bloch.eigsolve(5, k0, recip_lattice, epsilon)
    f = numpy.sqrt(-numpy.real(n[0]))
    n_eff = norm(recip_lattice @ k0) / f

    v2e = bloch.hmn_2_exyz(k0, recip_lattice, epsilon)
    e_field = v2e(v[0])

    k, f = find_k(frequency=1/1550,
                  tolerance=(1/1550 - 1/1551),
                  direction=[1, 0, 0],
                  G_matrix=recip_lattice,
                  epsilon=epsilon,
                  band=0)

"""

from typing import Callable, Any, cast, Sequence
import logging
import numpy
from numpy import pi, real, trace
from numpy.fft import fftfreq
from numpy.typing import NDArray, ArrayLike
import scipy                            # type: ignore
import scipy.optimize                   # type: ignore
from scipy.linalg import norm           # type: ignore
import scipy.sparse.linalg as spalg     # type: ignore

from ..fdmath import fdfield_t, cfdfield_t


logger = logging.getLogger(__name__)


try:
    import pyfftw.interfaces.numpy_fft  # type: ignore
    import pyfftw.interfaces            # type: ignore
    import multiprocessing
    logger.info('Using pyfftw')

    pyfftw.interfaces.cache.enable()
    pyfftw.interfaces.cache.set_keepalive_time(3600)
    fftw_args = {
        #'threads': multiprocessing.cpu_count(),
        'overwrite_input': True,
        #'planner_effort': 'FFTW_PATIENT',
        }

    def fftn(*args: Any, **kwargs: Any) -> NDArray[numpy.complex128]:
        return pyfftw.interfaces.numpy_fft.fftn(*args, **kwargs, **fftw_args)

    def ifftn(*args: Any, **kwargs: Any) -> NDArray[numpy.complex128]:
        return pyfftw.interfaces.numpy_fft.ifftn(*args, **kwargs, **fftw_args)

except ImportError:
    from numpy.fft import fftn, ifftn
    logger.info('Using numpy fft')


def generate_kmn(
        k0: ArrayLike,
        G_matrix: ArrayLike,
        shape: Sequence[int],
        ) -> tuple[NDArray[numpy.float64], NDArray[numpy.float64], NDArray[numpy.float64]]:
    """
    Generate a (k, m, n) orthogonal basis for each k-vector in the simulation grid.

    Args:
        k0: [k0x, k0y, k0z], Bloch wavevector, in G basis.
        G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
        shape: [nx, ny, nz] shape of the simulation grid.

    Returns:
        `(|k|, m, n)` where `|k|` has shape `tuple(shape) + (1,)`
            and `m`, `n` have shape `tuple(shape) + (3,)`.
            All are given in the xyz basis (e.g. `|k|[0,0,0] = norm(G_matrix @ k0)`).
    """
    k0 = numpy.array(k0)
    G_matrix = numpy.array(G_matrix, copy=False)

    Gi_grids = numpy.array(numpy.meshgrid(*(fftfreq(n, 1 / n) for n in shape[:3]), indexing='ij'))
    Gi = numpy.moveaxis(Gi_grids, 0, -1)

    k_G = k0[None, None, None, :] - Gi
    k_xyz = numpy.moveaxis(G_matrix @ numpy.moveaxis(k_G, 3, 2), 3, 2)

    m = numpy.broadcast_to([0, 1, 0], tuple(shape[:3]) + (3,)).astype(float)
    n = numpy.broadcast_to([0, 0, 1], tuple(shape[:3]) + (3,)).astype(float)

    xy_non0 = numpy.any(k_xyz[:, :, :, 0:1] != 0, axis=3)
    if numpy.any(xy_non0):
        u = numpy.cross(k_xyz[xy_non0], [0, 0, 1])
        m[xy_non0, :] = u / norm(u, axis=1)[:, None]

    z_non0 = numpy.any(k_xyz != 0, axis=3)
    if numpy.any(z_non0):
        v = numpy.cross(k_xyz[z_non0], m[z_non0])
        n[z_non0, :] = v / norm(v, axis=1)[:, None]

    k_mag = norm(k_xyz, axis=3)[:, :, :, None]
    return k_mag, m, n


def maxwell_operator(
        k0: ArrayLike,
        G_matrix: ArrayLike,
        epsilon: fdfield_t,
        mu: fdfield_t | None = None
        ) -> Callable[[NDArray[numpy.complex128]], NDArray[numpy.complex128]]:
    """
    Generate the Maxwell operator

        conv(1/mu_k, ik x conv(1/eps_k, ik x ___))

    which is the spatial-frequency-space representation of

        1/mu * curl(1/eps * curl(___))

    The operator is a function that acts on a vector h_mn of size `2 * epsilon[0].size`

    See the `meanas.fdfd.bloch` docstring for more information.

    Args:
        k0: Bloch wavevector, `[k0x, k0y, k0z]`.
        G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
        epsilon: Dielectric constant distribution for the simulation.
                 All fields are sampled at cell centers (i.e., NOT Yee-gridded)
        mu: Magnetic permability distribution for the simulation.
            Default None (1 everywhere).

    Returns:
        Function which applies the maxwell operator to h_mn.
    """

    shape = epsilon[0].shape + (1,)
    k_mag, m, n = generate_kmn(k0, G_matrix, shape)

    epsilon = numpy.moveaxis(epsilon, 0, -1)
    if mu is not None:
        mu = numpy.moveaxis(mu, 0, -1)

    def operator(h: NDArray[numpy.complex128]) -> NDArray[numpy.complex128]:
        """
        Maxwell operator for Bloch eigenmode simulation.

        h is complex 2-field in (m, n) basis, vectorized

        Args:
            h: Raveled h_mn; size `2 * epsilon[0].size`.
                Altered in-place.

        Returns:
            Raveled conv(1/mu_k, ik x conv(1/eps_k, ik x h_mn)), returned
            and overwritten in-place of `h`.
        """
        hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]

        #{d,e,h}_xyz fields are complex 3-fields in (1/x, 1/y, 1/z) basis

        # cross product and transform into xyz basis
        d_xyz = (n * hin_m
               - m * hin_n) * k_mag         # noqa: E128

        # divide by epsilon
        temp = ifftn(d_xyz, axes=range(3))   # reuses d_xyz if using pyfftw
        temp /= epsilon
        e_xyz = fftn(temp, axes=range(3))

        # cross product and transform into mn basis
        b_m = numpy.sum(e_xyz * n, axis=3, keepdims=True) * -k_mag
        b_n = numpy.sum(e_xyz * m, axis=3, keepdims=True) * +k_mag

        if mu is None:
            h_m, h_n = b_m, b_n
        else:
            # transform from mn to xyz
            b_xyz = (m * b_m[:, :, :, None]
                   + n * b_n[:, :, :, None])    # noqa: E128

            # divide by mu
            temp = ifftn(b_xyz, axes=range(3))
            temp /= mu
            h_xyz = fftn(temp, axes=range(3))

            # transform back to mn
            h_m = numpy.sum(h_xyz * m, axis=3)
            h_n = numpy.sum(h_xyz * n, axis=3)

        h.shape = (h.size,)
        h = numpy.concatenate((h_m.ravel(), h_n.ravel()), axis=None, out=h)     # ravel and merge
        h.shape = (h.size, 1)
        return h

    return operator


def hmn_2_exyz(
        k0: ArrayLike,
        G_matrix: ArrayLike,
        epsilon: fdfield_t,
        ) -> Callable[[NDArray[numpy.complex128]], cfdfield_t]:
    """
    Generate an operator which converts a vectorized spatial-frequency-space
     `h_mn` into an E-field distribution, i.e.

        ifft(conv(1/eps_k, ik x h_mn))

    The operator is a function that acts on a vector `h_mn` of size `2 * epsilon[0].size`.

    See the `meanas.fdfd.bloch` docstring for more information.

    Args:
        k0: Bloch wavevector, `[k0x, k0y, k0z]`.
        G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
        epsilon: Dielectric constant distribution for the simulation.
                 All fields are sampled at cell centers (i.e., NOT Yee-gridded)

    Returns:
        Function for converting `h_mn` into `E_xyz`
    """
    shape = epsilon[0].shape + (1,)
    epsilon = numpy.moveaxis(epsilon, 0, -1)

    k_mag, m, n = generate_kmn(k0, G_matrix, shape)

    def operator(h: NDArray[numpy.complex128]) -> cfdfield_t:
        hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
        d_xyz = (n * hin_m
               - m * hin_n) * k_mag         # noqa: E128

        # divide by epsilon
        return numpy.array([ei for ei in numpy.moveaxis(ifftn(d_xyz, axes=range(3)) / epsilon, 3, 0)])         # TODO avoid copy

    return operator


def hmn_2_hxyz(
        k0: ArrayLike,
        G_matrix: ArrayLike,
        epsilon: fdfield_t
        ) -> Callable[[NDArray[numpy.complex128]], cfdfield_t]:
    """
    Generate an operator which converts a vectorized spatial-frequency-space
     `h_mn` into an H-field distribution, i.e.

        ifft(h_mn)

    The operator is a function that acts on a vector `h_mn` of size `2 * epsilon[0].size`.

    See the `meanas.fdfd.bloch` docstring for more information.

    Args:
        k0: Bloch wavevector, `[k0x, k0y, k0z]`.
        G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
        epsilon: Dielectric constant distribution for the simulation.
                 Only `epsilon[0].shape` is used.

    Returns:
        Function for converting `h_mn` into `H_xyz`
    """
    shape = epsilon[0].shape + (1,)
    _k_mag, m, n = generate_kmn(k0, G_matrix, shape)

    def operator(h: NDArray[numpy.complex128]) -> cfdfield_t:
        hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
        h_xyz = (m * hin_m
               + n * hin_n)     # noqa: E128
        return numpy.array([ifftn(hi) for hi in numpy.moveaxis(h_xyz, 3, 0)])

    return operator


def inverse_maxwell_operator_approx(
        k0: ArrayLike,
        G_matrix: ArrayLike,
        epsilon: fdfield_t,
        mu: fdfield_t | None = None,
        ) -> Callable[[NDArray[numpy.complex128]], NDArray[numpy.complex128]]:
    """
    Generate an approximate inverse of the Maxwell operator,

        ik x conv(eps_k, ik x conv(mu_k, ___))

     which can be used to improve the speed of ARPACK in shift-invert mode.

    See the `meanas.fdfd.bloch` docstring for more information.

    Args:
        k0: Bloch wavevector, `[k0x, k0y, k0z]`.
        G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
        epsilon: Dielectric constant distribution for the simulation.
                 All fields are sampled at cell centers (i.e., NOT Yee-gridded)
        mu: Magnetic permability distribution for the simulation.
            Default None (1 everywhere).

    Returns:
        Function which applies the approximate inverse of the maxwell operator to `h_mn`.
    """
    shape = epsilon[0].shape + (1,)
    epsilon = numpy.moveaxis(epsilon, 0, -1)
    if mu is not None:
        mu = numpy.moveaxis(mu, 0, -1)

    k_mag, m, n = generate_kmn(k0, G_matrix, shape)

    def operator(h: NDArray[numpy.complex128]) -> NDArray[numpy.complex128]:
        """
        Approximate inverse Maxwell operator for Bloch eigenmode simulation.

        h is complex 2-field in (m, n) basis, vectorized

        Args:
            h: Raveled h_mn; size `2 * epsilon[0].size`.

        Returns:
            Raveled ik x conv(eps_k, ik x conv(mu_k, h_mn))
        """
        hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]

        #{d,e,h}_xyz fields are complex 3-fields in (1/x, 1/y, 1/z) basis

        if mu is None:
            b_m, b_n = hin_m, hin_n
        else:
            # transform from mn to xyz
            h_xyz = (m * hin_m[:, :, :, None]
                   + n * hin_n[:, :, :, None])  # noqa: E128

            # multiply by mu
            temp = ifftn(h_xyz, axes=range(3))
            temp *= mu
            b_xyz = fftn(temp, axes=range(3))

            # transform back to mn
            b_m = numpy.sum(b_xyz * m, axis=3)
            b_n = numpy.sum(b_xyz * n, axis=3)

        # cross product and transform into xyz basis
        e_xyz = (n * b_m
               - m * b_n) / k_mag  # noqa: E128

        # multiply by epsilon
        temp = ifftn(e_xyz, axes=range(3))
        temp *= epsilon
        d_xyz = fftn(temp, axes=range(3))

        # cross product and transform into mn basis   crossinv_t2c
        h_m = numpy.sum(d_xyz * n, axis=3, keepdims=True) / +k_mag
        h_n = numpy.sum(d_xyz * m, axis=3, keepdims=True) / -k_mag

        h.shape = (h.size,)
        h = numpy.concatenate((h_m, h_n), axis=None, out=h)
        h.shape = (h.size, 1)
        return h

    return operator


def find_k(
        frequency: float,
        tolerance: float,
        direction: ArrayLike,
        G_matrix: ArrayLike,
        epsilon: fdfield_t,
        mu: fdfield_t | None = None,
        band: int = 0,
        k_bounds: tuple[float, float] = (0, 0.5),
        k_guess: float | None = None,
        solve_callback: Callable[..., None] | None = None,
        iter_callback: Callable[..., None] | None = None,
        v0: NDArray[numpy.complex128] | None = None,
        ) -> tuple[float, float, NDArray[numpy.complex128], NDArray[numpy.complex128]]:
    """
    Search for a bloch vector that has a given frequency.

    Args:
        frequency: Target frequency.
        tolerance: Target frequency tolerance.
        direction: k-vector direction to search along.
        G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
        epsilon: Dielectric constant distribution for the simulation.
                 All fields are sampled at cell centers (i.e., NOT Yee-gridded)
        mu: Magnetic permability distribution for the simulation.
            Default None (1 everywhere).
        band: Which band to search in. Default 0 (lowest frequency).
        k_bounds: Minimum and maximum values for k. Default (0, 0.5).
        k_guess: Initial value for k.
        solve_callback: TODO
        iter_callback: TODO

    Returns:
        `(k, actual_frequency, eigenvalues, eigenvectors)`
        The found k-vector and its frequency, along with all eigenvalues and eigenvectors.
    """
    direction = numpy.array(direction) / norm(direction)

    k_bounds = tuple(sorted(k_bounds))    # type: ignore    # we know the length already...
    assert len(k_bounds) == 2

    if k_guess is None:
        k_guess = sum(k_bounds) / 2

    n = None
    v = v0

    def get_f(k0_mag: float, band: int = 0) -> float:
        nonlocal n, v
        k0 = direction * k0_mag                         # type: ignore
        n, v = eigsolve(band + 1, k0, G_matrix=G_matrix, epsilon=epsilon, mu=mu, y0=v, callback=iter_callback)
        f = numpy.sqrt(numpy.abs(numpy.real(n[band])))
        if solve_callback:
            solve_callback(k0_mag, n, v, f)
        return f

    res = scipy.optimize.minimize_scalar(
        lambda x: abs(get_f(x, band) - frequency),
        k_guess,
        method='Bounded',
        bounds=k_bounds,
        options={'xatol': abs(tolerance)},
        )

    assert n is not None
    assert v is not None
    return float(res.x * direction), float(res.fun + frequency), n, v


def eigsolve(
        num_modes: int,
        k0: ArrayLike,
        G_matrix: ArrayLike,
        epsilon: fdfield_t,
        mu: fdfield_t | None = None,
        tolerance: float = 1e-7,
        max_iters: int = 10000,
        reset_iters: int = 100,
        y0: ArrayLike | None = None,
        callback: Callable[..., None] | None = None,
        ) -> tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]:
    """
    Find the first (lowest-frequency) num_modes eigenmodes with Bloch wavevector
     k0 of the specified structure.

    Args:
        k0: Bloch wavevector, `[k0x, k0y, k0z]`.
        G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
        epsilon: Dielectric constant distribution for the simulation.
                 All fields are sampled at cell centers (i.e., NOT Yee-gridded)
        mu: Magnetic permability distribution for the simulation.
            Default `None` (1 everywhere).
        tolerance: Solver stops when fractional change in the objective
                   `trace(Z.H @ A @ Z @ inv(Z Z.H))` is smaller than the tolerance
        max_iters: TODO
        reset_iters: TODO
        callback: TODO
        y0: TODO, initial guess

    Returns:
        `(eigenvalues, eigenvectors)` where `eigenvalues[i]` corresponds to the
        vector `eigenvectors[i, :]`
    """
    k0 = numpy.array(k0, copy=False)

    h_size = 2 * epsilon[0].size

    kmag = norm(G_matrix @ k0)

    #
    # Generate the operators
    #
    mop = maxwell_operator(k0=k0, G_matrix=G_matrix, epsilon=epsilon, mu=mu)
    imop = inverse_maxwell_operator_approx(k0=k0, G_matrix=G_matrix, epsilon=epsilon, mu=mu)

    scipy_op = spalg.LinearOperator(dtype=complex, shape=(h_size, h_size), matvec=mop)
    scipy_iop = spalg.LinearOperator(dtype=complex, shape=(h_size, h_size), matvec=imop)

    y_shape = (h_size, num_modes)

    prev_E = 0.0
    d_scale = 1.0
    prev_traceGtKG = 0.0
    prev_theta = 0.5
    D = numpy.zeros(shape=y_shape, dtype=complex)

    Z: NDArray[numpy.complex128]
    if y0 is None:
        Z = numpy.random.rand(*y_shape) + 1j * numpy.random.rand(*y_shape)
    else:
        Z = numpy.array(y0, copy=False).T

    while True:
        Z *= num_modes / norm(Z)
        ZtZ = Z.conj().T @ Z
        try:
            U = numpy.linalg.inv(ZtZ)
        except numpy.linalg.LinAlgError:
            Z = numpy.random.rand(*y_shape) + 1j * numpy.random.rand(*y_shape)
            continue

        trace_U = real(trace(U))
        if trace_U > 1e8 * num_modes:
            Z = Z @ scipy.linalg.sqrtm(U).astype(numpy.complex128).conj().T
            prev_traceGtKG = 0
            continue
        break

    Zt = numpy.empty(Z.shape[::-1], dtype=numpy.complex128)
    AZ = numpy.empty(Z.shape, dtype=numpy.complex128)

    for i in range(max_iters):
        Zt = numpy.conj(Z.T, out=Zt)
        ZtZ = Zt @ Z
        U = numpy.linalg.inv(ZtZ)
        AZ = scipy_op @ Z.copy()
        ZtAZ = Zt @ AZ
        ZtAZU = ZtAZ @ U
        E_signed = real(trace(ZtAZU))
        sgn = numpy.sign(E_signed)
        E = numpy.abs(E_signed)

        # G = AZU projected onto the space orthonormal to Z via (1 - ZUZt)
        G = (AZ @ U - Z @ U @ ZtAZU) * sgn

        if i > 0 and abs(E - prev_E) < tolerance * 0.5 * (E + prev_E + 1e-7):
            logger.info(
                'Optimization succeded: '
                f'[change in trace] {abs(E - prev_E)} - 5e-8 '
                f'< {tolerance} [tolerance] * {(E + prev_E) / 2} [value of trace]'
                )
            break

        KG = scipy_iop @ G          # Preconditioned steepest descent direction
        traceGtKG = _rtrace_AtB(G, KG)

        if prev_traceGtKG == 0 or i % reset_iters == 0:
            logger.info('CG reset')
            gamma = 0.0
        else:
            gamma = traceGtKG / prev_traceGtKG
        prev_traceGtKG = traceGtKG

        D = gamma / d_scale * D + KG
        d_scale = num_modes / norm(D)
        D *= d_scale

        # Now know the direction (D), but need to find how far to go (alpha)
        # We are still minimizing E = tr((Z + alpha D)t A (Z + alpha D) U')
        #         = tr(ZtAZU' + alpha DtAZU' + alpha ZtADU' + alpha**2 DtADU')
        #         = tr((ZtAZ + 2 alpha sym(DtAZ) + alpha**2 DtAD) U')
        #         = tr(R U')
        #         = tr(R Qi) = tr(R inv(Q))

        # where
        #   R = ZtAZ + 2 alpha sym(DtAZ) + alpha**2 DtAD
        #
        #   Q = (Z + alpha D)t (Z + alpha D)
        #     = inv(ZtZ + alpha ZtD + alpha DtZ + alpha**2 DtD)
        #
        #   Qi = inv(Q) = U'

        AD = scipy_op @ D.copy()
        DtD = D.conj().T @ D
        DtAD = D.conj().T @ AD

        symZtD = _symmetrize(Zt @ D)
        symZtAD = _symmetrize(Zt @ AD)

        Qi_memo: list[float | None] = [None, None]

        def Qi_func(theta: float) -> float:
            nonlocal Qi_memo
            if Qi_memo[0] == theta:
                return cast(float, Qi_memo[1])

            c = numpy.cos(theta)
            s = numpy.sin(theta)
            Q = c * c * ZtZ + s * s * DtD + 2 * s * c * symZtD
            try:
                Qi = numpy.linalg.inv(Q)
            except numpy.linalg.LinAlgError:
                logger.info('taylor Qi')
                # if c or s small, taylor expand
                if c < 1e-4 * s and c != 0:
                    DtDi = numpy.linalg.inv(DtD)
                    Qi = DtDi / (s * s) - 2 * c / (s * s * s) * (DtDi @ (DtDi @ symZtD).conj().T)
                elif s < 1e-4 * c and s != 0:
                    ZtZi = numpy.linalg.inv(ZtZ)
                    Qi = ZtZi / (c * c) - 2 * s / (c * c * c) * (ZtZi @ (ZtZi @ symZtD).conj().T)
                else:
                    raise Exception('Inexplicable singularity in trace_func')
            Qi_memo[0] = theta
            Qi_memo[1] = cast(float, Qi)
            return cast(float, Qi)

        def trace_func(theta: float) -> float:
            c = numpy.cos(theta)
            s = numpy.sin(theta)
            Qi = Qi_func(theta)
            R = c * c * ZtAZ + s * s * DtAD + 2 * s * c * symZtAD
            trace = _rtrace_AtB(R, Qi)
            return numpy.abs(trace)

        if False:
            def trace_deriv(theta):
                Qi = Qi_func(theta)
                c2 = numpy.cos(2 * theta)
                s2 = numpy.sin(2 * theta)
                F = -0.5*s2 *  (ZtAZ - DtAD) + c2 * symZtAD
                trace_deriv = _rtrace_AtB(Qi, F)

                G = Qi @ F.conj().T @ Qi.conj().T
                H = -0.5*s2 * (ZtZ - DtD) + c2 * symZtD
                trace_deriv -= _rtrace_AtB(G, H)

                trace_deriv *= 2
                return trace_deriv * sgn

            U_sZtD = U @ symZtD

            dE = 2.0 * (_rtrace_AtB(U, symZtAD) -
                        _rtrace_AtB(ZtAZU, U_sZtD))

            d2E = 2 * (_rtrace_AtB(U, DtAD) -
                       _rtrace_AtB(ZtAZU, U @ (DtD - 4 * symZtD @ U_sZtD)) -
                   4 * _rtrace_AtB(U, symZtAD @ U_sZtD))

            # Newton-Raphson to find a root of the first derivative:
            theta = -dE / d2E

            if d2E < 0 or abs(theta) >= pi:
                theta = -abs(prev_theta) * numpy.sign(dE)

            # theta, new_E, new_dE = linmin(theta, E, dE, 0.1, min(tolerance, 1e-6), 1e-14, 0, -numpy.sign(dE) * K_PI, trace_func)
            theta, n, _, new_E, _, _new_dE = scipy.optimize.line_search(
                trace_func,
                trace_deriv,
                xk=theta,
                pk=numpy.ones((1, 1)),
                gfk=dE,
                old_fval=E,
                c1=min(tolerance, 1e-6),
                c2=0.1,
                amax=pi,
                )

        result = scipy.optimize.minimize_scalar(trace_func, bounds=(0, pi), tol=tolerance)
        new_E = result.fun
        theta = result.x

        improvement = numpy.abs(E - new_E) * 2 / numpy.abs(E + new_E)
        logger.info(f'linmin improvement {improvement}')
        Z *= numpy.cos(theta)
        Z += D * numpy.sin(theta)

        prev_theta = theta
        prev_E = E

        if callback:
            callback()

    #
    # Recover eigenvectors from Z
    #
    U = numpy.linalg.inv(ZtZ)
    Y = Z @ scipy.linalg.sqrtm(U).astype(numpy.complex128)
    W = Y.conj().T @ (scipy_op @ Y.copy())

    eigvals, W_eigvecs = numpy.linalg.eig(W)
    eigvecs = Y @ W_eigvecs

    for i in range(len(eigvals)):
        v = eigvecs[:, i]
        n = eigvals[i]
        v /= norm(v)
        Av = (scipy_op @ v.copy())[:, 0]
        eigness = norm(Av - (v.conj() @ Av) * v)
        f = numpy.sqrt(-numpy.real(n))
        df = numpy.sqrt(-numpy.real(n) + eigness)
        neff_err = kmag * (1 / df - 1 / f)
        logger.info(f'eigness {i}: {eigness}\n neff_err: {neff_err}')

    order = numpy.argsort(numpy.abs(eigvals))
    return eigvals[order], eigvecs.T[order]


'''
def linmin(x_guess, f0, df0, x_max, f_tol=0.1, df_tol=min(tolerance, 1e-6), x_tol=1e-14, x_min=0, linmin_func):
    if df0 > 0:
        x0, f0, df0 = linmin(-x_guess, f0, -df0, -x_max, f_tol, df_tol, x_tol, -x_min, lambda q, dq: -linmin_func(q, dq))
        return -x0, f0, -df0
    elif df0 == 0:
        return 0, f0, df0
    else:
        x = x_guess
        fx = f0
        dfx = df0

        isave = numpy.zeros((2,), numpy.intc)
        dsave = numpy.zeros((13,), float)

        x, fx, dfx, task = minpack2.dsrch(x, fx, dfx, f_tol, df_tol, x_tol, task,
                                          x_min, x_max, isave, dsave)
        for i in range(int(1e6)):
            if task != 'F':
                logging.info('search converged in {} iterations'.format(i))
                break
            fx = f(x, dfx)
            x, fx, dfx, task = minpack2.dsrch(x, fx, dfx, f_tol, df_tol, x_tol, task,
                                              x_min, x_max, isave, dsave)

        return x, fx, dfx
'''

def _rtrace_AtB(
        A: NDArray[numpy.complex128],
        B: NDArray[numpy.complex128] | float,
        ) -> float:
    return real(numpy.sum(A.conj() * B))

def _symmetrize(A: NDArray[numpy.complex128]) -> NDArray[numpy.complex128]:
    return (A + A.conj().T) * 0.5

