# Copyright (c) 2025 Yanhao Li
# yanhao.li@outlook.com

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.


import numpy as np
import torch
from skimage.util import view_as_blocks
from scipy import fft


def img2coef(img:np.ndarray, q_table:np.ndarray):
    """ Transform an image to blocks of dct coefficients, then quantize the coefficients
        by the quantization table

    Parameters
    ----------
    img : np.ndarray
        image of luminance channel, of shape (H, W)
    q_table : np.ndarray
        quantization table, of shape (8, 8)

    Returns
    -------
    np.ndarray
        dct coefficients, of shape (nr, nc, 8, 8)
    """
    # Shift pixel intensities to the centered JPEG working range.
    img = img - 128
    blocks = view_as_blocks(img, (8, 8))

    # Transform each block to the DCT domain and normalize by the quantization table.
    dct_blocks = fft.dctn(blocks, axes=(-1, -2), norm="ortho", workers=4)  # N, 8, 8
    dct_blocks /= q_table

    return dct_blocks


def coef2img(dct_coeffs:np.ndarray, q_table:np.ndarray) -> np.ndarray: 
    """ De-quantize the dct coefficients, then transform the dct coefficients to an image 

    Parameters
    ----------
    dct_blocks : np.ndarray
        quantized DCT coefficients, of shape (nr, nc, 8, 8)
    q_table : np.ndarray
        quantization table, of shape (8, 8)

    Returns
    -------
    np.ndarray
        An image of luminance channel, of shape (nr*8, nc*8)
    """
    # Recover the block grid shape and de-quantize the DCT coefficients.
    nr, nc = dct_coeffs.shape[:2]
    dct_blocks = dct_coeffs * q_table

    # Transform the blocks back to the image domain and stitch them together.
    restore_blocks = fft.idctn(dct_blocks, axes=(-1, -2), norm="ortho", workers=4)
    img = np.transpose(restore_blocks, (0, 2, 1, 3)).reshape(nr*8, nc*8)

    # Undo the JPEG intensity centering shift.
    img = img + 128

    return img


def _TV(u: torch.Tensor):
    """ Classical l2 total variation

    Parameters
    ----------
    u : torch.Tensor
        A 2D image, suppose the both sides are multiples of 8
    """
    # Build forward finite differences along rows and columns.
    ua = torch.cat((u[1:, :], u[-1:, :]), dim=0) - u
    ub = torch.cat((u[:, 1:], u[:, -1:]), dim=1) - u

    # Accumulate the isotropic total variation energy.
    j_u = torch.sum(torch.sqrt(ua**2 + ub**2 + 1e-8))

    return j_u
    

def _dirichlet(u: torch.Tensor):
    """ Squared l2 total variation

    Parameters
    ----------
    u : torch.Tensor
        A 2D image, suppose the both sides are multiples of 8
    """
    # Build forward finite differences along rows and columns.
    ua = torch.cat((u[1:, :], u[-1:, :]), dim=0) - u
    ub = torch.cat((u[:, 1:], u[:, -1:]), dim=1) - u

    # Accumulate the squared gradient energy.
    j_u = torch.sum(ua**2 + ub**2)

    return j_u


def adapated_alpha() -> torch.Tensor:
    # Return the periodic directional weights used by adapted TV.
    alpha_tab = torch.tensor([5.0, 2.0, 1.0, 1.0, 1.0, 2.0, 5.0, 10.0])
    return alpha_tab


def _aTV(u: torch.Tensor):
    """ Adapted total variation

    Parameters
    ----------
    u : torch.Tensor
        A 2D image, suppose the both sides are multiples of 8
    """
    # Read the image size and the periodic weights for each 8x8 block structure.
    rows, cols = u.shape[:2]

    alpha = adapated_alpha()
    
    # Build the weighted forward differences along rows and columns.
    ua = torch.cat((u[1:, :], u[-1:, :]), dim=0) - u
    alpha_a = alpha.repeat( rows // 8 + 1 )[:rows]
    ua = ua * alpha_a[..., None]

    ub = torch.cat((u[:, 1:], u[:, -1:]), dim=1) - u
    alpha_b = alpha.repeat( cols // 8 + 1)[:cols]
    ub = ub * alpha_b[None, ...]

    # Build the weighted backward differences along rows and columns.
    uc = u - torch.cat((u[:1, :], u[:-1, :]), dim=0)
    alpha_c = alpha.repeat( rows // 8 + 1 )[:rows]
    uc = uc * alpha_c[..., None]

    ud = u - torch.cat((u[:, :1], u[:, :-1]), dim=1)
    alpha_d = alpha.repeat( cols // 8 + 1)[:cols]
    ud = ud * alpha_d[None, ...]

    # Accumulate the adapted total variation energy.
    j_u = torch.sum(torch.sqrt(ua**2 + ub**2 + uc**2 + ud**2 + 1e-8))

    return j_u


def variational_cost(u: torch.Tensor, method:str):
    """ Compute the approximated total variation of a 2D image

    Parameters
    ----------
    u : torch.Tensor
        A 2D image, suppose the both sides are multiples of 8
    method : str
        Options: ["TV", "dirichlet", "aTV"]
        TV: classical total variation
        dirichlet: squared total variation
        aTV: adapted total variation
    """
    # Dispatch to the requested regularization functional.
    if method == "tv":
        return _TV(u)

    elif method == "atv":
        return _aTV(u)

    elif method == "dirichlet":
        return _dirichlet(u)


def step_size(method:str) -> float:
    # Select the default step size for each optimization method.
    if method.lower() == "tv":
        return 0.8
    elif method.lower() == "dirichlet":
        return 0.01
    elif method.lower() == "atv":
        return 0.1
    else:
        raise Exception(f"Unsupported method {method} for TV minimization.")
    

def psnr(img:np.ndarray, ref:np.ndarray):
    # Compute the mean squared error before converting it to PSNR.
    mse = np.mean((img - ref)**2)
    r = 20 * np.log10(255) - 10 * np.log10(mse)
    return r


def decompress(v0_coef:np.ndarray, q_table:np.ndarray, K:int, method:str, beta:float=-1, 
               height:int=None, width:int=None) -> np.ndarray:
    """ Decompress an image, described in Algorithm 1

    Parameters
    ----------
    v0_coef : np.ndarray
        quantized dct coefficients read from the jpeg file, of shape (nr, nc, 8, 8)
    q_table : np.ndarray
        quantization table, of shape (8, 8)
    K : int
        number of iterations to solve the TV-minimization problem
    method : str
        type of cost function, supported types are ["TV", "dirichlet", "aTV"]
    beta : float, optional
        step size of minimization, by default -1, which means using the default step size
    height : int, optional
        height of the decompressed image, by default None
    width : int, optional
        width of the decompressed image, by default None
    Returns
    -------
    np.ndarray
        Decompressed image, of shape (nr*8, nc*8)
    """

    # Fall back to the method-specific step size when beta is not provided.
    if beta <= 0:
        beta = step_size(method)

    # Determine the image size from the DCT coefficient grid if not provided.
    if height is None:
        height = v0_coef.shape[0] * 8
    if width is None:
        width = v0_coef.shape[1] * 8

    # Reconstruct the standard JPEG decoding as the starting point.
    u0 = coef2img(v0_coef, q_table)

    # Initialize the current iterate
    u = np.array(u0)
    j_opt = variational_cost(torch.from_numpy(u)[:height, :width], method=method)
    u_opt = u.copy()[:height, :width]

    # Run the projected gradient iterations described in Algorithm 1.
    for k in range(1, K+1):
        tk = beta / k

        # Rebuild the tensor with gradients enabled for the current iterate
        u = torch.from_numpy(u)
        u.requires_grad_()
        j_u = variational_cost(u[:height, :width], method=method)

        # Compute the gradient of the chosen regularization functional
        grad_u = torch.autograd.grad(outputs=j_u, inputs=u)[0]

        # Take a gradient step in the image domain
        v = u.clone()
        v -= tk * grad_u

        # Project the updated JPEG coefficients back to the quantization intervals
        v = v.detach().numpy()  # (h, w)
        v_coef = img2coef(v, q_table)  # (n, 8, 8)
        v_coef = np.clip(v_coef, v0_coef-0.5, v0_coef+0.5)
        u = coef2img(v_coef, q_table)

        # Update the best solution whenever the cost decreases
        if j_opt > j_u.item():
            j_opt = j_u.item()
            u_opt = u.copy()[:height, :width]

    return u_opt
