mocks

def radial_alignment_box_mock(
		n_centrals=300,
		n_sat=8,
		boxsize=205.0,
		sigma_sat=2.0,
		alignment_scatter=0.2,
		q_range=(0.3, 0.9),
		seed=42,
		margin=0.0,
		density_shapes=False,
):
	"""Generate a periodic-box mock with radially aligned satellites.

	Parameters
	----------
	n_centrals : int, optional
		Number of central galaxies (positions uniform in the box).
	n_sat : int, optional
		Number of satellites per central.
	boxsize : float, optional
		Periodic box size in Mpc/h.
	sigma_sat : float, optional
		Standard deviation (Mpc/h) of the Gaussian satellite offsets
		around their central.
	alignment_scatter : float, optional
		Standard deviation (radians) of the Gaussian noise angle added to
		the projected radial orientation. 0 gives perfect radial alignment.
	q_range : tuple of float, optional
		Uniform sampling range of the projected axis ratio q = b/a.
	seed : int, optional
		Seed for the random generator; fixes the full catalogue.
	margin : float, optional
		Keep centrals at least this far (Mpc/h) from every box face, so no
		periodically wrapped pair falls within the measured separation range.
		Needed when the same catalogue must also be measured without
		periodicity (e.g. the plane-parallel lightcone consistency check).
		Default 0 (fill the whole box).
	density_shapes : bool, optional
		If True, also give the **density** sample its own shapes, as
		"Axis_Direction_density_sample" and "q_density_sample", which is what a
		shape-shape ('++') correlation between two different catalogues needs.
		The density sample is centrals followed by satellites, so its satellite
		entries reuse the shapes those same objects already carry in the shape
		sample and the centrals get independent random orientations; the II
		signal therefore comes from the satellite-satellite pairs the two
		samples share. Default False, which leaves the catalogue byte-for-byte
		as before -- the random draws for the new arrays happen after every
		existing one.

	Returns
	-------
	dict
		Data dictionary in the MeasureIABox input format with the keys
		"Position" (centrals + satellites, the density sample),
		"Position_shape_sample" (satellites), "Axis_Direction" (projected
		2D unit vectors), "q", "LOS" (2, i.e. the z-axis), plus extras
		"boxsize" and "seed" for bookkeeping.

	"""
	rng = np.random.default_rng(seed)
	n_sats_total = n_centrals * n_sat

	centrals = rng.uniform(margin, boxsize - margin, (n_centrals, 3))
	offsets = rng.normal(0.0, sigma_sat, (n_sats_total, 3))
	satellites = (np.repeat(centrals, n_sat, axis=0) + offsets) % boxsize

	# projected (LOS = z) radial direction from central to satellite;
	# use the unwrapped offsets so periodic wrapping cannot flip directions
	radial = offsets[:, :2]
	norm = np.sqrt(np.sum(radial ** 2, axis=1))
	norm[norm < 1e-12] = 1e-12
	direction = radial / norm[:, None]

	# add alignment noise by rotating each direction by a Gaussian angle
	theta = rng.normal(0.0, alignment_scatter, n_sats_total)
	cos_t, sin_t = np.cos(theta), np.sin(theta)
	direction = np.column_stack([
		cos_t * direction[:, 0] - sin_t * direction[:, 1],
		sin_t * direction[:, 0] + cos_t * direction[:, 1],
	])

	q = rng.uniform(q_range[0], q_range[1], n_sats_total)
	positions = np.vstack([centrals, satellites])

	mock = {
		"Position": positions,
		"Position_shape_sample": satellites,
		"Axis_Direction": direction,
		"q": q,
		"LOS": 2,
		"boxsize": boxsize,
		"seed": seed,
	}

	if density_shapes:
		# Drawn last, so switching this on cannot perturb anything above it.
		# The density sample is [centrals, satellites]: the satellites are the
		# same objects as the shape sample and reuse their shapes, which is what
		# makes the cross correlation carry a real II signal; the centrals get
		# independent orientations and axis ratios.
		theta_c = rng.uniform(0.0, 2.0 * np.pi, n_centrals)
		direction_c = np.column_stack([np.cos(theta_c), np.sin(theta_c)])
		q_c = rng.uniform(q_range[0], q_range[1], n_centrals)
		mock["Axis_Direction_density_sample"] = np.vstack([direction_c, direction])
		mock["q_density_sample"] = np.concatenate([q_c, q])

	return mock

def halotools_inputs(mock, ellipticity="distortion"):
	"""Convert a radial-alignment mock to halotools ia_correlations inputs.

	Parameters
	----------
	mock : dict
		Output of radial_alignment_box_mock.
	ellipticity : str, optional
		'distortion' for e = (1-q^2)/(1+q^2) or 'ellipticity' for
		e = (1-q)/(1+q), matching the measureia option.

	Returns
	-------
	tuple
		(shape_positions, orientations, ellipticities, density_positions,
		period) ready for halotools gi_plus_projected.

	"""
	q = mock["q"]
	if ellipticity == "distortion":
		e = (1 - q ** 2) / (1 + q ** 2)
	elif ellipticity == "ellipticity":
		e = (1 - q) / (1 + q)
	else:
		raise ValueError(f"Unknown ellipticity definition: {ellipticity}")
	return (
		mock["Position_shape_sample"],
		mock["Axis_Direction"],
		e,
		mock["Position"],
		mock["boxsize"],
	)

def responsivity(mock, ellipticity="distortion"):
	"""Responsivity R = 1 - <e^2>/2 of the mock shape sample.

	measureia divides S+D by 2R; halotools does not, so
	w_g+^measureia * 2R = w_g+^halotools.
	"""
	q = mock["q"]
	if ellipticity == "distortion":
		e = (1 - q ** 2) / (1 + q ** 2)
	else:
		e = (1 - q) / (1 + q)
	return 1.0 - np.mean(e ** 2) / 2.0

def radial_alignment_lightcone_mock(
		n_centrals=400,
		n_sat=8,
		ra_range=(40.0, 50.0),
		dec_range=(-5.0, 5.0),
		r_range=(2450.0, 2650.0),
		sigma_sat=2.0,
		alignment_scatter=0.2,
		q_range=(0.3, 0.9),
		n_randoms_factor=5,
		seed=4242,
		density_shapes=False,
):
	"""Generate a lightcone mock with radially aligned satellites.

	Centrals are uniform in comoving volume inside a cone section
	(RA/DEC window x comoving-distance shell); satellites get Gaussian 3D
	comoving offsets around their central. Satellite ellipticities e1/e2
	follow the standard survey shear-catalogue convention that measureia
	(and TreeCorr) expect, with the major axis pointing at the satellite's
	own central plus Gaussian angle noise — i.e. radial alignment, giving
	e+ > 0 and w_g+ > 0 in the IA sign convention.

	The position (density) sample is the centrals; the shape sample is the
	satellites, so the two samples are disjoint (no self-pair corrections).
	Distances are comoving; the caller converts r <-> redshift with their
	cosmology (see r_com entries).

	density_shapes : bool, optional
		If True, also give the **density** sample (the centrals) its own shapes,
		as "e1_density_sample"/"e2_density_sample", which is what a shape-shape
		('++') correlation between two different catalogues needs. Each central
		is oriented radially about the cone axis, with the same
		`alignment_scatter`, so the centrals carry a coherent large-scale
		alignment that correlates with the satellites' radial one and gives the
		cross correlation a real, non-degenerate II signal. Default False, which
		leaves the catalogue byte-for-byte as before -- the extra random draws
		happen after every existing one.

	Returns
	-------
	tuple of dict
		(data, randoms_data, info): MeasureIALightcone-format data and
		randoms dictionaries (with Redshift left as comoving distance in
		"r_com*" keys for the caller to convert), and an info dict with
		q, the generator parameters and seed.

	"""
	rng = np.random.default_rng(seed)
	n_sats_total = n_centrals * n_sat

	ra_c, dec_c, r_c = _uniform_cone(rng, n_centrals, ra_range, dec_range, r_range)
	pos_c = _radec_to_cartesian(ra_c, dec_c, r_c)

	offsets = rng.normal(0.0, sigma_sat, (n_sats_total, 3))
	pos_s = np.repeat(pos_c, n_sat, axis=0) + offsets
	ra_s, dec_s, r_s = _cartesian_to_radec(pos_s)

	# local tangent basis at each satellite
	alpha, delta = np.radians(ra_s), np.radians(dec_s)
	east = np.column_stack([-np.sin(alpha), np.cos(alpha), np.zeros(n_sats_total)])
	north = np.column_stack([-np.sin(delta) * np.cos(alpha),
							 -np.sin(delta) * np.sin(alpha),
							 np.cos(delta)])
	n_hat = pos_s / r_s[:, None]

	# projected direction from satellite to central in the tangent plane
	d = np.repeat(pos_c, n_sat, axis=0) - pos_s
	d_perp = d - np.sum(d * n_hat, axis=1, keepdims=True) * n_hat
	phi_axis = np.arctan2(np.sum(d_perp * north, axis=1),
						  np.sum(d_perp * east, axis=1))
	phi_axis += rng.normal(0.0, alignment_scatter, n_sats_total)

	q = rng.uniform(q_range[0], q_range[1], n_sats_total)
	e = (1 - q ** 2) / (1 + q ** 2)
	# survey-convention components: the internal (east, north) frame angle
	# phi_axis maps to the survey frame as -phi_axis (opposite handedness)
	e1 = e * np.cos(2 * phi_axis)
	e2 = -e * np.sin(2 * phi_axis)

	n_rand_d = n_randoms_factor * n_centrals
	n_rand_s = n_randoms_factor * n_sats_total
	ra_rd, dec_rd, r_rd = _uniform_cone(rng, n_rand_d, ra_range, dec_range, r_range)
	ra_rs, dec_rs, r_rs = _uniform_cone(rng, n_rand_s, ra_range, dec_range, r_range)

	data = {
		"RA": ra_c, "DEC": dec_c, "r_com": r_c,
		"RA_shape_sample": ra_s, "DEC_shape_sample": dec_s, "r_com_shape_sample": r_s,
		"e1": e1, "e2": e2,
		"weight": np.ones(n_centrals), "weight_shape_sample": np.ones(n_sats_total),
	}

	if density_shapes:
		# Drawn after everything above, so switching this on perturbs nothing.
		# Centrals are oriented radially about the cone axis: a coherent
		# large-scale alignment that correlates with the satellites' radial one,
		# so the cross correlation carries a real II signal rather than noise.
		alpha_c, delta_c = np.radians(ra_c), np.radians(dec_c)
		east_c = np.column_stack([-np.sin(alpha_c), np.cos(alpha_c), np.zeros(n_centrals)])
		north_c = np.column_stack([-np.sin(delta_c) * np.cos(alpha_c),
								   -np.sin(delta_c) * np.sin(alpha_c),
								   np.cos(delta_c)])
		axis_ra, axis_dec = np.mean(ra_range), np.mean(dec_range)
		axis = _radec_to_cartesian(np.array([axis_ra]), np.array([axis_dec]),
								   np.array([1.0]))[0]
		n_c = pos_c / r_c[:, None]
		# projected direction from each central toward the cone axis
		d = axis[None, :] - np.sum(axis[None, :] * n_c, axis=1)[:, None] * n_c
		phi_c = np.arctan2(np.sum(d * north_c, axis=1), np.sum(d * east_c, axis=1))
		phi_c = phi_c + rng.normal(0.0, alignment_scatter, n_centrals)
		q_c = rng.uniform(q_range[0], q_range[1], n_centrals)
		e_c = (1 - q_c ** 2) / (1 + q_c ** 2)
		# same survey-convention mapping as the satellites above
		data["e1_density_sample"] = e_c * np.cos(2 * phi_c)
		data["e2_density_sample"] = -e_c * np.sin(2 * phi_c)
	randoms_data = {
		"RA": ra_rd, "DEC": dec_rd, "r_com": r_rd,
		"RA_shape_sample": ra_rs, "DEC_shape_sample": dec_rs, "r_com_shape_sample": r_rs,
		"weight": np.ones(n_rand_d), "weight_shape_sample": np.ones(n_rand_s),
	}
	info = {"q": q, "seed": seed, "sigma_sat": sigma_sat,
			"alignment_scatter": alignment_scatter}
	return data, randoms_data, info

def embed_box_mock_on_lightcone(mock, distance=3000.0, alpha0=45.0, delta0=0.0,
								n_randoms_factor=10, seed=777):
	"""Embed a box mock at a large comoving distance on a lightcone.

	The box is placed with its centre at comoving distance `distance` in the
	direction (alpha0, delta0), with the box axes mapped as
	x -> east, y -> north, z -> line of sight (matching the box LOS = 2
	convention). Projected axis directions are converted to survey-convention
	e1/e2 exactly, using each galaxy's own local (east, north) tangent frame,
	so the only differences between MeasureIABox on `mock` and
	MeasureIALightcone on the embedding are the plane-parallel approximation
	itself and the randoms-based (rather than analytic) estimator.

	Randoms are drawn uniformly in the FULL box cube and embedded the same
	way, mirroring the analytic-RR assumption of the box estimator. Use a
	mock generated with a margin >= the maximum measured separation so that
	no periodically wrapped pair contributes to the box measurement.

	Returns
	-------
	tuple of dict
		(data, randoms_data) in MeasureIALightcone format, with comoving
		distances in "r_com*" keys for the caller to convert to redshift.

	"""
	rng = np.random.default_rng(seed)
	L = mock["boxsize"]

	a0, d0 = np.radians(alpha0), np.radians(delta0)
	n0 = np.array([np.cos(d0) * np.cos(a0), np.cos(d0) * np.sin(a0), np.sin(d0)])
	east0 = np.array([-np.sin(a0), np.cos(a0), 0.0])
	north0 = np.array([-np.sin(d0) * np.cos(a0), -np.sin(d0) * np.sin(a0), np.cos(d0)])

	def embed(pos):
		rel = pos - L / 2.0
		return (rel[:, 0, None] * east0 + rel[:, 1, None] * north0
				+ (distance + rel[:, 2])[:, None] * n0)

	pos_d = embed(mock["Position"])
	pos_s = embed(mock["Position_shape_sample"])
	ra_d, dec_d, r_d = _cartesian_to_radec(pos_d)
	ra_s, dec_s, r_s = _cartesian_to_radec(pos_s)

	# exact per-galaxy conversion of the projected axis to survey e1/e2
	axis3d = (mock["Axis_Direction"][:, 0, None] * east0
			  + mock["Axis_Direction"][:, 1, None] * north0)
	alpha, delta = np.radians(ra_s), np.radians(dec_s)
	east = np.column_stack([-np.sin(alpha), np.cos(alpha), np.zeros(len(ra_s))])
	north = np.column_stack([-np.sin(delta) * np.cos(alpha),
							 -np.sin(delta) * np.sin(alpha),
							 np.cos(delta)])
	n_hat = pos_s / r_s[:, None]
	axis_perp = axis3d - np.sum(axis3d * n_hat, axis=1, keepdims=True) * n_hat
	phi_axis = np.arctan2(np.sum(axis_perp * north, axis=1),
						  np.sum(axis_perp * east, axis=1))
	q = mock["q"]
	e = (1 - q ** 2) / (1 + q ** 2)
	e1 = e * np.cos(2 * phi_axis)
	e2 = -e * np.sin(2 * phi_axis)

	n_rand_d = n_randoms_factor * len(ra_d)
	n_rand_s = n_randoms_factor * len(ra_s)
	cube_rd = rng.uniform(0.0, L, (n_rand_d, 3))
	cube_rs = rng.uniform(0.0, L, (n_rand_s, 3))
	pos_rd = embed(cube_rd)
	pos_rs = embed(cube_rs)
	ra_rd, dec_rd, r_rd = _cartesian_to_radec(pos_rd)
	ra_rs, dec_rs, r_rs = _cartesian_to_radec(pos_rs)

	data = {
		"RA": ra_d, "DEC": dec_d, "r_com": r_d,
		"RA_shape_sample": ra_s, "DEC_shape_sample": dec_s, "r_com_shape_sample": r_s,
		"e1": e1, "e2": e2,
		"weight": np.ones(len(ra_d)), "weight_shape_sample": np.ones(len(ra_s)),
	}
	randoms_data = {
		"RA": ra_rd, "DEC": dec_rd, "r_com": r_rd,
		"RA_shape_sample": ra_rs, "DEC_shape_sample": dec_rs, "r_com_shape_sample": r_rs,
		"weight": np.ones(n_rand_d), "weight_shape_sample": np.ones(n_rand_s),
	}
	cube_coords = {"randoms_position": cube_rd, "randoms_shape": cube_rs}
	return data, randoms_data, cube_coords

def subbox_labels(positions, boxsize, L):
	"""Subbox index (0..L^3-1) per position — the same partition the box
	jackknife uses, for supplying identical patches to the lightcone."""
	idx = np.floor(np.asarray(positions) / (boxsize / L)).astype(int)
	idx = np.clip(idx, 0, L - 1)
	return idx[:, 0] * L ** 2 + idx[:, 1] * L + idx[:, 2]