from lvec import LVec
import numpy as np

# Create a single Lorentz vector
v = LVec(px=1.0, py=2.0, pz=3.0, E=4.0)

# Access properties
print(f"Mass: {v.mass}")
print(f"pt: {v.pt}")

# Create from pt, eta, phi, mass
v2 = LVec.from_ptepm(pt=5.0, eta=0.0, phi=0.0, m=1.0)

# Vector operations
v3 = v1 + v2
v4 = v1 * 2.0

# Create vector with cache size limit (LRU eviction)
v = LVec(px=1.0, py=2.0, pz=3.0, E=4.0, max_cache_size=100)

# Create vector with TTL for all properties
v_ttl = LVec(px=1.0, py=2.0, pz=3.0, E=4.0, default_ttl=60)  # 60 second TTL

# Set TTL for specific property
v.set_ttl('pt', 10)  # 10 second TTL for transverse momentum

# Get cache statistics
stats = v.cache_stats
print(f"Cache hit ratio: {v.cache_hit_ratio}")

from lvec import LVec, Frame

# Create particles in the lab frame
p1_lab = LVec(px=0.0, py=0.0, pz=20.0, E=25.0)
p2_lab = LVec(px=0.0, py=0.0, pz=-15.0, E=20.0)
total_lab = p1_lab + p2_lab

# Create reference frames
lab_frame = Frame.rest(name="lab")  # Stationary frame
cm_frame = Frame.from_lvec(total_lab, name="cm")  # Center-of-mass frame

# Transform particles to center-of-mass frame
p1_cm = p1_lab.transform_frame(lab_frame, cm_frame)
p2_cm = p2_lab.transform_frame(lab_frame, cm_frame)

# Verify momentum conservation in CM frame
total_cm = p1_cm + p2_cm
print(f"Total momentum in CM: ({total_cm.px:.3f}, {total_cm.py:.3f}, {total_cm.pz:.3f})")
# Should output values close to (0, 0, 0)

# Alternative: directly transform to a frame
p1_cm_alt = p1_lab.to_frame(cm_frame)

from lvec import Vector2D

# Create a 2D vector
vec2d = Vector2D(x=3.0, y=4.0)

# Access properties
print(f"Magnitude: {vec2d.r}")
print(f"Angle (phi): {vec2d.phi}")

# Rotate vector
rotated = vec2d.rotate(angle=np.pi/4)  # 45 degrees rotation

from lvec import Vector3D

# Create a 3D vector
vec3d = Vector3D(x=1.0, y=2.0, z=3.0)

# Access properties
print(f"Magnitude: {vec3d.r}")
print(f"Theta: {vec3d.theta}")
print(f"Phi: {vec3d.phi}")

# Rotate around axis
rotated = vec3d.rotate(theta=np.pi/2, axis=[0, 1, 0])  # 90 degrees around y-axis

# Works with numpy arrays
px = np.array([1.0, 2.0, 3.0])
py = np.array([2.0, 3.0, 4.0])
pz = np.array([3.0, 4.0, 5.0])
E = np.array([4.0, 5.0, 6.0])
vectors = LVec(px, py, pz, E)

# Works with awkward arrays
import awkward as ak
vectors_ak = LVec(ak.Array(px), ak.Array(py), ak.Array(pz), ak.Array(E))