"""Câmera 3D das cenas do jornal (F-jornal1, J-jornal2) a partir da homografia medida no original.

Achado (17/10): com a página "plana" bem definida, o original bate com a câmera padrão do AE para comp 1080
(zoom 1500 px, centro óptico no centro) e o jornal girando em 3D: erro médio < 1 px. É o rig do autor: camadas
3D, sem camada de câmera.

Passos:
 1) homografia-<X>.json leva a tela do quadro-base para a tela de cada quadro (homografia_jornal.py);
 2) o plano verdadeiro = Hr · (tela do quadro-base - 540), com Hr = retificação de 4 parâmetros (a página pode
    não estar 100% de frente no quadro-base); Hr é escolhida para minimizar o erro de todos os quadros;
 3) por quadro: pose (R, t) do plano para a câmera fixa f = 1500 (solvePnP + mínimos quadrados);
 4) no começo da cena a pose é misturada com a pose "de frente" (para o encaixe encher a tela sem quina);
 5) saída para o AE: null 3D "PLANO" (ponto de ancoragem 0) com Posição = t + [540, 540, -1500] e
    Rotação X/Y/Z tais que Rx·Ry·Rz = R (ordem do AE medida no próprio AE).
Saída: camera/camera3d-<X>.json  {f, quadros: [{f, t, pos, rot:[x,y,z], erro_px}], Hr, encaixe}
Uso: python -I camera3d_jornal.py F|J
"""
import json, sys, numpy as np, cv2
from scipy.optimize import least_squares, minimize
sys.stdout.reconfigure(encoding="utf-8")
CENA = sys.argv[1]
FOC, C = 1500.0, 540.0
K = np.array([[FOC, 0, C], [0, FOC, C], [0, 0, 1.0]])
d = json.load(open(f"homografia-{CENA}.json"))
H = {int(k): np.array(v) for k, v in d["H"].items()}
fs = sorted(H)
# troca no 1º quadro inteiro da cena (10,80 s e 24,62 s caem entre quadros: 1º quadro visível = 260 e 591)
TROCA = {"F": 260 / 24, "J": 591 / 24}[CENA]
MISTURA = {"F": (260 / 24, 11.00), "J": (591 / 24, 24.80)}[CENA]
gx = np.linspace(60, 1020, 7)
GRADE = np.array([[x, y] for y in gx for x in gx], float)


def Hr_de(g):
    a, b, p, q = g
    return np.array([[1 + a, b, 0], [0, 1, 0], [p / 1000, q / 1000, 1.0]])


def plano_de_tela(f, g, uv):
    Hi = np.linalg.inv(H[f]); P = (Hi @ np.c_[uv, np.ones(len(uv))].T).T; x = P[:, :2] / P[:, 2:3] - C
    Q = (Hr_de(g) @ np.c_[x, np.ones(len(x))].T).T
    return Q[:, :2] / Q[:, 2:3]


def proj(p, X):
    R, _ = cv2.Rodrigues(p[:3]); Xc = (R @ np.c_[X, np.zeros(len(X))].T).T + p[3:6]
    return np.c_[FOC * Xc[:, 0] / Xc[:, 2] + C, FOC * Xc[:, 1] / Xc[:, 2] + C]


def pose(f, g, p0=None):
    X = plano_de_tela(f, g, GRADE); u = GRADE
    ok, rv, tv = cv2.solvePnP(np.c_[X, np.zeros(len(X))], u, K, None, flags=cv2.SOLVEPNP_IPPE)
    inits = [np.r_[rv.ravel(), tv.ravel()]]
    if p0 is not None: inits.append(p0)
    best = None
    for q0 in inits:
        s = least_squares(lambda p: (proj(p, X) - u).ravel(), q0, method="lm", max_nfev=200)
        e = float(np.sqrt(((proj(s.x, X) - u) ** 2).sum(1)).mean())
        if best is None or e < best[1]: best = (s.x, e)
    return best


cache = {}
def custo(g):
    tot = []
    for f in fs[::2]:
        p, e = pose(f, g, cache.get(f)); cache[f] = p; tot.append(e)
    return float(np.mean(np.array(tot) ** 2))


print("sem retificação: rms", custo(np.zeros(4)) ** .5)
r = minimize(custo, np.zeros(4), method="Nelder-Mead",
             options={"xatol": 1e-6, "fatol": 1e-4, "maxiter": 300,
                      "initial_simplex": np.array([[0, 0, 0, 0], [.02, 0, 0, 0], [0, .02, 0, 0], [0, 0, .05, 0], [0, 0, 0, .05]])})
g = r.x
print("retificação (Nelder-Mead)", np.round(g, 6), "rms", r.fun ** .5)
# ajuste conjunto (retificação + todas as poses) por mínimos quadrados esparsos
from scipy.sparse import lil_matrix
N = len(fs); NP = 4 + 6 * N
x0 = np.r_[g, np.concatenate([pose(f, g, cache.get(f))[0] for f in fs])]
def res_all(x):
    gg = x[:4]; out = []
    for i, f in enumerate(fs):
        X = plano_de_tela(f, gg, GRADE); out.append((proj(x[4 + 6 * i: 10 + 6 * i], X) - GRADE).ravel())
    return np.concatenate(out)
M = len(GRADE) * 2
J = lil_matrix((M * N, NP), dtype=int)
for i in range(N):
    J[i * M:(i + 1) * M, :4] = 1; J[i * M:(i + 1) * M, 4 + 6 * i: 10 + 6 * i] = 1
sol = least_squares(res_all, x0, jac_sparsity=J, x_scale="jac", max_nfev=200, verbose=0)
g = sol.x[:4]
for i, f in enumerate(fs): cache[f] = sol.x[4 + 6 * i: 10 + 6 * i]
rr = sol.fun.reshape(N, -1, 2); print("retificação (conjunto)", np.round(g, 6), "rms", float(np.sqrt((rr ** 2).sum(2)).mean()))
poses = {}
for f in fs:
    poses[f] = pose(f, g, cache.get(f))


def frontal(p):
    """Mesma pose, sem inclinação: a normal do plano alinhada ao eixo da câmera (mantém o giro no plano e t)."""
    R, _ = cv2.Rodrigues(p[:3])
    n = R[:, 2]  # normal do plano em coordenadas da câmera
    # rotação mínima que leva n para (0,0,1)
    z = np.array([0, 0, 1.0]); v = np.cross(n, z); s = np.linalg.norm(v); c = float(n @ z)
    if s < 1e-9: return p.copy()
    vx = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
    Rm = np.eye(3) + vx + vx @ vx * ((1 - c) / s ** 2)
    R2 = Rm @ R
    # mantém o ponto do plano que está no centro da tela no mesmo lugar e à mesma distância
    t = p[3:6]
    Xc0 = np.array([0, 0, 1.0]) * 1.0
    # ponto do plano no eixo óptico: R X + t = lambda * z
    A = np.c_[R[:, 0], R[:, 1], -z]; sol = np.linalg.solve(A, -t); X0 = np.array([sol[0], sol[1], 0]); lam = sol[2]
    t2 = lam * z - R2 @ X0
    return np.r_[cv2.Rodrigues(R2)[0].ravel(), t2]


def slerp_pose(pa, pb, w):
    Ra, _ = cv2.Rodrigues(pa[:3]); Rb, _ = cv2.Rodrigues(pb[:3])
    dr, _ = cv2.Rodrigues(Rb @ Ra.T); Rw = cv2.Rodrigues(dr * w)[0] @ Ra
    return np.r_[cv2.Rodrigues(Rw)[0].ravel(), (1 - w) * pa[3:6] + w * pb[3:6]]


a0, a1 = MISTURA
final = {}
for f in fs:
    p, e = poses[f]
    t = f / 24
    if t < a1:
        w = 0.0 if t <= a0 else (t - a0) / (a1 - a0)
        w = w * w * (3 - 2 * w)
        p = slerp_pose(frontal(p), p, w)
        X = plano_de_tela(f, g, GRADE); e = float(np.sqrt(((proj(p, X) - GRADE) ** 2).sum(1)).mean())
    final[f] = (p, e)


def euler_ae(R):
    """R = Rx(a)·Ry(b)·Rz(c) (ordem do AE) -> [a, b, c] em graus."""
    b = np.arcsin(np.clip(R[0, 2], -1, 1)); c = np.arctan2(-R[0, 1], R[0, 0]); a = np.arctan2(-R[1, 2], R[2, 2])
    return np.degrees([a, b, c])


out = []
prev = None
for f in fs:
    p, e = final[f]
    R, _ = cv2.Rodrigues(p[:3]); t = p[3:6]
    rot = euler_ae(R)
    if prev is not None:  # sem saltos de 360°
        rot = prev + ((rot - prev + 180) % 360 - 180)
    prev = rot
    out.append({"f": f, "t": round(f / 24, 5), "pos": (t + np.array([C, C, -FOC])).tolist(), "rot": rot.tolist(),
                "R": R.tolist(), "tc": t.tolist(), "erro_px": round(e, 2)})
    print(f, f"{f/24:7.3f}", "pos", np.round(t + [C, C, -FOC], 1), "rot", np.round(rot, 2), f"erro {e:.2f}")

# encaixe: quadrado na tela no instante da troca -> quadrilátero no plano (com a pose interpolada na troca)
f0 = max(f for f in fs if f / 24 <= TROCA); f1 = f0 + 1
w = (TROCA - f0 / 24) * 24
pt = slerp_pose(final[f0][0], final[f1][0], w)
R, _ = cv2.Rodrigues(pt[:3]); tt = pt[3:6]
cantos = []
for u, v in ((0, 0), (1080, 0), (1080, 1080), (0, 1080)):
    ray = np.array([(u - C) / FOC, (v - C) / FOC, 1.0])
    # R X + t = s ray, X=(x,y,0)
    A = np.c_[R[:, 0], R[:, 1], -ray]; sol = np.linalg.solve(A, -tt); cantos.append([sol[0], sol[1]])
cantos = np.array(cantos)
cen = cantos.mean(0); lado = float(np.linalg.norm(cantos[1] - cantos[0]))
giro = float(np.degrees(np.arctan2(cantos[1][1] - cantos[0][1], cantos[1][0] - cantos[0][0])))
print("encaixe no plano: cantos", np.round(cantos, 2).tolist(), "centro", np.round(cen, 2), "lado", round(lado, 3), "giro", round(giro, 3))
json.dump({"foco": FOC, "Hr": Hr_de(g).tolist(), "g": g.tolist(), "troca": TROCA, "mistura": MISTURA,
           "encaixe": {"cantos": cantos.tolist(), "centro": cen.tolist(), "lado": lado, "giro": giro},
           "quadros": out}, open(f"camera3d-{CENA}.json", "w"), indent=0)