#!/usr/bin/env python3
"""Reference exporter for OMEGA Marked Reference Geometry, Paper VII v1.6.0.

The formulas recorded in ``exact_formulas`` are normative.  JSON numbers and
OBJ vertices are IEEE-754 binary64 evaluations of those formulas.  The sampled
OBJ circle is a visualization polyline; the six marked points are evaluated
directly from the defining equations.
"""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path
from typing import Iterable, Sequence


VERSION = "OMEGA_Paper_VII_v1.6.0"
PHI = (1.0 + math.sqrt(5.0)) / 2.0
EPSILON = math.pi / 648000.0
DELTA = EPSILON / 2.0
SIN_DELTA = math.sin(DELTA)
COS_DELTA = math.cos(DELTA)
X_LEVELS = (SIN_DELTA, 1.0 + SIN_DELTA, PHI * PHI + SIN_DELTA)

Vector = tuple[float, float, float]
Matrix = tuple[Vector, Vector, Vector]


def dot(x: Sequence[float], y: Sequence[float]) -> float:
    return sum(a * b for a, b in zip(x, y))


def cross(x: Sequence[float], y: Sequence[float]) -> Vector:
    return (
        x[1] * y[2] - x[2] * y[1],
        x[2] * y[0] - x[0] * y[2],
        x[0] * y[1] - x[1] * y[0],
    )


def norm(x: Sequence[float]) -> float:
    return math.sqrt(dot(x, x))


def matvec(matrix: Matrix, vector: Sequence[float]) -> Vector:
    return tuple(dot(row, vector) for row in matrix)  # type: ignore[return-value]


def matmul(left: Matrix, right: Matrix) -> Matrix:
    columns = tuple(zip(*right))
    return tuple(
        tuple(dot(row, column) for column in columns) for row in left
    )  # type: ignore[return-value]


def rotation_matrix(rx: float, ry: float, rz: float) -> Matrix:
    """Return Rz(rz) Ry(ry) Rx(rx), with angles in radians."""

    cx, sx = math.cos(rx), math.sin(rx)
    cy, sy = math.cos(ry), math.sin(ry)
    cz, sz = math.cos(rz), math.sin(rz)
    rot_x: Matrix = ((1.0, 0.0, 0.0), (0.0, cx, -sx), (0.0, sx, cx))
    rot_y: Matrix = ((cy, 0.0, sy), (0.0, 1.0, 0.0), (-sy, 0.0, cy))
    rot_z: Matrix = ((cz, -sz, 0.0), (sz, cz, 0.0), (0.0, 0.0, 1.0))
    return matmul(rot_z, matmul(rot_y, rot_x))


def validate_parameters(scale: float, samples: int) -> None:
    if not math.isfinite(scale) or scale <= 0.0:
        raise ValueError("scale must be finite and strictly positive")
    if samples < 64:
        raise ValueError("samples must be at least 64")


def transform(
    point: Sequence[float], scale: float, rotation: Matrix, translation: Sequence[float]
) -> Vector:
    rotated = matvec(rotation, tuple(scale * value for value in point))
    return tuple(rotated[i] + translation[i] for i in range(3))  # type: ignore[return-value]


def canonical_frame() -> dict[str, Vector]:
    # This convention matches the paper's displayed XYZ coordinates.
    return {
        "u": (0.0, -1.0, 0.0),
        "v": (1.0, 0.0, 0.0),
        "n": (0.0, 0.0, 1.0),
    }


def canonical_landmarks() -> dict[str, Vector]:
    x0, x1, x2 = X_LEVELS
    c = COS_DELTA
    return {
        "alpha": (0.0, 0.0, 0.0),
        "d2": (-x2, -c, 0.0),
        "d1": (-x1, -c, 0.0),
        "d": (-x0, -c, 0.0),
        "b": (x0, -c, 0.0),
        "b1": (x1, -c, 0.0),
        "b2": (x2, -c, 0.0),
    }


def positive_bearings() -> tuple[float, float, float]:
    theta0 = DELTA
    theta1 = math.pi / 4.0 + EPSILON / 4.0
    theta2 = math.atan2(X_LEVELS[2], COS_DELTA)
    return theta0, theta1, theta2


def signed_bearings() -> dict[str, float]:
    theta0, theta1, theta2 = positive_bearings()
    values = (-theta2, -theta1, -theta0, theta0, theta1, theta2)
    names = ("d2", "d1", "d", "b", "b1", "b2")
    return dict(zip(names, values))


def angle_matrix() -> list[list[float]]:
    beta = list(signed_bearings().values())
    return [[abs(left - right) for right in beta] for left in beta]


def transformed_frame(rotation: Matrix) -> dict[str, Vector]:
    return {name: matvec(rotation, vector) for name, vector in canonical_frame().items()}


def export_obj(
    path: Path,
    *,
    scale: float = 1.0,
    translation: Vector = (0.0, 0.0, 0.0),
    rotation_deg: Vector = (0.0, 0.0, 0.0),
    samples: int = 4096,
) -> None:
    validate_parameters(scale, samples)
    rotation = rotation_matrix(*(math.radians(value) for value in rotation_deg))
    vertices: list[Vector] = []
    lines: list[list[int]] = []

    # Major circular arc, from b to d, in the orientation of the canonical frame.
    for index in range(samples):
        parameter = DELTA + (2.0 * math.pi - 2.0 * DELTA) * index / (samples - 1)
        point = (math.sin(parameter), -math.cos(parameter), 0.0)
        vertices.append(transform(point, scale, rotation, translation))
    lines.append(list(range(1, samples + 1)))

    x0, x1, x2 = X_LEVELS
    c = COS_DELTA
    for segment in (
        ((-x0, -c, 0.0), (-x1, -c, 0.0), (-x2, -c, 0.0)),
        ((x0, -c, 0.0), (x1, -c, 0.0), (x2, -c, 0.0)),
    ):
        indices: list[int] = []
        for point in segment:
            vertices.append(transform(point, scale, rotation, translation))
            indices.append(len(vertices))
        lines.append(indices)

    with path.open("w", encoding="utf-8", newline="\n") as stream:
        stream.write(f"# {VERSION} binary64 visualization export\n")
        stream.write(f"# epsilon_rad {EPSILON:.17g}\n")
        stream.write(f"# phi {PHI:.17g}\n")
        stream.write(f"# circle_samples {samples}\n")
        stream.write("o OMEGA_local_kernel\n")
        for vertex in vertices:
            stream.write("v %.17g %.17g %.17g\n" % vertex)
        for line in lines:
            stream.write("l " + " ".join(map(str, line)) + "\n")


def export_json(
    path: Path,
    *,
    scale: float = 1.0,
    translation: Vector = (0.0, 0.0, 0.0),
    rotation_deg: Vector = (0.0, 0.0, 0.0),
    samples: int = 4096,
) -> None:
    validate_parameters(scale, samples)
    rotation = rotation_matrix(*(math.radians(value) for value in rotation_deg))
    landmarks = {
        name: transform(point, scale, rotation, translation)
        for name, point in canonical_landmarks().items()
    }
    frame = transformed_frame(rotation)

    data = {
        "schema": "OMEGA_MARKED_REFERENCE_GEOMETRY/1.2",
        "paper_version": VERSION,
        "numeric_model": "IEEE-754 binary64 approximation",
        "normative_layer": "exact_formulas",
        "construction_inputs": {
            "radius_scale": "a>0",
            "seed_angle": "epsilon=pi/648000 rad=1 arcsecond",
            "frame": "ordered orthonormal pair (u,v)",
            "bilateral_steps": ["a", "a*phi"],
            "circular_component": "oriented major arc from b to d",
        },
        "exact_formulas": {
            "phi": "(1+sqrt(5))/2",
            "epsilon": "pi/648000",
            "delta": "epsilon/2",
            "s": "sin(delta)",
            "c": "cos(delta)",
            "x0": "s",
            "x1": "1+s",
            "x2": "phi^2+s",
            "golden_accumulation": "a+a*phi=a*(1+phi)=a*phi^2",
            "point": "alpha+a*(c*u+sigma*x_j*v)",
            "major_arc": "alpha+a*(cos(t)*u+sin(t)*v), delta<=t<=2*pi-delta",
            "left_foot": "d-a*lambda*v, 0<=lambda<=phi^2",
            "right_foot": "b+a*lambda*v, 0<=lambda<=phi^2",
            "gram_factorization": "G=c^2*1*1^T+z*z^T",
            "intrinsic_to_cartesian": "(xi_u,xi_v)->(xi_v,-xi_u)",
        },
        "constants_approx": {
            "epsilon_rad": EPSILON,
            "epsilon_arcsec": 1.0,
            "delta_rad": DELTA,
            "phi": PHI,
            "s": SIN_DELTA,
            "c": COS_DELTA,
            "x_levels": X_LEVELS,
        },
        "transform": {
            "scale": scale,
            "translation": translation,
            "rotation_deg_xyz": rotation_deg,
            "rotation_convention": "active Rz(rz) Ry(ry) Rx(rx)",
            "rotation_matrix": rotation,
        },
        "frame_approx": frame,
        "landmarks_approx": landmarks,
        "signed_bearings_rad_approx": signed_bearings(),
        "center_angle_matrix_rad_approx": angle_matrix(),
        "invariants": {
            "baseline_affine_dimension": 1,
            "center_augmented_affine_dimension": 2,
            "center_relative_matrix_shape": [3, 6],
            "gram_matrix_shape": [6, 6],
            "center_relative_matrix_rank": 2,
            "gram_matrix_rank": 2,
            "ambient_export_dimension": 3,
            "ordered_landmarks": ["d2", "d1", "d", "b", "b1", "b2"],
            "obj_circle_samples": samples,
            "major_arc_traversal": "counterclockwise from b to d in canonical Cartesian coordinates",
        },
    }
    path.write_text(json.dumps(data, indent=2, sort_keys=False) + "\n", encoding="utf-8")


def parse_vector(values: Iterable[float]) -> Vector:
    vector = tuple(values)
    if len(vector) != 3 or not all(math.isfinite(value) for value in vector):
        raise ValueError("expected three finite values")
    return vector  # type: ignore[return-value]


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--scale", type=float, default=1.0)
    parser.add_argument(
        "--translate", nargs=3, type=float, default=(0.0, 0.0, 0.0), metavar=("X", "Y", "Z")
    )
    parser.add_argument(
        "--rotate", nargs=3, type=float, default=(0.0, 0.0, 0.0), metavar=("RX", "RY", "RZ"),
        help="active XYZ angles in degrees, composed as Rz Ry Rx",
    )
    parser.add_argument("--samples", type=int, default=4096)
    parser.add_argument("--obj", default="OMEGA_local_kernel_v1_6_0.obj")
    parser.add_argument("--json", default="OMEGA_local_kernel_v1_6_0.json")
    arguments = parser.parse_args()

    translation = parse_vector(arguments.translate)
    rotation_deg = parse_vector(arguments.rotate)
    export_obj(
        Path(arguments.obj), scale=arguments.scale, translation=translation,
        rotation_deg=rotation_deg, samples=arguments.samples,
    )
    export_json(
        Path(arguments.json), scale=arguments.scale, translation=translation,
        rotation_deg=rotation_deg, samples=arguments.samples,
    )


if __name__ == "__main__":
    main()