"""Co-rotated strain utilities and the tied linear (PCA) strain basis.

3D analog of the 2D result: after analytic co-rotation, per-tet strain
lives in the 6-dim symmetric space; the tied basis compresses it to k
components shared across every tet, material, and scenario.

vec6 uses sqrt(2) off-diagonal scaling so Euclidean norm == Frobenius
norm of the symmetric matrix (keeps PCA variance and constraint energy
consistent).
"""
import numpy as np

_SQRT2 = np.sqrt(2.0)


def sym_to_vec6(S):
    """(N,3,3) symmetric -> (N,6): [xx, yy, zz, √2·xy, √2·yz, √2·xz]."""
    return np.stack([
        S[:, 0, 0], S[:, 1, 1], S[:, 2, 2],
        _SQRT2 * S[:, 0, 1], _SQRT2 * S[:, 1, 2], _SQRT2 * S[:, 0, 2],
    ], axis=1)


def vec6_to_sym(v):
    """(N,6) -> (N,3,3) symmetric."""
    S = np.empty((v.shape[0], 3, 3))
    S[:, 0, 0], S[:, 1, 1], S[:, 2, 2] = v[:, 0], v[:, 1], v[:, 2]
    S[:, 0, 1] = S[:, 1, 0] = v[:, 3] / _SQRT2
    S[:, 1, 2] = S[:, 2, 1] = v[:, 4] / _SQRT2
    S[:, 0, 2] = S[:, 2, 0] = v[:, 5] / _SQRT2
    return S


def corotated_strain(F, R):
    """e = sym(R^T F) - I as (N,6). The skew part of R^T F is extraction
    residual, not strain — discard it, exactly as in 2D."""
    RtF = np.einsum("nji,njk->nik", R, F)
    S = 0.5 * (RtF + RtF.transpose(0, 2, 1)) - np.eye(3)
    return sym_to_vec6(S)


class TiedStrainBasis:
    """One linear basis (6,k) shared by every tet — the tied embedding.

    Uncentered PCA (SVD of the raw sample matrix): rest strain is 0 and
    the constraint energy is a norm around 0, so no mean removal.
    """

    def __init__(self, B):
        self.B = B  # (6, k)

    @classmethod
    def fit(cls, samples, k):
        """samples (N,6) -> basis of top-k right singular vectors."""
        _, s, Vt = np.linalg.svd(samples, full_matrices=False)
        basis = cls(Vt[:k].T.copy())
        basis.singular_values = s
        return basis

    @property
    def k(self):
        return self.B.shape[1]

    def encode(self, e):
        return e @ self.B

    def decode(self, z):
        return z @ self.B.T

    def project(self, e):
        return self.decode(self.encode(e))

    def coverage(self, samples):
        """Fraction of strain energy reconstructed: 1 - ||e - BBᵀe||²/||e||²."""
        num = ((samples - self.project(samples)) ** 2).sum()
        den = (samples ** 2).sum()
        return 1.0 - num / max(den, 1e-30)