# Copyright 2021 CR.Sparse Development Team
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import NamedTuple, List, Dict


import jax.numpy as jnp
from jax import vmap, jit, lax
from jax.numpy.linalg import norm

from cr.sparse import hard_threshold, build_signal_from_indices_and_values


class BIHTState(NamedTuple):
    """Represents the state of the BIHT algorithm
    """
    # The non-zero values
    x_I: jnp.ndarray
    """Non-zero values"""
    I: jnp.ndarray
    """The support for non-zero values"""
    r: jnp.ndarray
    """The residuals"""
    n_mismatched_bits: jnp.ndarray
    # """The number of bits that are mismatched"""
    iterations: int
    """The number of iterations it took to complete"""


def biht(Phi, y, K, tau, max_iters=1000):
    r"""Solves the 1-bit compressive sensing problem :math:`\text{sgn} (\\Phi x) = y` using 
    Binary Iterative Hard Thresholding

    Args:
        Phi (jax.numpy.ndarray): A random dictionary of shape (M, N)
        y (jax.numpy.ndarray): The 1-bit measurements 
        K (int): Sparsity level of solution x (number of non-zero entries)
        tau (float): Step size for the x update step
        max_iters (int): Maximum number of iterations

    Returns:
        (BIHTState): A named tuple containing the solution x and other details


    We assume that :math:`x` is a K-sparse vector.

    We assume that the one-bit measurements are made as follows:

    .. math::

        y = \text{sgn} (\Phi x)

    Thus the vector y contains entries 1 and -1 for the signs of the entries in the
    measurement :math:`\Phi x`.


    The BIHT algorithm proceeds as follows:

    - Start with an estimate :math:`x = 0`
    - Compute the guess :math:`\hat{y} = \text{sgn} (\Phi x)`
    - Measure the residual :math:`r = y - \hat{y}`
    - Count the number of mismatched bits as number of places where r is non-zero.
    - Compute the correlation :math:`h = \Phi^T r`
    - Update x as :math:`x = x + \frac{\tau}{2} h`
    - Hard threshold x to keep only K largest entries
    - Repeat till convergence


    Example:
        >>> import cr.sparse as crs
        >>> import cr.sparse.dict as crdict
        >>> import cr.sparse.data as crdata
        >>> import cr.sparse.cs.cs1bit as cs1bit
        >>> M, N, K = 256, 512, 4
        >>> Phi = crdict.gaussian_mtx(cnb.KEYS[0], M, N, normalize_atoms=False)
        >>> x, omega = crdata.sparse_normal_representations(cnb.KEYS[1], N, K)
        >>> x = x / norm(x)
        >>> y = cs1bit.measure_1bit(Phi, x)
        >>> s0 = crdict.upper_frame_bound(Phi)
        >>> tau = 0.98 * s0
        >>> state = cs1bit.biht_jit(Phi, y, K, tau)
        >>> x_rec = crs.build_signal_from_indices_and_values(N, state.I, state.x_I)
        >>> x_rec = x_rec / norm(x_rec)
    """
    ## Initialize some constants for the algorithm
    M, N = Phi.shape

    def init():
        # Data for the initial approximation [r = y, x = 0]
        I = jnp.arange(0, K)
        x_I = jnp.zeros(K)
        # Assume initial estimate to be zero and compute residual
        # compute the 1 bit output based on current x estimate
        y_int = Phi[:,I] @ x_I
        y_hat = jnp.sign(y_int)
        # difference between actual y and current estimate
        r = y - y_hat
        n_mismatched_bits = jnp.sum(r != 0)
        return BIHTState(x_I=x_I, I=I, r=r, n_mismatched_bits=n_mismatched_bits, 
            iterations=0)

    def body(state):
        # Compute the correlation of the residual with the atoms of Phi
        h = Phi.T @ state.r
        # current approximation
        x = build_signal_from_indices_and_values(N, state.I, state.x_I)
        # update
        x = x + tau/2 * h
        # threshold
        I, x_I = hard_threshold(x, K)
        # Form the subdictionary of corresponding atoms
        Phi_I = Phi[:, I]
        # compute the 1 bit output based on current x estimate
        y_int = Phi_I @ x_I
        y_hat = jnp.sign(y_int)
        # Compute new residual
        # difference between actual y and current estimate
        r = y - y_hat
        # Compute residual norm squared
        n_mismatched_bits = jnp.sum(r != 0)
        # update new state
        return BIHTState(x_I=x_I, I=I, r=r, n_mismatched_bits=n_mismatched_bits,
            iterations=state.iterations+1)

    def cond(state):
        # limit on residual norm 
        a = state.n_mismatched_bits > 0
        # limit on number of iterations
        b = state.iterations < max_iters
        c = jnp.logical_and(a, b)
        # overall condition
        return c

    state = lax.while_loop(cond, body, init())
    return state


biht_jit = jit(biht, static_argnums=(2))

# Copyright 2021 CR.Sparse Development Team
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import NamedTuple, List, Dict


import jax.numpy as jnp
from jax import vmap, jit, lax
from jax.numpy.linalg import norm

from cr.sparse import hard_threshold, build_signal_from_indices_and_values


class BIHTState(NamedTuple):
    """Represents the state of the BIHT algorithm
    """
    # The non-zero values
    x_I: jnp.ndarray
    """Non-zero values"""
    I: jnp.ndarray
    """The support for non-zero values"""
    r: jnp.ndarray
    """The residuals"""
    n_mismatched_bits: jnp.ndarray
    # """The number of bits that are mismatched"""
    iterations: int
    """The number of iterations it took to complete"""


def biht(Phi, y, K, tau, max_iters=1000):
    r"""Solves the 1-bit compressive sensing problem :math:`\text{sgn} (\\Phi x) = y` using 
    Binary Iterative Hard Thresholding

    Args:
        Phi (jax.numpy.ndarray): A random dictionary of shape (M, N)
        y (jax.numpy.ndarray): The 1-bit measurements 
        K (int): Sparsity level of solution x (number of non-zero entries)
        tau (float): Step size for the x update step
        max_iters (int): Maximum number of iterations

    Returns:
        (BIHTState): A named tuple containing the solution x and other details


    We assume that :math:`x` is a K-sparse vector.

    We assume that the one-bit measurements are made as follows:

    .. math::

        y = \text{sgn} (\Phi x)

    Thus the vector y contains entries 1 and -1 for the signs of the entries in the
    measurement :math:`\Phi x`.


    The BIHT algorithm proceeds as follows:

    - Start with an estimate :math:`x = 0`
    - Compute the guess :math:`\hat{y} = \text{sgn} (\Phi x)`
    - Measure the residual :math:`r = y - \hat{y}`
    - Count the number of mismatched bits as number of places where r is non-zero.
    - Compute the correlation :math:`h = \Phi^T r`
    - Update x as :math:`x = x + \frac{\tau}{2} h`
    - Hard threshold x to keep only K largest entries
    - Repeat till convergence


    Example:
        >>> import cr.sparse as crs
        >>> import cr.sparse.dict as crdict
        >>> import cr.sparse.data as crdata
        >>> import cr.sparse.cs.cs1bit as cs1bit
        >>> M, N, K = 256, 512, 4
        >>> Phi = crdict.gaussian_mtx(cnb.KEYS[0], M, N, normalize_atoms=False)
        >>> x, omega = crdata.sparse_normal_representations(cnb.KEYS[1], N, K)
        >>> x = x / norm(x)
        >>> y = cs1bit.measure_1bit(Phi, x)
        >>> s0 = crdict.upper_frame_bound(Phi)
        >>> tau = 0.98 * s0
        >>> state = cs1bit.biht_jit(Phi, y, K, tau)
        >>> x_rec = crs.build_signal_from_indices_and_values(N, state.I, state.x_I)
        >>> x_rec = x_rec / norm(x_rec)
    """
    ## Initialize some constants for the algorithm
    M, N = Phi.shape

    def init():
        # Data for the initial approximation [r = y, x = 0]
        I = jnp.arange(0, K)
        x_I = jnp.zeros(K)
        # Assume initial estimate to be zero and compute residual
        # compute the 1 bit output based on current x estimate
        y_int = Phi[:,I] @ x_I
        y_hat = jnp.sign(y_int)
        # difference between actual y and current estimate
        r = y - y_hat
        n_mismatched_bits = jnp.sum(r != 0)
        return BIHTState(x_I=x_I, I=I, r=r, n_mismatched_bits=n_mismatched_bits, 
            iterations=0)

    def body(state):
        # Compute the correlation of the residual with the atoms of Phi
        h = Phi.T @ state.r
        # current approximation
        x = build_signal_from_indices_and_values(N, state.I, state.x_I)
        # update
        x = x + tau/2 * h
        # threshold
        I, x_I = hard_threshold(x, K)
        # Form the subdictionary of corresponding atoms
        Phi_I = Phi[:, I]
        # compute the 1 bit output based on current x estimate
        y_int = Phi_I @ x_I
        y_hat = jnp.sign(y_int)
        # Compute new residual
        # difference between actual y and current estimate
        r = y - y_hat
        # Compute residual norm squared
        n_mismatched_bits = jnp.sum(r != 0)
        # update new state
        return BIHTState(x_I=x_I, I=I, r=r, n_mismatched_bits=n_mismatched_bits,
            iterations=state.iterations+1)

    def cond(state):
        # limit on residual norm 
        a = state.n_mismatched_bits > 0
        # limit on number of iterations
        b = state.iterations < max_iters
        c = jnp.logical_and(a, b)
        # overall condition
        return c

    state = lax.while_loop(cond, body, init())
    return state


biht_jit = jit(biht, static_argnums=(2))

# Copyright 2021 CR.Sparse Development Team # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # #     https://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. Represents the state of the BIHT algorithm # The non-zero values Non-zero values The support for non-zero values The residuals # """The number of bits that are mismatched""" The number of iterations it took to complete Solves the 1-bit compressive sensing problem :math:`\text{sgn} (\\Phi x) = y` using      Binary Iterative Hard Thresholding      Args:         Phi (jax.numpy.ndarray): A random dictionary of shape (M, N)         y (jax.numpy.ndarray): The 1-bit measurements          K (int): Sparsity level of solution x (number of non-zero entries)         tau (float): Step size for the x update step         max_iters (int): Maximum number of iterations      Returns:         (BIHTState): A named tuple containing the solution x and other details       We assume that :math:`x` is a K-sparse vector.      We assume that the one-bit measurements are made as follows:      .. math::          y = \text{sgn} (\Phi x)      Thus the vector y contains entries 1 and -1 for the signs of the entries in the     measurement :math:`\Phi x`.       The BIHT algorithm proceeds as follows:      - Start with an estimate :math:`x = 0`     - Compute the guess :math:`\hat{y} = \text{sgn} (\Phi x)`     - Measure the residual :math:`r = y - \hat{y}`     - Count the number of mismatched bits as number of places where r is non-zero.     - Compute the correlation :math:`h = \Phi^T r`     - Update x as :math:`x = x + \frac{\tau}{2} h`     - Hard threshold x to keep only K largest entries     - Repeat till convergence       Example:         >>> import cr.sparse as crs         >>> import cr.sparse.dict as crdict         >>> import cr.sparse.data as crdata         >>> import cr.sparse.cs.cs1bit as cs1bit         >>> M, N, K = 256, 512, 4         >>> Phi = crdict.gaussian_mtx(cnb.KEYS[0], M, N, normalize_atoms=False)         >>> x, omega = crdata.sparse_normal_representations(cnb.KEYS[1], N, K)         >>> x = x / norm(x)         >>> y = cs1bit.measure_1bit(Phi, x)         >>> s0 = crdict.upper_frame_bound(Phi)         >>> tau = 0.98 * s0         >>> state = cs1bit.biht_jit(Phi, y, K, tau)         >>> x_rec = crs.build_signal_from_indices_and_values(N, state.I, state.x_I)         >>> x_rec = x_rec / norm(x_rec) ## Initialize some constants for the algorithm # Data for the initial approximation [r = y, x = 0] # Assume initial estimate to be zero and compute residual # compute the 1 bit output based on current x estimate # difference between actual y and current estimate # Compute the correlation of the residual with the atoms of Phi # current approximation # update # threshold # Form the subdictionary of corresponding atoms # compute the 1 bit output based on current x estimate # Compute new residual # difference between actual y and current estimate # Compute residual norm squared # update new state # limit on residual norm # limit on number of iterations # overall condition