#!/usr/bin/env python3
"""Free-fall camera with parallel (= geodesic Fermi-Walker) transport.

Ingoing Cartesian Kerr-Schild, signature -+++, G=c=M=1. Requires numpy/scipy.
The renderer's Schwarzschild backend fixes M=1 as well.
"""
import argparse
import csv
import math
from pathlib import Path
import sys

import numpy as np
from scipy.integrate import solve_ivp

ETA = np.diag([-1., 1., 1., 1.])
HEADER = ['t', 'tau', 'x', 'y', 'z'] + [f'e{a}{c}' for a in range(4) for c in 'txyz']


def metric_connection(x):
    """Analytic g and Christoffels; dg[k,mu,nu] = partial_k g_mu_nu."""
    r = np.linalg.norm(x)
    if not np.isfinite(r) or r <= 0:
        raise ValueError('metric undefined at r=0 or nonfinite position')
    n = x / r
    ell = np.r_[1., n]
    f = 2 / r
    g = ETA + f * np.outer(ell, ell)
    raised = ETA @ ell
    inverse = ETA - f * np.outer(raised, raised)
    dg = np.zeros((4, 4, 4))
    for k in range(3):
        dl = np.r_[0., (np.eye(3)[k] - n[k] * n) / r]
        dg[k+1] = f * (np.outer(dl, ell) + np.outer(ell, dl)
                         - n[k] / r * np.outer(ell, ell))
    connection = .5 * np.einsum('ml,alb->mab', inverse,
                                dg + dg.transpose(2, 1, 0) - dg.transpose(1, 0, 2))
    return g, connection


def initial_state(position, velocity, ra=None, dec=None, roll=0., tetrad=None, t0=0.):
    g, _ = metric_connection(position)
    u = np.r_[1., velocity]
    q = u @ g @ u
    if not np.isfinite(q) or q >= 0:
        raise ValueError('coordinate velocity must be future timelike: g(1,v;1,v) < 0')
    u /= math.sqrt(-q)
    if tetrad is not None:
        e = np.asarray(tetrad, dtype=float)
        if e.shape != (4, 4) or not np.all(np.isfinite(e)):
            raise ValueError('initial tetrad must contain four rows of four finite components')
        if not np.allclose(e[0], u, rtol=1e-9, atol=1e-9):
            raise ValueError('initial tetrad e0 must agree with the specified coordinate velocity')
        if np.max(np.abs(e @ g @ e.T - ETA)) > 1e-8:
            raise ValueError('initial tetrad must be Lorentz orthonormal')
        if np.linalg.det(e) <= 0:
            raise ValueError('initial tetrad must have forward cross up = right orientation')
    else:
        if ra is None:
            direction = -np.asarray(position) / np.linalg.norm(position)
            ra = math.degrees(math.atan2(direction[1], direction[0])) % 360
            dec = math.degrees(math.asin(direction[2]))
        a, d = np.deg2rad([ra, dec])
        e = np.array([u, [0, math.cos(d)*math.cos(a), math.cos(d)*math.sin(a), math.sin(d)],
                      [0, -math.sin(d)*math.cos(a), -math.sin(d)*math.sin(a), math.cos(d)],
                      [0, math.sin(a), -math.cos(a), 0]])
        for i in range(1, 4):
            for _ in range(2):
                for j in range(i):
                    e[i] -= (-1 if j == 0 else 1) * (e[i] @ g @ e[j]) * e[j]
            e[i] /= math.sqrt(e[i] @ g @ e[i])
        angle = math.radians(roll)
        up, right = e[2].copy(), e[3].copy()
        e[2] = math.cos(angle)*up + math.sin(angle)*right
        e[3] = -math.sin(angle)*up + math.cos(angle)*right
    return np.r_[t0, position, e.ravel()]


def rhs(tau, state):
    _, connection = metric_connection(state[1:4])
    e = state[4:].reshape(4, 4)
    # e0 is u: transporting all four legs also integrates the timelike geodesic.
    de = -np.einsum('mab,a,ib->im', connection, e[0], e)
    return np.r_[e[0], de.ravel()]


def integrate(state, duration, fps, stop_radius=1e-3, rtol=1e-10, atol=1e-12):
    if np.linalg.norm(state[1:4]) <= stop_radius:
        raise ValueError('initial radius must exceed stop-radius')
    if duration == 0:
        return np.array([0.]), state[None, :], None

    def stop(tau, y):
        return np.linalg.norm(y[1:4]) - stop_radius
    stop.terminal = True
    stop.direction = -1
    solution = solve_ivp(rhs, (0., duration), state, method='DOP853',
                         rtol=rtol, atol=atol, events=stop, dense_output=True)
    if not solution.success:
        raise ValueError(f'integration failed at tau={solution.t[-1]:.17g}: {solution.message}')
    end = solution.t[-1]
    # Include tau=0 and every full 1/fps interval; never add an off-cadence endpoint.
    count = int(math.floor(np.nextafter(end * fps, np.inf))) + 1
    tau = np.arange(count, dtype=float) / fps
    tau = tau[tau <= end]
    states = solution.sol(tau).T
    if not np.all(np.isfinite(states)) or np.any(np.diff(states[:, 0]) <= 0):
        raise ValueError('nonfinite trajectory or coordinate time lost monotonicity')
    return tau, states, end if solution.t_events[0].size else None


def main():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    p.add_argument('--output', type=Path, required=True, help='21-column movie CSV')
    p.add_argument('--position', nargs=3, type=float, required=True, metavar=('X', 'Y', 'Z'), help='initial KS Cartesian position in M')
    p.add_argument('--velocity', nargs=3, type=float, default=[0., 0., 0.], help='coordinate dx/dt, dy/dt, dz/dt (not local 3-speed)')
    p.add_argument('--look-ra-deg', type=float, help='initial forward RA; default points toward origin')
    p.add_argument('--look-dec-deg', type=float, help='initial forward Dec; specify together with RA')
    p.add_argument('--roll-deg', type=float, default=0., help='initial roll, same sign as renderer')
    p.add_argument('--tetrad', type=float, nargs=16, help='explicit row-major e0,e1,e2,e3 in t,x,y,z; replaces look/roll')
    p.add_argument('--t0', type=float, default=0., help='initial KS coordinate time')
    p.add_argument('--fps', type=float, default=30., help='samples per unit proper time M')
    p.add_argument('--duration', type=float, required=True, help='requested elapsed proper time in M')
    p.add_argument('--stop-radius', type=float, default=1e-3, help='numerical singularity guard in M, strictly between 0 and 2')
    p.add_argument('--rtol', type=float, default=1e-10, help='DOP853 relative error tolerance')
    p.add_argument('--atol', type=float, default=1e-12, help='DOP853 absolute error tolerance')
    args = p.parse_args()
    try:
        numbers = [*args.position, *args.velocity, args.roll_deg, args.t0, args.fps,
                   args.duration, args.stop_radius, args.rtol, args.atol]
        numbers += [v for v in (args.look_ra_deg, args.look_dec_deg) if v is not None]
        if not all(math.isfinite(v) for v in numbers):
            raise ValueError('all numeric arguments must be finite')
        if args.fps <= 0 or args.duration < 0 or not 0 < args.stop_radius < 2 or min(args.rtol, args.atol) <= 0:
            raise ValueError('require fps,rtol,atol > 0, duration >= 0 and 0 < stop-radius < 2')
        if (args.look_ra_deg is None) != (args.look_dec_deg is None):
            raise ValueError('specify both look-ra-deg and look-dec-deg')
        if args.look_ra_deg is not None and not (0 <= args.look_ra_deg < 360 and abs(args.look_dec_deg) <= 90):
            raise ValueError('require 0 <= RA < 360 and -90 <= Dec <= 90')
        if args.tetrad is not None and (args.look_ra_deg is not None or args.roll_deg != 0):
            raise ValueError('tetrad conflicts with look/roll')
        e = None if args.tetrad is None else np.array(args.tetrad).reshape(4, 4)
        state = initial_state(args.position, args.velocity, args.look_ra_deg,
                              args.look_dec_deg, args.roll_deg, e, args.t0)
        tau, states, stopped = integrate(state, args.duration, args.fps,
                                         args.stop_radius, args.rtol, args.atol)
        error = max(np.max(np.abs(s[4:].reshape(4, 4) @ metric_connection(s[1:4])[0]
                                  @ s[4:].reshape(4, 4).T - ETA)) for s in states)
        if error > 1e-6:
            raise ValueError(f'tetrad norm drift {error:.3g} exceeds 1e-6; tighten tolerances')
        with args.output.open('w', newline='') as stream:
            writer = csv.writer(stream)
            writer.writerow(HEADER)
            for t, s in zip(tau, states):
                writer.writerow(format(v, '.17g') for v in np.r_[s[0], t, s[1:]])
        print(f'Wrote {len(tau)} frames; tau=0..{tau[-1]:.17g}, t={states[0,0]:.17g}..{states[-1,0]:.17g}; max tetrad error={error:.3g}', file=sys.stderr)
        if stopped is not None:
            print(f'Stopped early at r={args.stop_radius:g} M, tau={stopped:.17g} (numerical guard before r=0).', file=sys.stderr)
        print('Render with --observer-track <CSV> --movie-track-samples --frames-dir <DIR>; encode at the chosen fps.', file=sys.stderr)
    except (ValueError, OSError, OverflowError) as exc:
        p.error(str(exc))


if __name__ == '__main__':
    main()