"""linalg_lab.py - Linear algebra toolkit from scratch (pure Python, no numpy).

Contains: Matrix class, LU (partial pivoting), solve/det/inverse,
Householder QR, least squares, Jacobi eigen-solver (symmetric),
SVD, and PCA / polynomial-fit demos. Run the file to execute self-tests.
"""
import math
import random


class Matrix:
    def __init__(self, rows):
        self.a = [list(map(float, r)) for r in rows]
        self.m, self.n = len(self.a), len(self.a[0])

    @staticmethod
    def eye(n):
        return Matrix([[1.0 if i == j else 0.0 for j in range(n)] for i in range(n)])

    @staticmethod
    def diag(vals):
        n = len(vals)
        return Matrix([[vals[i] if i == j else 0.0 for j in range(n)] for i in range(n)])

    @property
    def T(self):
        return Matrix([[self.a[i][j] for i in range(self.m)] for j in range(self.n)])

    def __matmul__(self, o):
        assert self.n == o.m, "shape mismatch"
        cols = o.T.a
        return Matrix([[sum(x * y for x, y in zip(r, c)) for c in cols] for r in self.a])

    def __sub__(self, o):
        return Matrix([[x - y for x, y in zip(r, s)] for r, s in zip(self.a, o.a)])

    def norm(self):  # Frobenius norm
        return math.sqrt(sum(x * x for r in self.a for x in r))

    def __repr__(self):
        return "\n".join("[" + " ".join(f"{x:9.4f}" for x in r) + "]" for r in self.a)


def col(v):
    return Matrix([[x] for x in v])


# ---------- LU decomposition: PA = LU ----------
def lu(A):
    n = A.n
    U = [r[:] for r in A.a]
    L = [[0.0] * n for _ in range(n)]
    perm, sign = list(range(n)), 1
    for k in range(n):
        p = max(range(k, n), key=lambda i: abs(U[i][k]))
        if abs(U[p][k]) < 1e-12:
            raise ValueError("matrix is singular")
        if p != k:
            U[k], U[p] = U[p], U[k]
            L[k], L[p] = L[p], L[k]
            perm[k], perm[p] = perm[p], perm[k]
            sign = -sign
        for i in range(k + 1, n):
            L[i][k] = U[i][k] / U[k][k]
            for j in range(k, n):
                U[i][j] -= L[i][k] * U[k][j]
    for i in range(n):
        L[i][i] = 1.0
    return Matrix(L), Matrix(U), perm, sign


def solve(A, b):
    L, U, perm, _ = lu(A)
    n = A.n
    y = [0.0] * n
    for i in range(n):
        y[i] = b[perm[i]] - sum(L.a[i][j] * y[j] for j in range(i))
    x = [0.0] * n
    for i in reversed(range(n)):
        x[i] = (y[i] - sum(U.a[i][j] * x[j] for j in range(i + 1, n))) / U.a[i][i]
    return x


def det(A):
    _, U, _, s = lu(A)
    return s * math.prod(U.a[i][i] for i in range(A.n))


def inverse(A):
    n = A.n
    cols = [solve(A, [1.0 if i == j else 0.0 for i in range(n)]) for j in range(n)]
    return Matrix(cols).T


# ---------- Householder QR: A = QR ----------
def qr(A):
    m, n = A.m, A.n
    R = [r[:] for r in A.a]
    Q = Matrix.eye(m).a
    for k in range(min(m - 1, n)):
        x = [R[i][k] for i in range(k, m)]
        nx = math.sqrt(sum(t * t for t in x))
        if nx < 1e-14:
            continue
        v = x[:]
        v[0] += math.copysign(nx, x[0])
        nv = math.sqrt(sum(t * t for t in v))
        v = [t / nv for t in v]
        for j in range(n):  # R <- H R
            d = sum(v[i] * R[k + i][j] for i in range(m - k))
            for i in range(m - k):
                R[k + i][j] -= 2 * v[i] * d
        for r in range(m):  # Q <- Q H
            d = sum(Q[r][k + i] * v[i] for i in range(m - k))
            for i in range(m - k):
                Q[r][k + i] -= 2 * d * v[i]
    return Matrix(Q), Matrix(R)


def lstsq(A, b):
    """Least squares min ||Ax - b|| via QR (more stable than normal equations)."""
    Q, R = qr(A)
    n = A.n
    qtb = [sum(Q.a[i][j] * b[i] for i in range(A.m)) for j in range(n)]
    x = [0.0] * n
    for i in reversed(range(n)):
        x[i] = (qtb[i] - sum(R.a[i][j] * x[j] for j in range(i + 1, n))) / R.a[i][i]
    return x


# ---------- Jacobi eigenvalue algorithm (symmetric) ----------
def eigh(A, tol=1e-18, max_sweeps=100):
    n = A.n
    S = [r[:] for r in A.a]
    V = Matrix.eye(n).a
    for _ in range(max_sweeps):
        off = sum(S[i][j] ** 2 for i in range(n) for j in range(n) if i != j)
        if off < tol:
            break
        for p in range(n - 1):
            for q in range(p + 1, n):
                if abs(S[p][q]) < 1e-300:
                    continue
                theta = (S[q][q] - S[p][p]) / (2 * S[p][q])
                t = math.copysign(1, theta) / (abs(theta) + math.sqrt(theta * theta + 1))
                c = 1 / math.sqrt(t * t + 1)
                s = t * c
                for k in range(n):
                    a, b = S[k][p], S[k][q]
                    S[k][p], S[k][q] = c * a - s * b, s * a + c * b
                for k in range(n):
                    a, b = S[p][k], S[q][k]
                    S[p][k], S[q][k] = c * a - s * b, s * a + c * b
                for k in range(n):
                    a, b = V[k][p], V[k][q]
                    V[k][p], V[k][q] = c * a - s * b, s * a + c * b
    vals = [S[i][i] for i in range(n)]
    order = sorted(range(n), key=lambda i: -vals[i])
    return [vals[i] for i in order], Matrix([[V[k][i] for i in order] for k in range(n)])


# ---------- SVD via eigen-decomposition of A^T A ----------
def svd(A):
    """Thin SVD: A = U diag(s) V^T. (Small singular values lose precision.)"""
    vals, V = eigh(A.T @ A)
    s = [math.sqrt(max(v, 0.0)) for v in vals]
    AV = A @ V
    U = Matrix([[AV.a[i][j] / s[j] if s[j] > 1e-12 else 0.0 for j in range(A.n)]
                for i in range(A.m)])
    return U, s, V


# ---------- Demos ----------
def pca_demo(n_points=300):
    random.seed(1)
    d = [0.8, 0.5, 0.33]
    X = []
    for _ in range(n_points):
        z = random.gauss(0, 3)
        X.append([z * d[i] + random.gauss(0, 0.5) for i in range(3)])
    mean = [sum(r[j] for r in X) / n_points for j in range(3)]
    Xc = Matrix([[r[j] - mean[j] for j in range(3)] for r in X])
    cov = Matrix([[x / (n_points - 1) for x in r] for r in (Xc.T @ Xc).a])
    vals, V = eigh(cov)
    total = sum(vals)
    print("PCA explained variance:", [f"{v / total:.1%}" for v in vals])


def polyfit_demo():
    random.seed(2)
    xs = [i / 5 for i in range(-10, 11)]
    ys = [2 - 3 * x + 0.5 * x * x + random.gauss(0, 0.1) for x in xs]
    A = Matrix([[x ** j for j in range(3)] for x in xs])
    print("Fit coeffs (true 2, -3, 0.5):", [round(c, 3) for c in lstsq(A, ys)])


def _check(name, err, tol=1e-8):
    print(f"{'PASS' if err < tol else 'FAIL'}  {name:<26} err={err:.2e}")


if __name__ == "__main__":
    random.seed(7)
    n = 5
    rnd = lambda r, c: Matrix([[random.uniform(-1, 1) for _ in range(c)] for _ in range(r)])
    A, B = rnd(n, n), rnd(n, n)
    b = [random.uniform(-1, 1) for _ in range(n)]

    _check("LU solve", (A @ col(solve(A, b)) - col(b)).norm())
    _check("inverse", (A @ inverse(A) - Matrix.eye(n)).norm())
    Q, R = qr(A)
    _check("QR reconstruct", (Q @ R - A).norm())
    _check("Q orthogonal", (Q.T @ Q - Matrix.eye(n)).norm())

    S = B.T @ B
    vals, V = eigh(S)
    _check("eigh  V^T S V = D", (V.T @ S @ V - Matrix.diag(vals)).norm())

    M = rnd(6, 4)
    U, s, Vm = svd(M)
    _check("SVD reconstruct", (U @ Matrix.diag(s) @ Vm.T - M).norm(), 1e-6)
    print(f"det(A) = {det(A):.6f}\n")

    polyfit_demo()
    pca_demo()