"""
Filename: unwrap.py
Author: Andrew Humphreys
Created: 2026-02-06

Description:
    Implement the matlab code from  "Robust 2D Phase Unwrapping..."
    in python.

Dependencies:
    - Python 3.14.2
    - numpy
"""

# ==============================
# Imports
# ==============================

import numpy as np
import matplotlib.pyplot as plt
import scipy

# ==============================
# Constants
# ==============================

pi = np.pi

# ==============================
# Functions
# ==============================

def _wrap(angle):
    """
    Wrap angle(s) to the interval (-period/2, period/2]

    Parameters
    ----------
    angle : array_like
        Angle or angles in radians.
    period : float, optional
        Size of the range over which to wrap the input.

    Returns
    -------
    wrapped : ndarray
        Angles wrapped to (-period/2, priod/2]

    """
    angle = np.asarray(angle)
    return (angle + pi) % (2 * pi) - pi

def _diff(x, n=1, axis=None):
    """
    MATLAB-like diff function.

    Parameters
    ----------
    x : array_like
        Input array.
    n : int, optional
        Order of the difference (default is 1).
    axis : int, optional
        Dimension along which to take the difference.
        If None, uses the first non-singleton axis (MATLAB behavior).

    Returns
    -------
    ndarray
        The n-th forward difference of x.

    Notes
    -----
    Attribution: ChatGPT generated code
    """
    x = np.asarray(x)

    if axis is None:
        # MATLAB: first non-singleton dimension
        axis = next((i for i, s in enumerate(x.shape) if s > 1), 0)

    if n < 0:
        raise ValueError("n must be non-negative")

    y = x
    for _ in range(n):
        y = np.take(y, range(1, y.shape[axis]), axis=axis) \
          - np.take(y, range(0, y.shape[axis] - 1), axis=axis)

    return y

def _dct2(x):
    """
    Python implementation of the MatLab dct2 method.

    Parameters
    ----------
    x : ndarray of shape (M,N)
        input matrix

    Returns
    ndarray of shape (M,N)
        The two-dimensional dct of x

    Notes
    -----
    The matlab definition of dct2 is simplified as follows,
    when only one argument is provided:
        
    >>> function b = dct2(arg1)
    >>>     b = dct(dct(arg1).').';
    >>>     return;

    Where `.'` is the transpose operator. Essentially, all it is doing
    is taking the dct in the first dimension, then in the seond dimension,
    and then returning. Our implementation follows this behavior.

    Using `norm="ortho"` should make the dct function equivalent to the MatLab
    implementation, according to the SciPy documentation.
    """

    # take dct along each axis:
    x = scipy.fft.dct(scipy.fft.dct(x, axis=0, norm='ortho'), axis=1, norm='ortho')
    return x

def _idct2(x):
    """
    Python implementation of the MatLab idct2 method.

    Parameters
    ----------
    x : ndarray of shape (M,N)
        input matrix

    Returns
    ndarray of shape (M,N)
        The two-dimensional inverse dct of x

    Notes
    -----
    The matlab definition of idct2 is simplified as follows,
    when only one argument is provided:
        
    >>> function b = idct2(arg1)
    >>>     b = idct(idct(arg1).').';
    >>>     return;

    Where `.'` is the transpose operator. Essentially, all it is doing
    is taking the dct in the first dimension, then in the second dimension,
    and then returning. Our implementation follows this behavior.

    Using `norm="ortho"` should make the dct function equivalent to the MatLab
    implementation, according to the SciPy documentation.
    """
    # take idct along each axis:

    x = scipy.fft.idct(scipy.fft.idct(x, axis=0, norm='ortho'), axis=1, norm='ortho')
    return x

def unwrap(phase_wrap):
    """
    2D Phase unwrapping algorithm.

    Parameters
    ----------
    phase_wrap : numpy.ndarray of size (M,N)
        Input 2D array of the wrapped phase.

    Returns
    -------
    phase_unwrapped : numpy.ndarray of size (M,N)
        Array of unwrapped phase values.
    N : int
        Number of unwrapping iterations.

    """
    phi1 = _unwrap_tie(phase_wrap)
    phi1 = phi1 + np.mean(phase_wrap) - np.mean(phi1)
    k1 = np.int32(np.round((phi1-phase_wrap) / 2 / pi)) # calculate integer k,
    phase_unwrap = phase_wrap + 2* k1 * pi              # convert to int for while loop
    residue = _wrap(phase_unwrap - phi1)
    phi1 = phi1 + _unwrap_tie(residue)
    phi1 = phi1 + np.mean(phase_wrap) - np.mean(phi1) # adjust piston
    k2 = np.int32(np.round((phi1-phase_wrap) / 2 / pi))
    phase_unwrap = phase_wrap + 2 * k2 * pi
    residue = _wrap(phase_unwrap-phi1)
    N = 0
    while np.sum(np.abs(k1-k2)) > 0:
        k1=k2
        phic = _unwrap_tie(residue)
        phi1 = phi1 + phic
        phi1 = phi1 + np.mean(phase_wrap) - np.mean(phi1) # adjust piston
        k2 = np.int32(np.round((phi1-phase_wrap)/2/pi))
        phase_unwrap = phase_wrap + 2*k2*pi
        residue = _wrap(phase_unwrap - phi1)
        N = N+1

    return (phase_unwrap, N)




def _unwrap_tie(phase_wrap):
    psi = np.exp(1j * phase_wrap)
    xdim, ydim = psi.shape

    # Take the difference along the x axis,
    # and then pad on left/right with zeros:
    edx = np.concatenate(
            [np.zeros((xdim,1)),
             _diff(psi, 1, 1),                               # See notes [1] 
             np.zeros((xdim,1))],
            axis=1
    )

    edy = np.concatenate(
            [np.zeros((1,ydim)),
             _diff(psi, 1, 0),
             np.zeros((1,ydim))],
            axis=0
    )

    # calculate laplacian using the finite difference
    lap = _diff(edx, 1, 1) + _diff(edy, 1, 0)

    # rho is the right hand side of equation (4) in the paper.
    rho = (np.conjugate(psi) * lap).imag                     # see notes [2]

    return solvePoisson(rho)

def solvePoisson(rho):
    dctRho = _dct2(rho)          # shape (N,M)... so dct function only takes real inputs. Why is rho complex?
    N, M = rho.shape
    # matlab actually uses cartesian indexing?!
    I, J = np.meshgrid(np.arange(0,M),np.arange(0,N), indexing='xy')
    
    # We are dividing by zero here, so it is desirable to supress it.
    # First, let's store the current behavior for handling errors:
    err =  np.geterr()['divide'] 
    # supress the divide-by-zero error:
    np.seterr(divide='ignore')

    dctPhi = dctRho / 2 / (np.cos(pi * I/M) + np.cos(pi * J/N) - 2)

    # Restore the original configuration:
    np.seterr(err)

    # handle divide-by-zero values
    dctPhi[0,0] = 0

    return _idct2(dctPhi)

__all__ = ["unwrap"]



# ================================================================================= #
#                                                                                   #
#              888b      88                                                         #
#              8888b     88                ,d                                       #
#              88 `8b    88                88                                       #
#              88  `8b   88   ,adPPYba,  MM88MMM  ,adPPYba,  ,adPPYba,              #  
#              88   `8b  88  a8"     "8a   88    a8P_____88  I8[    ""              #
#              88    `8b 88  8b       d8   88    8PP"""""""   `"Y8ba,               #
#              88     `8888  "8a,   ,a8"   88,   "8b,   ,aa  aa    ]8I              #
#              88      `888   `"YbbdP"'    "Y888  `"Ybbd8"'  `"YbbdP"'              #
#                                                                                   #
# ================================================================================= #
                                                         
r"""
[1] The original code composes the first-order difference function diff(X, 1, d) with the
    function wrap. The first-order difference function is called with a complex input, "psi".
    psi is defined as the exponential of the angle, phi. The first order difference function will
    then return a matrix of complex values itself. The issue is that "wrap" is designed to take
    real-valued inputs that represent angles; it's behavior for complex valued inputs is not formally
    defined. It will work like this:
        
        is the real component of the number in (-pi, pi]?
            => Yes: Then return the original input.
            => No: attempt to wrap it, and return.

    However, when the function attempts to wrap it the modulo function `mod` it fails as it
    is not written to handle complex inputs. We know the authors code will always work as exponential
    is contained within the unit circle; the implication of this is that "wrap" does *nothing*.
    As such the wrapToPi(...) call can (and should) safely be removed.

[2] I have a couple of comments for this. The first is that, exp{-i * phase_wrapped} is equal to the 
    complex conjugate of exp{i * phase_wrapped}. The second is a reminder that the matlab notation ".*"
    represents element-wise multiplication. However, in NumPy element-wise multiplication is the default.



"""
