"""Public helper module.

Conceptually, this module groups linear projection and orthogonal-subspace
operations used by coding-direction and movement analyses.

It exists as a separate unit so matrix geometry, subspace decomposition, and
trial-matrix orientation fixes can stay independent from higher-level
decoding or plotting workflows.

It connects normalized neural activity matrices to movement-subspace and
coding-direction pipelines.
"""

from __future__ import annotations

from typing import Dict, Tuple

import numpy as np

from .stats_utils import zscore_matlab_cols

__all__ = [
    "ensure_trials_by_cells",
    "gram_schmidt_columns",
    "normalize_vec",
    "orthogonal_subspaces_pg",
    "pg_decompose_python",
    "project_trials",
    "qr_orth",
]


def normalize_vec(v: np.ndarray) -> np.ndarray:
    """Return an L2-normalized vector.

    Parameters
    ----------
    v:
        One-dimensional vector-like payload.

    Returns
    -------
    np.ndarray
        Normalized vector. Invalid or zero-norm inputs return zeros with the
        same shape.
    """

    # MATLAB works in double precision; keep float64 so coding-direction
    # magnitudes match Code_M.
    v = np.asarray(v, dtype=np.float64).reshape(-1)
    nrm = np.linalg.norm(v)
    if (not np.isfinite(nrm)) or (nrm == 0):
        return np.zeros_like(v)
    return v / nrm



def gram_schmidt_columns(mat: np.ndarray) -> np.ndarray:
    """Orthogonalize matrix columns exactly like MATLAB ``fn_gram_schmidt_process``.

    Code_M's ``fn_gram_schmidt_process.m`` performs a *non-normalizing* modified
    Gram-Schmidt: NaNs are zeroed, then for each column ``ii`` its projection is
    subtracted from every later column ``jj`` using ``proj(u,v)=(u·v/u·u)·u``.
    The per-column ``v/norm(v)`` normalization is commented out in the source, so
    the returned columns are orthogonal but NOT unit length. Re-normalizing here
    (as classical Gram-Schmidt would) rescales the downstream coding-direction
    projections, so preserve the orthogonal-but-unnormalized behaviour and stay
    in float64.
    """

    v = np.asarray(mat, dtype=np.float64).copy()
    v[np.isnan(v)] = 0.0
    n_cols = v.shape[1]
    for ii in range(n_cols):
        u = v[:, ii]
        uu = float(np.dot(u, u))
        if uu == 0.0:
            continue
        for jj in range(ii + 1, n_cols):
            coeff = float(np.dot(u, v[:, jj])) / uu
            v[:, jj] = v[:, jj] - coeff * u
    return v



def project_trials(currsig: np.ndarray, cd_vec: np.ndarray) -> np.ndarray:
    """Project trial-by-time neural activity on one coding-direction vector.

    Parameters
    ----------
    currsig:
        Neural activity tensor with shape `(time, trials, neurons)`.
    cd_vec:
        Projection vector with shape `(neurons,)`.

    Returns
    -------
    np.ndarray
        Projected activity with shape `(time, trials)`.
    """

    n_time, n_trials, _ = currsig.shape
    out = np.zeros((n_time, n_trials), dtype=np.float32)
    # Project each trial independently so the output preserves the original
    # `(time, trials)` structure expected by downstream plotting code.
    for itrial in range(n_trials):
        curr = currsig[:, itrial, :]
        out[:, itrial] = curr @ cd_vec
    return out



def qr_orth(X: np.ndarray) -> np.ndarray:
    """Return a QR-orthonormalized basis with a deterministic sign convention."""

    Q, R = np.linalg.qr(np.asarray(X, dtype=np.float64))
    signs = np.sign(np.diag(R))
    signs[signs == 0] = 1.0
    return Q * signs[np.newaxis, :]



def _top_eigvals_sym(C: np.ndarray, k: int) -> np.ndarray:
    vals = np.linalg.eigvalsh(np.asarray(C, dtype=np.float64))
    vals = np.sort(vals)[::-1]
    return vals[:k]



def orthogonal_subspaces_pg(
    C1: np.ndarray,
    d1: int,
    C2: np.ndarray,
    d2: int,
    alpha: float = 0.0,
    seed: int = 101,
    max_iter: int = 250,
    tol: float = 1e-6,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Solve the orthogonal-subspace decomposition problem.

    Parameters
    ----------
    C1, C2:
        Square covariance matrices with the same shape.
    d1, d2:
        Target dimensions of the two orthogonal subspaces.
    alpha:
        Optional regularization weight in the objective.
    seed:
        Seed used for the initial Stiefel point.

    Returns
    -------
    np.ndarray, np.ndarray, np.ndarray
        Joint basis matrix `Q` and the two selector matrices `P1`, `P2`.
    """

    del max_iter, tol

    C1 = np.asarray(C1, dtype=np.float64)
    C2 = np.asarray(C2, dtype=np.float64)

    if C1.shape != C2.shape or C1.shape[0] != C1.shape[1]:
        raise ValueError("C1 and C2 must be same-size square covariance matrices.")

    n = C1.shape[0]
    if n < (d1 + d2):
        raise ValueError(f"Not enough neurons for d1+d2: n={n}, d1+d2={d1 + d2}")

    dmax = max(d1, d2)
    eigvals1 = _top_eigvals_sym(C1, dmax)
    eigvals2 = _top_eigvals_sym(C2, dmax)

    s1 = float(np.sum(eigvals1[:d1]))
    s2 = float(np.sum(eigvals2[:d2]))
    if (not np.isfinite(s1)) or s1 <= 0:
        s1 = 1.0
    if (not np.isfinite(s2)) or s2 <= 0:
        s2 = 1.0

    P1 = np.vstack([np.eye(d1, dtype=np.float64), np.zeros((d2, d1), dtype=np.float64)])
    P2 = np.vstack([np.zeros((d1, d2), dtype=np.float64), np.eye(d2, dtype=np.float64)])

    try:
        import autograd.numpy as anp
        from pymanopt import Problem, function
        from pymanopt.manifolds import Stiefel
        from pymanopt.optimizers import TrustRegions
    except Exception as exc:
        raise RuntimeError(
            "orthogonal_subspaces_pg requires pymanopt and autograd to match the source trust-region solver"
        ) from exc

    manifold = Stiefel(n, d1 + d2)
    P1_ag = anp.asarray(P1)
    P2_ag = anp.asarray(P2)
    C1_ag = anp.asarray(C1)
    C2_ag = anp.asarray(C2)
    s1_ag = anp.asarray(s1)
    s2_ag = anp.asarray(s2)
    alpha_ag = anp.asarray(float(alpha))

    @function.autograd(manifold)
    def cost(Qm):
        Q1 = Qm @ P1_ag
        Q2 = Qm @ P2_ag
        return (
            -0.5 * anp.trace(Q1.T @ C1_ag @ Q1) / s1_ag
            - 0.5 * anp.trace(Q2.T @ C2_ag @ Q2) / s2_ag
            - alpha_ag * (anp.linalg.norm(Q1) + anp.linalg.norm(Q2))
        )

    problem = Problem(manifold, cost)
    optimizer = TrustRegions(verbosity=0, miniter=3)

    rng_local = np.random.default_rng(seed)
    # Start from a deterministic orthonormal basis so repeated runs can be
    # compared without solver-initialization drift.
    q0 = qr_orth(rng_local.standard_normal((n, d1 + d2), dtype=np.float64))
    result = optimizer.run(problem, initial_point=q0)
    Q = np.asarray(result.point, dtype=np.float64)
    return Q, P1, P2



def pg_decompose_python(
    seq: np.ndarray,
    move_mask: np.ndarray,
    d_null: int = 10,
    d_potent: int = 10,
) -> Dict[str, np.ndarray]:
    """Decompose activity into potent and null movement subspaces.

    Parameters
    ----------
    seq:
        Neural activity tensor with shape `(bins, trials, neurons)`.
    move_mask:
        Boolean movement mask with shape `(bins, trials)`.
    d_null, d_potent:
        Target dimensionality of the null and potent projections.

    Returns
    -------
    dict[str, np.ndarray]
        Dictionary containing `potent` and `null` projection tensors.
    """

    seq = np.asarray(seq, dtype=np.float32)
    move_mask = np.asarray(move_mask, dtype=bool)

    if seq.ndim != 3:
        raise ValueError(f"seq must be 3D (bins, trials, neurons), got {seq.shape}")
    if move_mask.ndim != 2:
        raise ValueError(f"move_mask must be 2D (bins, trials), got {move_mask.shape}")

    n_bins = min(seq.shape[0], move_mask.shape[0])
    n_trials = min(seq.shape[1], move_mask.shape[1])
    n_neurons = seq.shape[2]

    if n_bins <= 0 or n_trials <= 0 or n_neurons <= 0:
        return {"potent": np.zeros((0, 0, 0), dtype=np.float32), "null": np.zeros((0, 0, 0), dtype=np.float32)}

    seq = seq[:n_bins, :n_trials, :]
    move_mask = move_mask[:n_bins, :n_trials]

    # Flatten in column-major order so the `(bins, trials, neurons)` tensor is
    # linearized in the same conceptual order as the rest of the pipeline.
    temp = np.reshape(seq, (n_bins * n_trials, n_neurons), order="F").astype(np.float64, copy=False)
    full_cat = zscore_matlab_cols(temp)
    full = np.reshape(full_cat, (n_bins, n_trials, n_neurons), order="F").astype(np.float32, copy=False)

    move_mask_flat = np.reshape(move_mask, (-1,), order="F")
    n_move = int(np.sum(move_mask_flat))
    n_nomove = int(np.sum(~move_mask_flat))
    if n_move < 2 or n_nomove < 2:
        return {"potent": np.zeros((0, 0, 0), dtype=np.float32), "null": np.zeros((0, 0, 0), dtype=np.float32)}

    N_null = full_cat[~move_mask_flat, :]
    N_pot = full_cat[move_mask_flat, :]

    cov_null = np.asarray(np.cov(N_null, rowvar=False), dtype=np.float64)
    cov_pot = np.asarray(np.cov(N_pot, rowvar=False), dtype=np.float64)

    if cov_null.ndim != 2 or cov_pot.ndim != 2:
        return {"potent": np.zeros((0, 0, 0), dtype=np.float32), "null": np.zeros((0, 0, 0), dtype=np.float32)}

    cov_null = np.nan_to_num(cov_null, nan=0.0, posinf=0.0, neginf=0.0)
    cov_pot = np.nan_to_num(cov_pot, nan=0.0, posinf=0.0, neginf=0.0)
    cov_null = 0.5 * (cov_null + cov_null.T)
    cov_pot = 0.5 * (cov_pot + cov_pot.T)

    try:
        Q, P1, P2 = orthogonal_subspaces_pg(cov_pot, d_potent, cov_null, d_null, alpha=0.0, seed=101)
    except Exception:
        return {"potent": np.zeros((0, 0, 0), dtype=np.float32), "null": np.zeros((0, 0, 0), dtype=np.float32)}

    Q_pot = (Q @ P1).astype(np.float32, copy=False)
    Q_null = (Q @ P2).astype(np.float32, copy=False)

    proj_pot = np.tensordot(full, Q_pot, axes=([2], [0])).astype(np.float32, copy=False)
    proj_null = np.tensordot(full, Q_null, axes=([2], [0])).astype(np.float32, copy=False)

    return {"potent": proj_pot, "null": proj_null}



def ensure_trials_by_cells(X: np.ndarray, n_trials_hint: int) -> np.ndarray:
    """Coerce a matrix to `(n_trials, n_cells)` using a trial-count hint.

    Parameters
    ----------
    X:
        Input matrix or vector.
    n_trials_hint:
        Expected number of trials used to choose whether a transpose or reshape
        is needed.

    Returns
    -------
    np.ndarray
        Matrix with trials on axis 0 and cells/features on axis 1.
    """

    X = np.asarray(X, dtype=np.float32)
    if X.ndim == 1:
        return X.reshape(1, -1) if int(n_trials_hint) <= 1 else X.reshape(int(n_trials_hint), -1)
    if X.ndim != 2:
        return X.reshape(int(n_trials_hint), -1)
    if X.shape[0] == int(n_trials_hint):
        return X
    if X.shape[1] == int(n_trials_hint):
        return X.T
    return X.reshape(int(n_trials_hint), -1)