"""Apply normalization transforms and convert to polar coordinates — plan §7.3–§7.4.

Exact mode follows the reference implementation used to write GTC files (gtc2vcf ``raw_x_y2norm_x_y`` /
``norm_x_y2ilmn_theta_r``, MIT): cos θ and sin θ are computed in double and rounded to float once per bin, the
affine transform runs in float32, negative results are clamped to 0 **before** scaling, and R = |x| + |y|.
BeadArrayFiles (BSD-2) documents the same transform; it differs only in using double precision internally.
"""

from __future__ import annotations

import math
from collections.abc import Sequence

import numpy as np
import numpy.typing as npt

from simurg.signal.normalization.fit import M_2_PI, XForm
from simurg.signal.normalization.numerics import TINY

FloatArr = npt.NDArray[np.float32]
XFormRecords = npt.NDArray[np.void]  # GTC normalization_transforms record array


def transforms_to_arrays(xforms: Sequence[XForm] | XFormRecords) -> dict[str, FloatArr]:
    """Per-bin parameter arrays from XForm objects or a GTC ``normalization_transforms`` record array."""
    if isinstance(xforms, np.ndarray):
        get = {
            k: xforms[k].astype(np.float32) for k in ("offset_x", "offset_y", "scale_x", "scale_y", "shear", "theta")
        }
    else:
        get = {
            k: np.array([getattr(x, k) for x in xforms], dtype=np.float32)
            for k in ("offset_x", "offset_y", "scale_x", "scale_y", "shear", "theta")
        }
    get["cos"] = np.array([math.cos(float(t)) for t in get["theta"]], dtype=np.float64).astype(np.float32)
    get["sin"] = np.array([math.sin(float(t)) for t in get["theta"]], dtype=np.float64).astype(np.float32)
    return get


def normalize(
    raw_x: npt.NDArray[np.uint16],
    raw_y: npt.NDArray[np.uint16],
    lookup: npt.NDArray[np.int32],
    xforms: Sequence[XForm] | XFormRecords,
) -> tuple[FloatArr, FloatArr]:
    """Normalized X/Y per locus; ``lookup[i]`` is locus i's transform index (Manifest.normalization_lookups)."""
    p = transforms_to_arrays(xforms)
    ox, oy = p["offset_x"][lookup], p["offset_y"][lookup]
    c, s = p["cos"][lookup], p["sin"][lookup]
    sh, sx, sy = p["shear"][lookup], p["scale_x"][lookup], p["scale_y"][lookup]
    tx = raw_x.astype(np.float32) - ox
    ty = raw_y.astype(np.float32) - oy
    tx2 = c * tx + s * ty
    ty2 = (-s) * tx + c * ty
    tx3 = tx2 - sh * ty2
    with np.errstate(divide="ignore", invalid="ignore"):
        nx = np.where(tx3 < 0, np.float32(0), tx3 / sx).astype(np.float32)
        ny = np.where(ty2 < 0, np.float32(0), ty2 / sy).astype(np.float32)
    return nx, ny


def polar(norm_x: FloatArr, norm_y: FloatArr) -> tuple[FloatArr, FloatArr]:
    """(θ, R): θ = (2/π)·atan2(y, x), R = |x| + |y|; both NaN where x = y = 0 (plan §7.4)."""
    theta = (np.arctan2(norm_y.astype(np.float64), norm_x.astype(np.float64)) * M_2_PI).astype(np.float32)
    r = np.where((norm_x < 0) & (norm_y < 0), TINY, np.abs(norm_x) + np.abs(norm_y)).astype(np.float32)
    zero = (norm_x == 0) & (norm_y == 0)
    theta[zero] = np.nan
    r[zero] = np.nan
    return theta, r