
"""
pyFANG - Python implementation of TFANG (Focused Angular N-body event Generator)

Self-contained single-file implementation with Numba JIT optimization.

Authors: Arik Kreisel, Itay Horin

Required packages: numpy, numba
Optional: matplotlib (for plotting)
"""

from __future__ import annotations

import numpy as np
from numba import njit
from dataclasses import dataclass
from typing import List, Tuple, Optional
import warnings


PI = np.pi
TWO_PI = 2.0 * PI
FOUR_PI = 4.0 * PI

POSITION_TOLERANCE = 1e-6
MOMENTUM_TOLERANCE = 1e-12

MODE_POINT = 2.0
MODE_CIRCLE = 0.0

kPi = PI
kTwoPi = TWO_PI
kFourPi = FOUR_PI
kPositionTolerance = POSITION_TOLERANCE
kMomentumTolerance = MOMENTUM_TOLERANCE
kModePoint = MODE_POINT
kModeCircle = MODE_CIRCLE



@njit(fastmath=True)
def is_point(shape: float) -> bool:
    """Check if shape parameter indicates point generation mode."""
    return shape == MODE_POINT


@njit(fastmath=True)
def is_circle(shape: float) -> bool:
    """Check if shape parameter indicates circle generation mode."""
    return shape == MODE_CIRCLE


@njit(fastmath=True)
def is_strip(shape: float) -> bool:
    """Check if shape parameter indicates strip generation mode."""
    return 0.0 < shape <= 1.0


@njit(fastmath=True)
def is_ring(shape: float) -> bool:
    """Check if shape parameter indicates ring generation mode."""
    return shape < 0.0


@njit(fastmath=True)
def calc_km_factor(x: float, y: float) -> float:
    """Phase space kinematic function F(x,y) = sqrt((1-x-y)^2 - 4xy)."""
    arg = (1.0 - x - y)**2 - 4.0 * x * y
    if arg < 0:
        return 0.0
    return np.sqrt(arg)


@njit(fastmath=True)
def factorial_numba(n: int) -> float:
    """Factorial function compatible with Numba."""
    if n <= 1:
        return 1.0
    result = 1.0
    for i in range(2, n + 1):
        result *= float(i)
    return result


@njit(fastmath=True)
def normalize_vec3(v: np.ndarray) -> np.ndarray:
    """Normalize a 3-vector."""
    mag = np.sqrt(v[0]**2 + v[1]**2 + v[2]**2)
    if mag < 1e-15:
        return np.array([0.0, 0.0, 1.0])
    return v / mag


@njit(fastmath=True)
def vec3_theta(v: np.ndarray) -> float:
    """Get polar angle of a 3-vector."""
    r = np.sqrt(v[0]**2 + v[1]**2 + v[2]**2)
    if r < 1e-15:
        return 0.0
    cos_theta = v[2] / r
    if cos_theta > 1.0:
        cos_theta = 1.0
    if cos_theta < -1.0:
        cos_theta = -1.0
    return np.arccos(cos_theta)


@njit(fastmath=True)
def vec3_phi(v: np.ndarray) -> float:
    """Get azimuthal angle of a 3-vector."""
    return np.arctan2(v[1], v[0])


@njit(fastmath=True)
def cross_product(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    """Cross product of two 3-vectors."""
    return np.array([
        a[1] * b[2] - a[2] * b[1],
        a[2] * b[0] - a[0] * b[2],
        a[0] * b[1] - a[1] * b[0]
    ])



@njit(fastmath=True)
def lorentz_from_pxpypze(px: float, py: float, pz: float, e: float) -> np.ndarray:
    """Create Lorentz vector (px, py, pz, m) from (px, py, pz, E)."""
    p2 = px*px + py*py + pz*pz
    m2 = e*e - p2
    m = np.sqrt(max(0.0, m2))
    return np.array([px, py, pz, m])


@njit(fastmath=True)
def lorentz_energy(p: np.ndarray) -> float:
    """Energy component of Lorentz vector."""
    return np.sqrt(p[0]**2 + p[1]**2 + p[2]**2 + p[3]**2)


@njit(fastmath=True)
def lorentz_p_mag(p: np.ndarray) -> float:
    """Magnitude of 3-momentum."""
    return np.sqrt(p[0]**2 + p[1]**2 + p[2]**2)


@njit(fastmath=True)
def lorentz_beta(p: np.ndarray) -> float:
    """Velocity beta = v/c = P/E."""
    e = lorentz_energy(p)
    if e == 0:
        return 0.0
    return lorentz_p_mag(p) / e


@njit(fastmath=True)
def lorentz_gamma(p: np.ndarray) -> float:
    """Lorentz factor gamma = E/M."""
    if p[3] == 0:
        return 1e15
    return lorentz_energy(p) / p[3]


@njit(fastmath=True)
def lorentz_boost_to_cm(p: np.ndarray) -> np.ndarray:
    """Return boost vector to center-of-mass frame."""
    e = lorentz_energy(p)
    if e == 0:
        return np.array([0.0, 0.0, 0.0])
    return np.array([-p[0]/e, -p[1]/e, -p[2]/e])


@njit(fastmath=True)
def lorentz_boost(p: np.ndarray, beta: np.ndarray) -> np.ndarray:
    """Apply a Lorentz boost to a 4-vector. Returns new 4-vector (px, py, pz, m)."""
    beta_mag = np.sqrt(beta[0]**2 + beta[1]**2 + beta[2]**2)
    if beta_mag < 1e-15:
        return p.copy()
    if beta_mag >= 1.0:
        return p.copy()

    gamma = 1.0 / np.sqrt(1.0 - beta_mag**2)
    e = lorentz_energy(p)
    beta_dot_p = beta[0]*p[0] + beta[1]*p[1] + beta[2]*p[2]
    factor = (gamma - 1.0) * beta_dot_p / (beta_mag**2) + gamma * e

    new_px = p[0] + factor * beta[0]
    new_py = p[1] + factor * beta[1]
    new_pz = p[2] + factor * beta[2]
    new_e = gamma * (e + beta_dot_p)

    return lorentz_from_pxpypze(new_px, new_py, new_pz, new_e)


@njit(fastmath=True)
def lorentz_boost_to_cm_frame(p: np.ndarray, system: np.ndarray) -> np.ndarray:
    """Boost vector p to the rest frame of system."""
    return lorentz_boost(p, lorentz_boost_to_cm(system))


@njit(fastmath=True)
def lorentz_boost_from_cm_frame(p: np.ndarray, system: np.ndarray) -> np.ndarray:
    """Boost vector p from the rest frame of system to the lab frame."""
    beta = lorentz_boost_to_cm(system)
    return lorentz_boost(p, -beta)



@njit(fastmath=True)
def two_body_numba(S: np.ndarray, m1: float, m2: float,
                   rnd1: float, rnd2: float) -> tuple:
    """Generate isotropic two-body decay. Returns tuple of (p1, p2) as 4-vectors."""
    cst = 2.0 * rnd1 - 1.0
    snt = np.sqrt(1.0 - cst**2)
    phi = TWO_PI * rnd2

    M = S[3]
    E1 = (M**2 - m2**2 + m1**2) / (2.0 * M)
    p_sq = E1**2 - m1**2
    if p_sq < 0:
        p_sq = 0.0
    sp = np.sqrt(p_sq)

    p1_cm = np.array([sp * snt * np.cos(phi), sp * snt * np.sin(phi), sp * cst, m1])
    p2_cm = np.array([-sp * snt * np.cos(phi), -sp * snt * np.sin(phi), -sp * cst, m2])

    p1 = lorentz_boost_from_cm_frame(p1_cm, S)
    p1[3] = m1
    p2 = lorentz_boost_from_cm_frame(p2_cm, S)
    p2[3] = m2

    return p1, p2


@njit(fastmath=True)
def t_gen_point_space_numba(S1: np.ndarray, m1: float, m2: float,
                            v_det: np.ndarray) -> tuple:
    """Calculate 4-momentum for particle constrained to a lab-frame direction."""
    P_mag = lorentz_p_mag(S1)
    if P_mag < 1e-15:
        V_Su = np.array([0.0, 0.0, 1.0])
    else:
        V_Su = np.array([S1[0] / P_mag, S1[1] / P_mag, S1[2] / P_mag])
    V_Su = normalize_vec3(V_Su)

    beta_VS = -lorentz_beta(S1) * V_Su
    v_det = normalize_vec3(v_det)
    cos_LAB = V_Su[0]*v_det[0] + V_Su[1]*v_det[1] + V_Su[2]*v_det[2]

    m_CM = S1[3]
    E_CM = lorentz_energy(S1)
    p_CM = lorentz_p_mag(S1)
    gamma1 = lorentz_gamma(S1)

    CM_E3 = (m_CM**2 - m2**2 + m1**2) / (2.0 * m_CM)

    if CM_E3**2 - m1**2 < 0:
        return 0, np.zeros(0), np.zeros((0, 4)), np.zeros((0, 4))

    CM_p3 = np.sqrt(CM_E3**2 - m1**2)

    aa = p_CM**2 * cos_LAB**2 - E_CM**2
    bb = 2.0 * p_CM * cos_LAB * CM_E3 * m_CM
    cc = (CM_E3 * m_CM)**2 - m1**2 * E_CM**2

    discriminant = bb**2 - 4.0 * aa * cc
    if discriminant < 0:
        return 0, np.zeros(0), np.zeros((0, 4)), np.zeros((0, 4))

    sqrt_disc = np.sqrt(discriminant)
    sol1 = (-bb + sqrt_disc) / (2.0 * aa)
    sol2 = (-bb - sqrt_disc) / (2.0 * aa)

    solutions = 0
    p3_LAB = np.zeros(2)

    if sol1 > 0:
        p3_LAB[solutions] = sol1
        solutions += 1
    if sol2 > 0 and abs(sol2 - sol1) > 1e-12:
        p3_LAB[solutions] = sol2
        solutions += 1

    if solutions == 0:
        return 0, np.zeros(0), np.zeros((0, 4)), np.zeros((0, 4))

    jack_pdf = np.zeros(solutions)
    p_det = np.zeros((solutions, 4))
    p_d2 = np.zeros((solutions, 4))

    pdf_CM = 1.0 / FOUR_PI

    for idx in range(solutions):
        p_lab = p3_LAB[idx]
        p_det_vec = np.array([p_lab * v_det[0], p_lab * v_det[1], p_lab * v_det[2], m1])

        p3_CM = lorentz_boost(p_det_vec, beta_VS)
        p3_CM[3] = m1

        p4_CM = np.array([-p3_CM[0], -p3_CM[1], -p3_CM[2], m2])
        p_d2_vec = lorentz_boost(p4_CM, -beta_VS)
        p_d2_vec[3] = m2

        cos_CM = (p3_CM[0]*V_Su[0] + p3_CM[1]*V_Su[1] + p3_CM[2]*V_Su[2]) / CM_p3
        qqq = p_CM * lorentz_energy(p3_CM) / (E_CM * lorentz_p_mag(p3_CM))

        if abs(cos_LAB) > 0.99 and abs(cos_CM) > 0.99:
            Jack = (gamma1**2 * cos_CM**2 * (1.0 + qqq * cos_CM)**2 / (cos_LAB**2))
        else:
            sin2_CM = (1.0 - cos_CM) * (1.0 + cos_CM)
            sin2_LAB = (1.0 - cos_LAB) * (1.0 + cos_LAB)
            if sin2_LAB < 1e-15:
                Jack = 0.0
            else:
                Jack = (sin2_CM / sin2_LAB * np.sqrt(sin2_CM / sin2_LAB) / gamma1 / (1.0 + qqq * cos_CM))

        jack_pdf[idx] = abs(pdf_CM * Jack)
        p_det[idx] = p_det_vec
        p_d2[idx] = p_d2_vec

    return solutions, jack_pdf, p_det, p_d2


@njit(fastmath=True)
def t_gen_vec_numba(omega: float, ratio: float, v_center: np.ndarray,
                    rnd1: float, rnd2: float) -> np.ndarray:
    """Generate random direction vector within specified solid angle."""
    if omega > FOUR_PI or omega < 0:
        omega = FOUR_PI
    if ratio > 1.0:
        ratio = 0.0

    v_center = normalize_vec3(v_center)

    if is_circle(ratio):
        cst = 1.0 - omega / TWO_PI * rnd1
        phi = TWO_PI * rnd2

        if abs(v_center[0]) < POSITION_TOLERANCE and abs(v_center[1]) < POSITION_TOLERANCE:
            if v_center[2] > 0:
                theta = np.arccos(cst)
            else:
                theta = np.arccos(-cst)
            return np.array([np.sin(theta) * np.cos(phi), np.sin(theta) * np.sin(phi), np.cos(theta)])
        else:
            Vz = np.array([0.0, 0.0, 1.0])
            newZ = v_center
            newY = normalize_vec3(cross_product(newZ, Vz))
            newX = normalize_vec3(cross_product(newY, newZ))

            theta = np.arccos(cst)
            v_gen_x = np.sin(theta) * np.cos(phi)
            v_gen_y = np.sin(theta) * np.sin(phi)
            v_gen_z = np.cos(theta)

            return np.array([
                newX[0]*v_gen_x + newY[0]*v_gen_y + newZ[0]*v_gen_z,
                newX[1]*v_gen_x + newY[1]*v_gen_y + newZ[1]*v_gen_z,
                newX[2]*v_gen_x + newY[2]*v_gen_y + newZ[2]*v_gen_z
            ])

    elif is_strip(ratio):
        Dphi = ratio * TWO_PI
        Dcos = omega / Dphi
        phi0 = vec3_phi(v_center)
        cst0 = np.cos(vec3_theta(v_center))

        if cst0 > (1.0 - Dcos / 2.0):
            cst0 = 1.0 - Dcos / 2.0
        if cst0 < (-1.0 + Dcos / 2.0):
            cst0 = -1.0 + Dcos / 2.0

        cst = cst0 - Dcos / 2.0 + Dcos * rnd1
        phi = phi0 - Dphi / 2.0 + Dphi * rnd2
        theta = np.arccos(cst)

        return np.array([np.sin(theta) * np.cos(phi), np.sin(theta) * np.sin(phi), np.cos(theta)])

    elif is_ring(ratio):
        cst = 1.0 - omega / TWO_PI
        phi = TWO_PI * rnd2

        if abs(v_center[0]) < POSITION_TOLERANCE and abs(v_center[1]) < POSITION_TOLERANCE:
            if v_center[2] > 0:
                theta = np.arccos(cst)
            else:
                theta = np.arccos(-cst)
            return np.array([np.sin(theta) * np.cos(phi), np.sin(theta) * np.sin(phi), np.cos(theta)])
        else:
            Vz = np.array([0.0, 0.0, 1.0])
            newZ = v_center
            newY = normalize_vec3(cross_product(newZ, Vz))
            newX = normalize_vec3(cross_product(newY, newZ))

            theta = np.arccos(cst)
            v_gen_x = np.sin(theta) * np.cos(phi)
            v_gen_y = np.sin(theta) * np.sin(phi)
            v_gen_z = np.cos(theta)

            return np.array([
                newX[0]*v_gen_x + newY[0]*v_gen_y + newZ[0]*v_gen_z,
                newX[1]*v_gen_x + newY[1]*v_gen_y + newZ[1]*v_gen_z,
                newX[2]*v_gen_x + newY[2]*v_gen_y + newZ[2]*v_gen_z
            ])

    return v_center


@njit(fastmath=True)
def gen_fang_unconstrained_numba(n_body: int, S: np.ndarray, masses: np.ndarray,
                                  random_array: np.ndarray) -> tuple:
    """Fast unconstrained N-body phase space generation."""
    m_S = S[3]
    mall = np.sum(masses)

    if mall >= m_S:
        return np.zeros((n_body, 4)), 0.0

    if n_body == 2:
        wh_PS = calc_km_factor(masses[0]**2 / (S[3]**2), masses[1]**2 / (S[3]**2)) * PI / 2.0
        p1, p2 = two_body_numba(S, masses[0], masses[1], random_array[0], random_array[1])
        result = np.zeros((2, 4))
        result[0] = p1
        result[1] = p2
        return result, wh_PS

    n_random = n_body - 2
    rrr = random_array[:n_random].copy()
    rrr.sort()

    m_V = np.zeros(n_body - 2)
    for i in range(n_body - 2):
        m_A = 0.0
        for j in range(i+1, n_body):
            m_A += masses[j]
        m_V[i] = rrr[n_body - 2 - i - 1] * (m_S - mall) + m_A

    wh_PS = m_V[0] * calc_km_factor(masses[0]**2 / (S[3]**2), m_V[0]**2 / (S[3]**2))

    for i in range(n_body - 3):
        wh_PS *= m_V[i+1] * calc_km_factor(masses[i+1]**2 / (m_V[i]**2), m_V[i+1]**2 / (m_V[i]**2))

    if n_body > 2:
        wh_PS *= calc_km_factor(masses[n_body-2]**2 / (m_V[n_body-3]**2), masses[n_body-1]**2 / (m_V[n_body-3]**2))

    wh_PS *= (PI**(n_body - 1) / 2.0 * (m_S - mall)**(n_body - 2) / factorial_numba(n_body - 2))

    result = np.zeros((n_body, 4))
    rnd_idx = n_random
    p1, p2 = two_body_numba(S, masses[0], m_V[0], random_array[rnd_idx], random_array[rnd_idx+1])
    result[0] = p1
    p_V = p2
    rnd_idx += 2

    for i in range(n_body - 3):
        p1, p2 = two_body_numba(p_V, masses[i+1], m_V[i+1], random_array[rnd_idx], random_array[rnd_idx+1])
        result[i+1] = p1
        p_V = p2
        rnd_idx += 2

    p1, p2 = two_body_numba(p_V, masses[n_body-2], masses[n_body-1], random_array[rnd_idx], random_array[rnd_idx+1])
    result[n_body-2] = p1
    result[n_body-1] = p2

    return result, wh_PS


@njit(fastmath=True)
def batch_generate_unconstrained(n_events: int, n_body: int, S: np.ndarray,
                                 masses: np.ndarray, random_numbers: np.ndarray) -> tuple:
    """Batch generation of unconstrained events."""
    all_momenta = np.zeros((n_events, n_body, 4))
    all_weights = np.zeros(n_events)

    for i in range(n_events):
        random_array = random_numbers[i]
        momenta, weight = gen_fang_unconstrained_numba(n_body, S, masses, random_array)
        all_momenta[i] = momenta
        all_weights[i] = weight

    return all_momenta, all_weights



def IsPoint(shape: float) -> bool:
    return is_point(shape)

def IsCircle(shape: float) -> bool:
    return is_circle(shape)

def IsStrip(shape: float) -> bool:
    return is_strip(shape)

def IsRing(shape: float) -> bool:
    return is_ring(shape)

def CalcKMFactor(x: float, y: float) -> float:
    return calc_km_factor(x, y)



@dataclass
class LorentzVector:
    """A 4-momentum vector (px, py, pz, m)."""
    px: float = 0.0
    py: float = 0.0
    pz: float = 0.0
    m: float = 0.0

    def to_array(self) -> np.ndarray:
        return np.array([self.px, self.py, self.pz, self.m])

    @classmethod
    def from_array(cls, arr: np.ndarray) -> 'LorentzVector':
        return cls(arr[0], arr[1], arr[2], arr[3])

    @property
    def E(self) -> float:
        return lorentz_energy(self.to_array())

    @property
    def P(self) -> float:
        return lorentz_p_mag(self.to_array())

    @property
    def M(self) -> float:
        return self.m



@dataclass
class Detector:
    """Base class for detector constraints.

    Every detector has a direction vector (3D) specifying where the
    detector is pointing. Subclasses define the detector shape and
    angular coverage.
    """
    direction: np.ndarray

    def _omega(self) -> float:
        """Return the solid angle parameter for the internal engine."""
        raise NotImplementedError

    def _ratio(self) -> float:
        """Return the shape mode parameter for the internal engine."""
        raise NotImplementedError

    def __post_init__(self):
        self.direction = np.asarray(self.direction, dtype=np.float64)


@dataclass
class CircleDetector(Detector):
    """Circular cone detector.

    Particles are generated anywhere inside a cone of the given solid
    angle centered on the direction vector.
    """
    solid_angle: float = 0.0

    def _omega(self) -> float:
        return self.solid_angle

    def _ratio(self) -> float:
        return MODE_CIRCLE


@dataclass
class PointDetector(Detector):
    """Fixed-direction point detector.

    The particle momentum direction is fixed exactly to the given
    direction vector. No solid angle parameter is needed.
    """

    def _omega(self) -> float:
        return 1.0

    def _ratio(self) -> float:
        return MODE_POINT


@dataclass
class StripDetector(Detector):
    """Strip-shaped detector.

    A rectangular region in (theta, phi) space. azimuthal_coverage
    is the fraction of 2*pi azimuthal angle covered (0 < value <= 1).
    """
    solid_angle: float = 0.0
    azimuthal_coverage: float = 0.0

    def _omega(self) -> float:
        return self.solid_angle

    def _ratio(self) -> float:
        return self.azimuthal_coverage


@dataclass
class RingDetector(Detector):
    """Ring/annular detector.

    An annular region around the direction vector. opening controls
    the ring geometry (positive value, negated internally).
    """
    solid_angle: float = 0.0
    opening: float = 0.0

    def _omega(self) -> float:
        return self.solid_angle

    def _ratio(self) -> float:
        return -abs(self.opening)


def gen_fang_fast(n_body: int, S: LorentzVector, masses: np.ndarray,
                  n_events: int = 1, seed: int = None) -> Tuple[np.ndarray, np.ndarray]:
    """Fast batch generation for unconstrained phase space."""
    if seed is None:
        seed = np.random.randint(0, 2**31)
    rng = np.random.default_rng(seed)
    n_random_per_event = (n_body - 2) + 2 * (n_body - 1)
    random_numbers = rng.random((n_events, n_random_per_event))
    S_array = S.to_array()
    return batch_generate_unconstrained(n_events, n_body, S_array, masses, random_numbers)


def calculate_phase_space_fast(n_body: int, S: LorentzVector, masses: np.ndarray,
                               n_events: int = 100000, seed: int = None) -> Tuple[float, float]:
    """Fast phase space volume calculation."""
    all_momenta, all_weights = gen_fang_fast(n_body, S, masses, n_events, seed)
    sum_w = np.sum(all_weights)
    sum_w2 = np.sum(all_weights**2)
    n_success = len(all_weights)
    if n_success == 0:
        return 0.0, 0.0
    phase_space = sum_w / n_success
    error = np.sqrt(sum_w2) / n_success
    return phase_space, error



class Node_t:
    """Binary tree node for tracking multiple kinematic solutions."""

    def __init__(self, p1: np.ndarray, p2: np.ndarray, weight: float, parent: Optional['Node_t']):
        self.fPV = p2.copy()
        self.fPDet = p1.copy()
        self.fWeight = weight
        self.fLeft: Optional[Node_t] = None
        self.fRight: Optional[Node_t] = None
        self.fParent = parent



def DeleteTree(node: Optional[Node_t]) -> None:
    """Recursively delete a tree."""
    if node is None:
        return
    DeleteTree(node.fLeft)
    DeleteTree(node.fRight)
    node.fLeft = None
    node.fRight = None
    node.fParent = None


def CreateFirst(node: Optional[Node_t], p1: np.ndarray, p2: np.ndarray, weight: float) -> Node_t:
    """Create the first (root) node of the tree."""
    if node is None:
        return Node_t(p1, p2, weight, None)
    return node


def CreateRight(node: Optional[Node_t], tmp: Optional[Node_t],
                p1: np.ndarray, p2: np.ndarray, weight: float) -> Node_t:
    """Create a right child node."""
    if node is None:
        return Node_t(p1, p2, weight, tmp)
    node.fRight = CreateRight(node.fRight, node, p1, p2, weight)
    return node


def CreateLeft(node: Optional[Node_t], tmp: Optional[Node_t],
               p1: np.ndarray, p2: np.ndarray, weight: float) -> Node_t:
    """Create a left child node."""
    if node is None:
        return Node_t(p1, p2, weight, tmp)
    node.fLeft = CreateLeft(node.fLeft, node, p1, p2, weight)
    return node


def CollectPaths(nBody: int, node: Optional[Node_t],
                 path: List[np.ndarray], paths: List[List[np.ndarray]]) -> None:
    """Collect all root-to-leaf paths for 4-momenta."""
    if node is None:
        return
    path.append(node.fPDet.copy())
    if node.fLeft is None and node.fRight is None and len(path) == nBody + 1:
        paths.append([p.copy() for p in path])
    else:
        CollectPaths(nBody, node.fLeft, path, paths)
        CollectPaths(nBody, node.fRight, path, paths)
    path.pop()


def CollectPathsWeights(nBody: int, node: Optional[Node_t],
                        path: List[float], paths: List[List[float]]) -> None:
    """Collect all root-to-leaf paths for weights."""
    if node is None:
        return
    path.append(node.fWeight)
    if node.fLeft is None and node.fRight is None and len(path) == nBody + 1:
        paths.append(path.copy())
    else:
        CollectPathsWeights(nBody, node.fLeft, path, paths)
        CollectPathsWeights(nBody, node.fRight, path, paths)
    path.pop()



def TwoBody(S: np.ndarray, m1: float, m2: float,
            rng: np.random.Generator) -> Tuple[np.ndarray, np.ndarray]:
    """Generate isotropic two-body decay."""
    rnd1 = rng.random()
    rnd2 = rng.random()
    return two_body_numba(S, m1, m2, rnd1, rnd2)


def TGenPointSpace(S1: np.ndarray, m1: float, m2: float,
                   vDet: np.ndarray) -> Tuple[bool, int, np.ndarray, np.ndarray, np.ndarray]:
    """Calculate 4-momentum for particle constrained to a lab-frame direction."""
    solutions, jackPDF, pDet, pD2 = t_gen_point_space_numba(S1, m1, m2, vDet)
    hit = solutions > 0
    return hit, solutions, jackPDF, pDet, pD2


def TGenVec(Omega: float, Ratio: float, Vcenter: np.ndarray,
            rng: np.random.Generator) -> np.ndarray:
    """Generate random direction vector within specified solid angle."""
    if Omega > kFourPi or Omega < 0:
        Omega = kFourPi
        warnings.warn("TFANG::TGenVec: Omega out of range, set to 4pi")
    if Ratio > 1.0:
        Ratio = 0.0
        warnings.warn("TFANG::TGenVec: Ratio out of range, set to 0")
    rnd1 = rng.random()
    rnd2 = rng.random()
    return t_gen_vec_numba(Omega, Ratio, Vcenter, rnd1, rnd2)


def TwoBodyWithRandoms(S: np.ndarray, m1: float, m2: float,
                       rnd1: float, rnd2: float) -> Tuple[np.ndarray, np.ndarray]:
    """Generate isotropic two-body decay using pre-generated random numbers."""
    return two_body_numba(S, m1, m2, rnd1, rnd2)


def TGenVecWithRandoms(Omega: float, Ratio: float, Vcenter: np.ndarray,
                       rnd1: float, rnd2: float) -> np.ndarray:
    """Generate random direction vector using pre-generated randoms."""
    if Omega > kFourPi or Omega < 0:
        Omega = kFourPi
    if Ratio > 1.0:
        Ratio = 0.0
    return t_gen_vec_numba(Omega, Ratio, Vcenter, rnd1, rnd2)


def calc_n_random_per_event(nBody: int, nDet: int) -> int:
    """Calculate the number of random numbers needed per event."""
    return (nBody - 2) + 2 * nDet + 2 * (nBody - 1 - nDet)



def GenFANG(nBody: int, S: np.ndarray, masses: np.ndarray,
            Om: Optional[np.ndarray], Ratio: Optional[np.ndarray],
            V3Det: List[np.ndarray], rng: np.random.Generator) -> Tuple[List[List[np.ndarray]], List[float]]:
    """Generate phase-space events with angular constraints."""
    nDet = len(V3Det)
    mS = S[3]
    wh = 1.0
    mall = np.sum(masses)

    if mall >= mS:
        warnings.warn(f"TFANG::GenFANG: Sum of decay masses ({mall}) >= parent mass ({mS})")
        return [], []

    VecVecP: List[List[np.ndarray]] = []
    vecWi: List[float] = []

    if nBody == 2:
        whPS = CalcKMFactor(masses[0]**2 / S[3]**2, masses[1]**2 / S[3]**2) * kPi / 2.0

        if nDet == 1:
            if IsPoint(Ratio[0]):
                V3 = normalize_vec3(V3Det[0])
            else:
                V3 = TGenVec(Om[0], Ratio[0], normalize_vec3(V3Det[0]), rng)

            Hit, solutions, jackPDF, pDet, pD2 = TGenPointSpace(S, masses[0], masses[1], V3)
            if not Hit:
                return [], []

            for l in range(solutions):
                vecP = [pDet[l].copy(), pD2[l].copy()]
                VecVecP.append(vecP)
                vecWi.append(jackPDF[l] * whPS)
        else:
            p1, p2 = TwoBody(S, masses[0], masses[1], rng)
            vecP = [p1, p2]
            wh = 1.0
            vecWi.append(wh * whPS)
            VecVecP.append(vecP)

        return VecVecP, vecWi

    rrr = np.sort(rng.random(nBody - 2))

    mV = np.zeros(nBody - 2)
    for i in range(nBody - 2):
        mB = 0.0
        mA = 0.0
        for l in range(i + 1):
            mB += masses[l]
        for l in range(i + 1, nBody):
            mA += masses[l]
        mV[i] = rrr[nBody - 2 - i - 1] * (mS - mall) + mA

        if i > 0 and (mV[i - 1] - mV[i]) < masses[i]:
            warnings.warn(f"TFANG::GenFANG: Virtual mass constraint violated at i={i}")

    whPS = mV[0] * CalcKMFactor(masses[0]**2 / S[3]**2, mV[0]**2 / S[3]**2)

    for i in range(nBody - 3):
        whPS *= mV[i + 1] * CalcKMFactor(masses[i + 1]**2 / (mV[i]**2), mV[i + 1]**2 / (mV[i]**2))

    if nBody > 2:
        whPS *= CalcKMFactor(masses[nBody - 2]**2 / (mV[nBody - 3]**2), masses[nBody - 1]**2 / (mV[nBody - 3]**2))

    whPS *= (kPi**(nBody - 1) / 2.0 * (mS - mall)**(nBody - 2) / factorial_numba(nBody - 2))

    if nDet == 0:
        vecP: List[np.ndarray] = []
        p1, p2 = TwoBody(S, masses[0], mV[0], rng)
        vecP.append(p1)
        pV = p2

        for i in range(nBody - 3):
            p1, p2 = TwoBody(pV, masses[i + 1], mV[i + 1], rng)
            vecP.append(p1)
            pV = p2

        p1, p2 = TwoBody(pV, masses[nBody - 2], masses[nBody - 1], rng)
        vecP.append(p1)
        vecP.append(p2)

        wh = 1.0
        vecWi.append(wh * whPS)
        VecVecP.append(vecP)
        return VecVecP, vecWi

    root = CreateFirst(None, S, S, 1.0)
    cur = root
    level = 0
    branch: List[int] = []

    while level < (nBody - 1):
        if level < nDet and level < (nBody - 2):
            pV = cur.fPV
            if IsPoint(Ratio[level]):
                V3 = normalize_vec3(V3Det[level])
            else:
                V3 = TGenVec(Om[level], Ratio[level], normalize_vec3(V3Det[level]), rng)

            Hit, solutions, jackPDF, pDet, pD2 = TGenPointSpace(pV, masses[level], mV[level], V3)

            if solutions == 0 and len(branch) == 0:
                level = nBody - 1
                continue
            if solutions == 0 and len(branch) > 0:
                while level > branch[-1]:
                    level -= 1
                    cur = cur.fParent
                cur = cur.fLeft
                level += 1
                branch.pop()
                continue
            if solutions == 1:
                CreateRight(cur, None, pDet[0], pD2[0], jackPDF[0])
                cur = cur.fRight
                level += 1
                continue
            elif solutions == 2:
                branch.append(level)
                CreateLeft(cur, None, pDet[1], pD2[1], jackPDF[1])
                CreateRight(cur, None, pDet[0], pD2[0], jackPDF[0])
                cur = cur.fRight
                level += 1
                continue

        if level < nDet and level == (nBody - 2):
            pV = cur.fPV
            if IsPoint(Ratio[level]):
                V3 = normalize_vec3(V3Det[level])
            else:
                V3 = TGenVec(Om[level], Ratio[level], normalize_vec3(V3Det[level]), rng)

            Hit, solutions, jackPDF, pDet, pD2 = TGenPointSpace(pV, masses[level], masses[level + 1], V3)

            if solutions == 0 and len(branch) == 0:
                level = nBody - 1
                continue
            if solutions == 0 and len(branch) > 0:
                while level > branch[-1]:
                    level -= 1
                    cur = cur.fParent
                cur = cur.fLeft
                level += 1
                branch.pop()
                continue
            if solutions == 1:
                CreateRight(cur, None, pDet[0], pD2[0], jackPDF[0])
                cur = cur.fRight
                CreateRight(cur, None, pD2[0], S, 1.0)
                cur = cur.fRight
                level += 1
                continue
            elif solutions == 2:
                CreateRight(cur, None, pDet[0], pD2[0], jackPDF[0])
                cur = cur.fRight
                CreateRight(cur, None, pD2[0], S, 1.0)
                cur = cur.fParent
                CreateLeft(cur, None, pDet[1], pD2[1], jackPDF[1])
                cur = cur.fLeft
                CreateRight(cur, None, pD2[1], S, 1.0)
                cur = cur.fRight
                level += 1

                if level == (nBody - 1) and len(branch) > 0:
                    while level > branch[-1]:
                        level -= 1
                        cur = cur.fParent
                    cur = cur.fParent
                    cur = cur.fLeft
                    level += 1
                    branch.pop()
                    continue
                continue

        if level >= nDet and level < nBody - 2:
            pV = cur.fPV
            p1, p2 = TwoBody(pV, masses[level], mV[level], rng)
            CreateRight(cur, None, p1, p2, 1.0)
            cur = cur.fRight
            level += 1
            continue

        if level >= nDet and level == nBody - 2:
            pV = cur.fPV
            p1, p2 = TwoBody(pV, masses[level], masses[level + 1], rng)
            CreateRight(cur, None, p1, p2, 1.0)
            cur = cur.fRight
            CreateRight(cur, None, p2, S, 1.0)
            cur = cur.fRight
            level += 1

            if level == (nBody - 1) and len(branch) > 0:
                while level > branch[-1]:
                    level -= 1
                    cur = cur.fParent
                cur = cur.fParent
                cur = cur.fLeft
                level += 1
                branch.pop()
                continue
            continue

        if level == (nBody - 1) and len(branch) > 0:
            while level > branch[-1]:
                level -= 1
                cur = cur.fParent
            cur = cur.fParent
            cur = cur.fLeft
            level += 1
            branch.pop()
            continue

    pathsJ: List[List[float]] = []
    pathsP: List[List[np.ndarray]] = []
    vecJ: List[float] = []
    vecP_collect: List[np.ndarray] = []

    CollectPathsWeights(nBody, root, vecJ, pathsJ)
    CollectPaths(nBody, root, vecP_collect, pathsP)

    DeleteTree(root)
    root = None

    if len(pathsP) == 0:
        return [], []

    for i in range(len(pathsJ)):
        vecJ = pathsJ[i]
        vecP = pathsP[i]
        vecP = vecP[1:]
        VecVecP.append(vecP)

        wh = 1.0
        for j in range(1, nDet + 1):
            wh *= vecJ[j]
        vecWi.append(wh * whPS)

    return VecVecP, vecWi



def GenFANGWithRandoms(nBody: int, S: np.ndarray, masses: np.ndarray,
                       Om: Optional[np.ndarray], Ratio: Optional[np.ndarray],
                       V3Det: List[np.ndarray], random_array: np.ndarray,
                       rng: Optional[np.random.Generator] = None) -> Tuple[List[List[np.ndarray]], List[float]]:
    """Generate phase-space events with angular constraints using pre-generated random numbers."""
    nDet = len(V3Det)
    mS = S[3]
    wh = 1.0
    mall = np.sum(masses)

    if mall >= mS:
        return [], []

    VecVecP: List[List[np.ndarray]] = []
    vecWi: List[float] = []

    rnd_idx = [0]

    def get_random():
        idx = rnd_idx[0]
        rnd_idx[0] += 1
        if idx < len(random_array):
            return random_array[idx]
        elif rng is not None:
            return rng.random()
        else:
            return np.random.random()

    if nBody == 2:
        whPS = CalcKMFactor(masses[0]**2 / S[3]**2, masses[1]**2 / S[3]**2) * kPi / 2.0

        if nDet == 1:
            if IsPoint(Ratio[0]):
                V3 = normalize_vec3(V3Det[0])
            else:
                V3 = TGenVecWithRandoms(Om[0], Ratio[0], normalize_vec3(V3Det[0]), get_random(), get_random())

            Hit, solutions, jackPDF, pDet, pD2 = TGenPointSpace(S, masses[0], masses[1], V3)
            if not Hit:
                return [], []

            for l in range(solutions):
                vecP = [pDet[l].copy(), pD2[l].copy()]
                VecVecP.append(vecP)
                vecWi.append(jackPDF[l] * whPS)
        else:
            p1, p2 = TwoBodyWithRandoms(S, masses[0], masses[1], get_random(), get_random())
            vecP = [p1, p2]
            wh = 1.0
            vecWi.append(wh * whPS)
            VecVecP.append(vecP)

        return VecVecP, vecWi

    rrr_list = [get_random() for _ in range(nBody - 2)]
    rrr = np.sort(np.array(rrr_list))

    mV = np.zeros(nBody - 2)
    for i in range(nBody - 2):
        mB = 0.0
        mA = 0.0
        for l in range(i + 1):
            mB += masses[l]
        for l in range(i + 1, nBody):
            mA += masses[l]
        mV[i] = rrr[nBody - 2 - i - 1] * (mS - mall) + mA

    whPS = mV[0] * CalcKMFactor(masses[0]**2 / S[3]**2, mV[0]**2 / S[3]**2)

    for i in range(nBody - 3):
        whPS *= mV[i + 1] * CalcKMFactor(masses[i + 1]**2 / (mV[i]**2), mV[i + 1]**2 / (mV[i]**2))

    if nBody > 2:
        whPS *= CalcKMFactor(masses[nBody - 2]**2 / (mV[nBody - 3]**2), masses[nBody - 1]**2 / (mV[nBody - 3]**2))

    whPS *= (kPi**(nBody - 1) / 2.0 * (mS - mall)**(nBody - 2) / factorial_numba(nBody - 2))

    if nDet == 0:
        vecP: List[np.ndarray] = []
        p1, p2 = TwoBodyWithRandoms(S, masses[0], mV[0], get_random(), get_random())
        vecP.append(p1)
        pV = p2

        for i in range(nBody - 3):
            p1, p2 = TwoBodyWithRandoms(pV, masses[i + 1], mV[i + 1], get_random(), get_random())
            vecP.append(p1)
            pV = p2

        p1, p2 = TwoBodyWithRandoms(pV, masses[nBody - 2], masses[nBody - 1], get_random(), get_random())
        vecP.append(p1)
        vecP.append(p2)

        wh = 1.0
        vecWi.append(wh * whPS)
        VecVecP.append(vecP)
        return VecVecP, vecWi

    root = CreateFirst(None, S, S, 1.0)
    cur = root
    level = 0
    branch: List[int] = []

    while level < (nBody - 1):
        if level < nDet and level < (nBody - 2):
            pV = cur.fPV
            if IsPoint(Ratio[level]):
                V3 = normalize_vec3(V3Det[level])
            else:
                V3 = TGenVecWithRandoms(Om[level], Ratio[level], normalize_vec3(V3Det[level]), get_random(), get_random())

            Hit, solutions, jackPDF, pDet, pD2 = TGenPointSpace(pV, masses[level], mV[level], V3)

            if solutions == 0 and len(branch) == 0:
                level = nBody - 1
                continue
            if solutions == 0 and len(branch) > 0:
                while level > branch[-1]:
                    level -= 1
                    cur = cur.fParent
                cur = cur.fLeft
                level += 1
                branch.pop()
                continue
            if solutions == 1:
                CreateRight(cur, None, pDet[0], pD2[0], jackPDF[0])
                cur = cur.fRight
                level += 1
                continue
            elif solutions == 2:
                branch.append(level)
                CreateLeft(cur, None, pDet[1], pD2[1], jackPDF[1])
                CreateRight(cur, None, pDet[0], pD2[0], jackPDF[0])
                cur = cur.fRight
                level += 1
                continue

        if level < nDet and level == (nBody - 2):
            pV = cur.fPV
            if IsPoint(Ratio[level]):
                V3 = normalize_vec3(V3Det[level])
            else:
                V3 = TGenVecWithRandoms(Om[level], Ratio[level], normalize_vec3(V3Det[level]), get_random(), get_random())

            Hit, solutions, jackPDF, pDet, pD2 = TGenPointSpace(pV, masses[level], masses[level + 1], V3)

            if solutions == 0 and len(branch) == 0:
                level = nBody - 1
                continue
            if solutions == 0 and len(branch) > 0:
                while level > branch[-1]:
                    level -= 1
                    cur = cur.fParent
                cur = cur.fLeft
                level += 1
                branch.pop()
                continue
            if solutions == 1:
                CreateRight(cur, None, pDet[0], pD2[0], jackPDF[0])
                cur = cur.fRight
                CreateRight(cur, None, pD2[0], S, 1.0)
                cur = cur.fRight
                level += 1
                continue
            elif solutions == 2:
                CreateRight(cur, None, pDet[0], pD2[0], jackPDF[0])
                cur = cur.fRight
                CreateRight(cur, None, pD2[0], S, 1.0)
                cur = cur.fParent
                CreateLeft(cur, None, pDet[1], pD2[1], jackPDF[1])
                cur = cur.fLeft
                CreateRight(cur, None, pD2[1], S, 1.0)
                cur = cur.fRight
                level += 1

                if level == (nBody - 1) and len(branch) > 0:
                    while level > branch[-1]:
                        level -= 1
                        cur = cur.fParent
                    cur = cur.fParent
                    cur = cur.fLeft
                    level += 1
                    branch.pop()
                    continue
                continue

        if level >= nDet and level < nBody - 2:
            pV = cur.fPV
            p1, p2 = TwoBodyWithRandoms(pV, masses[level], mV[level], get_random(), get_random())
            CreateRight(cur, None, p1, p2, 1.0)
            cur = cur.fRight
            level += 1
            continue

        if level >= nDet and level == nBody - 2:
            pV = cur.fPV
            p1, p2 = TwoBodyWithRandoms(pV, masses[level], masses[level + 1], get_random(), get_random())
            CreateRight(cur, None, p1, p2, 1.0)
            cur = cur.fRight
            CreateRight(cur, None, p2, S, 1.0)
            cur = cur.fRight
            level += 1

            if level == (nBody - 1) and len(branch) > 0:
                while level > branch[-1]:
                    level -= 1
                    cur = cur.fParent
                cur = cur.fParent
                cur = cur.fLeft
                level += 1
                branch.pop()
                continue
            continue

        if level == (nBody - 1) and len(branch) > 0:
            while level > branch[-1]:
                level -= 1
                cur = cur.fParent
            cur = cur.fParent
            cur = cur.fLeft
            level += 1
            branch.pop()
            continue

    pathsJ: List[List[float]] = []
    pathsP: List[List[np.ndarray]] = []
    vecJ: List[float] = []
    vecP_collect: List[np.ndarray] = []

    CollectPathsWeights(nBody, root, vecJ, pathsJ)
    CollectPaths(nBody, root, vecP_collect, pathsP)

    DeleteTree(root)
    root = None

    if len(pathsP) == 0:
        return [], []

    for i in range(len(pathsJ)):
        vecJ = pathsJ[i]
        vecP = pathsP[i]
        vecP = vecP[1:]
        VecVecP.append(vecP)

        wh = 1.0
        for j in range(1, nDet + 1):
            wh *= vecJ[j]
        vecWi.append(wh * whPS)

    return VecVecP, vecWi



class TFANG:
    """User interface class for the FANG phase space generator."""

    def __init__(self, rng: Optional[np.random.Generator] = None, seed: int = None):
        self._fNBody: int = 0
        self._fS: np.ndarray = np.zeros(4)
        self._fMasses: np.ndarray = np.array([])
        self._detectors: List[Detector] = []
        self._fVecVecP: List[List[np.ndarray]] = []
        self._fVecWi: List[float] = []

        if rng is not None:
            self._fRng = rng
            self._fOwnRng = False
        else:
            if seed is None:
                seed = 0
            self._fRng = np.random.default_rng(seed)
            self._fOwnRng = True

    def SetDecay(self, S: np.ndarray, masses: np.ndarray) -> bool:
        """Set decay configuration."""
        if masses is None or len(masses) < 2:
            warnings.warn(f"TFANG::SetDecay: masses must have at least 2 elements, got {0 if masses is None else len(masses)}")
            return False

        nBody = len(masses)
        totalMass = 0.0
        for i in range(nBody):
            if masses[i] < 0.0:
                warnings.warn(f"TFANG::SetDecay: Negative mass at index {i}: {masses[i]}")
                return False
            totalMass += masses[i]

        if totalMass >= S[3]:
            warnings.warn(f"TFANG::SetDecay: Sum of decay masses ({totalMass}) >= parent mass ({S[3]})")
            return False

        self._fS = S.copy()
        self._fNBody = nBody
        self._fMasses = masses.copy()
        self._fVecVecP = []
        self._fVecWi = []
        return True

    def AddDetector(self, detector: Detector) -> None:
        """Add a detector constraint for the next particle in order."""
        if self._fNBody == 0:
            warnings.warn("TFANG::AddDetector: SetDecay must be called before AddDetector")
        self._detectors.append(detector)

    def ClearDetectors(self) -> None:
        """Remove all detector constraints."""
        self._detectors = []

    def _decompose_detectors(self):
        """Convert detector objects into arrays for the internal engine."""
        if not self._detectors:
            return None, None, []
        Om = np.array([d._omega() for d in self._detectors])
        Ratio = np.array([d._ratio() for d in self._detectors])
        V3Det = [d.direction.copy() for d in self._detectors]
        return Om, Ratio, V3Det

    def SetSeed(self, seed: int) -> None:
        """Set random number generator seed."""
        self._fRng = np.random.default_rng(seed)

    def Generate(self) -> int:
        """Generate a single phase space event."""
        if self._fNBody < 2:
            warnings.warn("TFANG::Generate: SetDecay must be called before Generate")
            return 0

        self._fVecVecP = []
        self._fVecWi = []

        Om, Ratio, V3Det = self._decompose_detectors()

        self._fVecVecP, self._fVecWi = GenFANG(
            self._fNBody, self._fS, self._fMasses,
            Om, Ratio, V3Det, self._fRng
        )

        if len(self._fVecVecP) == 0:
            return 0
        return len(self._fVecVecP)

    def GenerateWithRandoms(self, random_array: np.ndarray) -> int:
        """Generate a single phase space event using pre-generated random numbers."""
        if self._fNBody < 2:
            warnings.warn("TFANG::GenerateWithRandoms: SetDecay must be called first")
            return 0

        self._fVecVecP = []
        self._fVecWi = []

        Om, Ratio, V3Det = self._decompose_detectors()

        self._fVecVecP, self._fVecWi = GenFANGWithRandoms(
            self._fNBody, self._fS, self._fMasses,
            Om, Ratio, V3Det, random_array, self._fRng
        )

        if len(self._fVecVecP) == 0:
            return 0
        return len(self._fVecVecP)

    def GetNRandomsPerEvent(self) -> int:
        """Get the estimated number of random numbers needed per event."""
        nDet = len(self._detectors)
        return calc_n_random_per_event(self._fNBody, nDet)

    def GenerateBatch(self, nEvents: int, seed: int = None) -> Tuple[np.ndarray, np.ndarray]:
        """Generate multiple phase space events.

        Works with or without detectors. When detectors are set,
        the returned number of events may differ from nEvents
        (constrained generation can yield 0 or multiple solutions per call).
        """
        if self._fNBody < 2:
            warnings.warn("TFANG::GenerateBatch: SetDecay must be called before GenerateBatch")
            return np.array([]), np.array([])

        if len(self._detectors) == 0:
            if seed is None:
                seed = self._fRng.integers(0, 2**31)
            S = LorentzVector(self._fS[0], self._fS[1], self._fS[2], self._fS[3])
            return gen_fang_fast(self._fNBody, S, self._fMasses, nEvents, seed)

        Om, Ratio, V3Det = self._decompose_detectors()
        all_momenta = []
        all_weights = []

        for _ in range(nEvents):
            vecVecP, vecWi = GenFANG(
                self._fNBody, self._fS, self._fMasses,
                Om, Ratio, V3Det, self._fRng
            )
            for sol_idx in range(len(vecWi)):
                event = np.stack(vecVecP[sol_idx])
                all_momenta.append(event)
                all_weights.append(vecWi[sol_idx])

        if len(all_momenta) == 0:
            return np.empty((0, self._fNBody, 4)), np.empty(0)

        return np.array(all_momenta), np.array(all_weights)

    def GetNSolutions(self) -> int:
        """Get number of solutions from last generation."""
        return len(self._fVecVecP)

    def GetWeight(self, iSolution: int = 0) -> float:
        """Get weight for a specific solution."""
        if len(self._fVecWi) == 0:
            warnings.warn("TFANG::GetWeight: No event generated yet")
            return 0.0
        if iSolution < 0 or iSolution >= len(self._fVecWi):
            warnings.warn(f"TFANG::GetWeight: Solution index {iSolution} out of range")
            return 0.0
        return self._fVecWi[iSolution]

    def GetDecay(self, iSolution_or_iParticle: int, iParticle: int = None) -> np.ndarray:
        """Get 4-momentum of particle."""
        if iParticle is None:
            iSolution = 0
            iParticle = iSolution_or_iParticle
        else:
            iSolution = iSolution_or_iParticle

        if iSolution < 0 or iSolution >= len(self._fVecVecP):
            warnings.warn(f"TFANG::GetDecay: Solution index {iSolution} out of range")
            return np.zeros(4)

        vecP = self._fVecVecP[iSolution]
        if iParticle < 0 or iParticle >= len(vecP):
            warnings.warn(f"TFANG::GetDecay: Particle index {iParticle} out of range")
            return np.zeros(4)
        return vecP[iParticle].copy()

    def GetDecays(self, iSolution: int) -> List[np.ndarray]:
        """Get all 4-momenta for a solution."""
        if iSolution < 0 or iSolution >= len(self._fVecVecP):
            warnings.warn(f"TFANG::GetDecays: Solution index {iSolution} out of range")
            return []
        return [p.copy() for p in self._fVecVecP[iSolution]]

    def GetPhaseSpace(self, nEvents: int = 1000000) -> Tuple[float, float]:
        """Calculate full (unconstrained) phase space integral."""
        if self._fNBody < 2:
            warnings.warn("TFANG::GetPhaseSpace: SetDecay must be called first")
            return 0.0, 0.0
        if nEvents <= 0:
            warnings.warn(f"TFANG::GetPhaseSpace: nEvents must be positive")
            return 0.0, 0.0

        S = LorentzVector(self._fS[0], self._fS[1], self._fS[2], self._fS[3])
        seed = self._fRng.integers(0, 2**31)
        return calculate_phase_space_fast(self._fNBody, S, self._fMasses, nEvents, seed)

    def GetPartialPhaseSpace(self, nEvents: int = 1000000) -> Tuple[float, float]:
        """Calculate partial (constrained) phase space integral."""
        if self._fNBody < 2:
            warnings.warn("TFANG::GetPartialPhaseSpace: SetDecay must be called first")
            return 0.0, 0.0
        if len(self._detectors) == 0:
            warnings.warn("TFANG::GetPartialPhaseSpace: No detectors set")
            return 0.0, 0.0
        if nEvents <= 0:
            warnings.warn(f"TFANG::GetPartialPhaseSpace: nEvents must be positive")
            return 0.0, 0.0

        totalOmega = 1.0
        for det in self._detectors:
            totalOmega *= det._omega()

        localRng = np.random.default_rng(self._fRng.integers(0, 2**31) + 1)

        sumW = 0.0
        sumW2 = 0.0
        nSuccess = 0

        Om, Ratio, V3Det = self._decompose_detectors()

        for i in range(nEvents):
            vecVecP, vecWi = GenFANG(
                self._fNBody, self._fS, self._fMasses,
                Om, Ratio, V3Det, localRng
            )
            if len(vecVecP) > 0:
                for w in vecWi:
                    sumW += w
                    sumW2 += w * w
                    nSuccess += 1

        if nSuccess == 0:
            warnings.warn("TFANG::GetPartialPhaseSpace: No successful events")
            return 0.0, 0.0

        mean = sumW / nSuccess
        variance = sumW2 / nSuccess - mean * mean
        phaseSpace = totalOmega * mean
        error = totalOmega * np.sqrt(variance / nSuccess)
        return phaseSpace, error

    def IsConstrained(self) -> bool:
        """Check if generator has any detectors set."""
        return len(self._detectors) > 0

    def GetNDetectors(self) -> int:
        """Get number of detectors."""
        return len(self._detectors)

    def GetNBody(self) -> int:
        """Get number of particles in decay."""
        return self._fNBody



if __name__ == "__main__":
    print("pyFANG - Python implementation of TFANG loaded successfully!")
    print("Use TFANG class for phase space event generation")
