import torch
import math


def double_phase_amp_phase(amplitudes, phases, three_pi=True, mean_adjust=True):
    """converts amplitude and phase to double phase coding

    amplitudes:  per-pixel amplitudes of the complex field
    phases:  per-pixel phases of the complex field
    three_pi:  if True, outputs values in a 3pi range, instead of 2pi
    mean_adjust:  if True, centers the phases in the range of interest
    """
    # normalize
    amplitudes = amplitudes / amplitudes.max()

    phases_a = phases - torch.acos(amplitudes)
    phases_b = phases + torch.acos(amplitudes)

    phases_out = phases_a
    phases_out[..., ::2, 1::2] = phases_b[..., ::2, 1::2]
    phases_out[..., 1::2, ::2] = phases_b[..., 1::2, ::2]

    if three_pi:
        max_phase = 3 * math.pi
    else:
        max_phase = 2 * math.pi

    if mean_adjust:
        phases_out -= phases_out.mean()

    return (phases_out + max_phase / 2) % max_phase - max_phase / 2