"""Earthquake shake table, in Python: a uniform n-storey shear building (equal masses m, equal storey stiffness k, fixed base, free top) under
harmonic ground acceleration.  Modes against the exact formula, Newmark's average acceleration method, the stability limits of the
central difference (omega h <= 2) and linear acceleration (omega h <= 2 sqrt 3) schemes, and the energy balance.  Standard library only."""
from math import sin, cos, pi, sqrt

def build(n, T1, m=2.0e5, zeta=0.02):
    k = m * (2 * pi / T1 / (2 * sin(pi / (4 * n + 2)))) ** 2
    K = [[0.0] * n for _ in range(n)]
    for i in range(n):
        K[i][i] = k + (k if i + 1 < n else 0)
        if i + 1 < n: K[i][i + 1] = K[i + 1][i] = -k
    return k, m, K

def jacobi(A):
    n = len(A); A = [r[:] for r in A]
    for _ in range(60):
        off = sum(A[i][j] ** 2 for i in range(n) for j in range(i + 1, n))
        if off < 1e-24: break
        for p in range(n):
            for q in range(p + 1, n):
                if abs(A[p][q]) < 1e-300: continue
                th = (A[q][q] - A[p][p]) / (2 * A[p][q]); t = (1 if th >= 0 else -1) / (abs(th) + sqrt(th * th + 1)); c = 1 / sqrt(t * t + 1); s = t * c
                for r in range(n): A[r][p], A[r][q] = c * A[r][p] - s * A[r][q], s * A[r][p] + c * A[r][q]
                for r in range(n): A[p][r], A[q][r] = c * A[p][r] - s * A[q][r], s * A[p][r] + c * A[q][r]
    return sorted(A[i][i] for i in range(n))

def solve(A, b):
    n = len(b); M = [A[i][:] + [b[i]] for i in range(n)]
    for i in range(n):
        piv = max(range(i, n), key=lambda r: abs(M[r][i])); M[i], M[piv] = M[piv], M[i]
        for r in range(i + 1, n):
            f = M[r][i] / M[i][i]
            for c in range(i, n + 1): M[r][c] -= f * M[i][c]
    x = [0.0] * n
    for i in range(n - 1, -1, -1): x[i] = (M[i][n] - sum(M[i][j] * x[j] for j in range(i + 1, n))) / M[i][i]
    return x

def newmark(n, T1, f, amp, zeta, h, T, beta=0.25, gamma=0.5):
    """returns peak roof displacement and the energy balance error; ground acceleration amp*sin(2 pi f t)"""
    k, m, K = build(n, T1); w2 = jacobi([[K[i][j] / m for j in range(n)] for i in range(n)])
    wa, wb = sqrt(w2[0]), sqrt(w2[min(n - 1, 2)] if n > 1 else w2[0] * 2.25)
    a0, a1 = 2 * zeta * wa * wb / (wa + wb), 2 * zeta / (wa + wb)
    C = [[a0 * m * (i == j) + a1 * K[i][j] for j in range(n)] for i in range(n)]
    u = [0.0] * n; v = [0.0] * n; ag0 = amp * sin(0.0); a = [-ag0] * n
    mv = lambda A, x: [sum(A[i][j] * x[j] for j in range(n)) for i in range(n)]
    Keff = [[m * (i == j) + gamma * h * C[i][j] + beta * h * h * K[i][j] for j in range(n)] for i in range(n)]
    peak = 0.0; EI = ED = EW = 0.0; Ip = Dp = Wp = 0.0
    for s in range(1, int(T / h) + 1):
        t = s * h; ag = amp * sin(2 * pi * f * t)
        up = [u[i] + h * v[i] + h * h * (0.5 - beta) * a[i] for i in range(n)]; vp = [v[i] + h * (1 - gamma) * a[i] for i in range(n)]
        rhs = [-m * ag - c - kk for c, kk in zip(mv(C, vp), mv(K, up))]
        an = solve(Keff, [rhs[i] - 0 for i in range(n)])            # (M + gamma h C + beta h^2 K) a_new = -M ag - C vp - K up
        u = [up[i] + beta * h * h * an[i] for i in range(n)]; v = [vp[i] + gamma * h * an[i] for i in range(n)]; a = an
        peak = max(peak, abs(u[-1]))
        Ik = -m * sum(v) * ag; Dk = sum(x * y for x, y in zip(v, mv(C, v))); Wk = sum(x * y for x, y in zip(v, mv(K, u)))
        EI += 0.5 * h * (Ip + Ik); ED += 0.5 * h * (Dp + Dk); EW += 0.5 * h * (Wp + Wk); Ip, Dp, Wp = Ik, Dk, Wk
    EK = 0.5 * m * sum(x * x for x in v)
    return peak, abs(EK + ED + EW - EI) / max(EI, 1e-30)

def central_difference_stable(n, T1, f, amp, zeta, h, T):
    """explicit central difference; returns True if the roof displacement stays finite and small"""
    k, m, K = build(n, T1); w2 = jacobi([[K[i][j] / m for j in range(n)] for i in range(n)])
    wa, wb = sqrt(w2[0]), sqrt(w2[min(n - 1, 2)])
    a0, a1 = 2 * zeta * wa * wb / (wa + wb), 2 * zeta / (wa + wb)
    C = [[a0 * m * (i == j) + a1 * K[i][j] for j in range(n)] for i in range(n)]
    Lh = [[m / h ** 2 * (i == j) + C[i][j] / (2 * h) for j in range(n)] for i in range(n)]
    u = [0.0] * n; up = [0.5 * h * h * 0.0] * n
    mv = lambda A, x: [sum(A[i][j] * x[j] for j in range(n)) for i in range(n)]
    for s in range(int(T / h)):
        ag = amp * sin(2 * pi * f * s * h); Ku = mv(K, u); Cu = mv(C, up)
        rhs = [-m * ag - Ku[i] + 2 * m * u[i] / h ** 2 - m * up[i] / h ** 2 + Cu[i] / (2 * h) for i in range(n)]
        un = solve(Lh, rhs); up, u = u, un
        if abs(u[-1]) > 1e3 or u[-1] != u[-1]: return False
    return True

if __name__ == "__main__":
    n, T1 = 5, 0.5; k, m, K = build(n, T1)
    w = [sqrt(x) for x in jacobi([[K[i][j] / m for j in range(n)] for i in range(n)])]
    exact = [2 * sqrt(k / m) * sin((2 * j - 1) * pi / (2 * (2 * n + 1))) for j in range(1, n + 1)]
    print("periods (s):", [round(2 * pi / x, 4) for x in w]); print("max relative error against the exact formula:", max(abs(a - b) / b for a, b in zip(w, exact)))
    ref = newmark(n, T1, 1 / T1, 0.5, 0.02, 0.0025, 8)[0]
    for h in (0.02, 0.01, 0.005):
        p, e = newmark(n, T1, 1 / T1, 0.5, 0.02, h, 8); print(f"h = {h}: peak roof {p:.5f} m (error vs h = 0.0025: {abs(p - ref):.2e}), energy error {e:.1e}")
    wmax = w[-1]; print("omega_max =", round(wmax, 1), "-> central difference needs h <", round(2 / wmax, 4))
    for h in (0.019, 0.022, 0.026, 0.03): print(f"central difference, h = {h}:", "stable" if central_difference_stable(n, T1, 1 / T1, 0.5, 0.02, h, 4) else "blows up")