# 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 os

import skimage
import jpeglib
import numpy as np

from tv_decompress import decompress, coef2img, psnr


def float2uint8(x):
    """
    Convert a floating point image array to uint8 format.

    Parameters:
        x (np.ndarray): Input image array.

    Returns:
        np.ndarray: Image array in uint8 format, values clipped to [0, 255].
    """
    return np.uint8(np.round(np.clip(x, 0, 255)))


def normalize_residual(img):
    """
    Normalize the residual image for visualization.

    Parameters:
        img (np.ndarray): Residual image array.

    Returns:
        np.ndarray: Normalized residual image scaled to [0, 255].
    """
    # Rescale the signed residual to a symmetric display range.
    vmax = max(np.max(img), -np.min(img))
    img = img / vmax * 0.5 + 0.5
    img = img * 255
    return img


def method_print_name(method: str) -> str:
    """
    Map an internal solver name to a user-facing label.

    Parameters:
        method (str): Internal method name.

    Returns:
        str: User-facing label for console output.
    """
    # Keep the printed method names aligned with the IPOL script.
    mapping = {
        "tv": "classical TV",
        "atv": "adapted TV",
        "dirichlet": "Dirichlet integral",
    }
    return mapping[method]


def main_with_jpeg(
    input: str,
    output: str,
    method: str,
    K: int,
    beta: float,
    u_ref=None,
    output_diff=None
):
    """
    Perform JPEG decompression using standard and TV minimization-based methods.

    Parameters:
        input (str): Path to the input JPEG image.
        output (str): Path to save the decompressed image.
        method (str): TV minimization method ('tv', 'atv', or 'dirichlet').
        K (int): Number of iterations for TV minimization.
        beta (float): Regularization parameter for TV minimization.
        u_ref (np.ndarray, optional): Reference image for comparison.
        output_diff (str, optional): Path to save the difference image.

    Raises:
        Exception: If input is not a JPEG image.
    """
    # Validate that this entry point only receives JPEG inputs.
    if not input.endswith((".jpg", ".jpeg")):
        error_msg = (
            f"When no additional JPEG compression is processed on the image, "
            f"the input image is required to be a jpeg image, while `{input}` is found as input. "
            "Decoding failed, please try again."
        )
        raise Exception(error_msg)

    # Load the quantized DCT coefficients and quantization table.
    img_data = jpeglib.read_dct(input)
    dct_coeffs = img_data.Y
    q_table = img_data.qt[0]

    height, width = img_data.height, img_data.width

    # Reconstruct the standard JPEG decoding baseline.
    u_baseline = coef2img(dct_coeffs, q_table)[:height, :width]

    # clip u_baseline to [0, 255]
    u_baseline = np.clip(u_baseline, 0, 255)
    
    # Run the selected variational decompression method.
    u_TV = decompress(dct_coeffs, q_table, K=K, method=method, beta=beta,
                      height=height, width=width)
    
    # clip u_TV to [0, 255]
    u_TV = np.clip(u_TV, 0, 255)

    print("debug: u_TV min/max =", u_TV.min(), u_TV.max())
    skimage.io.imsave(output, u_TV.astype(np.uint8))
    print(f"{method_print_name(method)} decompressiion done, result saved to {output}")
    print("JPEG decoding successful.")

    # Compare against the reference image when it is available.
    if u_ref is not None:
        # Build a residual image for visual inspection.
        diff_TV = u_TV - u_ref
        diff_TV = normalize_residual(diff_TV)
        if output_diff is not None:
            skimage.io.imsave(output_diff, np.clip(diff_TV, 0, 255).astype(np.uint8))
            print(f"Difference image saved to {output_diff}")

        # Report PSNR for the JPEG baseline and the optimized result.
        psnr_baseline = psnr(u_baseline, u_ref)
        psnr_TV = psnr(u_TV, u_ref)
        print(f"PSNR (standard JPEG): {psnr_baseline:.2f}")
        print(f"PSNR ({method_print_name(method)}): {psnr_TV:.2f}")

def main(
    input: str,
    output: str,
    method: str,
    K: int,
    beta: float,
    add_compression: bool,
    Q: int
):
    """
    Main entry point for JPEG decompression and optional compression.

    Parameters:
        input (str): Path to the input image.
        output (str): Path to save the decompressed image.
        method (str): TV minimization method ('tv', 'atv', or 'dirichlet').
        K (int): Number of iterations for TV minimization.
        beta (float): Regularization parameter for TV minimization.
        add_compression (bool): Whether to apply additional JPEG compression.
        Q (int): JPEG compression quality (if add_compression is True).
    """
    # Optionally create a JPEG input from an uncompressed source image.
    if add_compression:
        new_input = input + ".jpg"
        img = skimage.io.imread(input)  # load image
        if len(img.shape) == 3:
            u_ref = skimage.color.rgb2gray(img) * 255  # convert to grayscale
        else:
            u_ref = img
        skimage.io.imsave(new_input, float2uint8(u_ref), quality=Q)
        print("JPEG-compressed image saved to", new_input, f"with quality Q={Q}")

        # Save the residual next to the output image for comparison.
        output_diff = os.path.splitext(output)[0] + "_diff.png"
        main_with_jpeg(new_input, output, method, K, beta, u_ref=u_ref, output_diff=output_diff)
    else:
        # Process an existing JPEG input without a reference image.
        main_with_jpeg(input, output, method, K, beta)


def parse_args():
    """
    Parse command-line arguments.

    Returns:
        argparse.Namespace: Parsed arguments.
    """
    import argparse

    # Define the command-line interface for the decompression script.
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", type=str, required=True,
                        help="filename of the jpeg image in .jpg or .jpeg")
    parser.add_argument("--output", type=str, required=False, default="decompressed.png",
                        help="filename of the decompressed image")
    parser.add_argument("--method", type=str, required=True, choices=["tv", "atv", "dirichlet"],
                        help="method to use for TV minimization")
    parser.add_argument("--K", type=int, required=False, default=5,
                        help="number of iterations to solve the TV-minimization problem")
    parser.add_argument("--beta", type=float, required=False, default=-1,
                        help="regularization parameter for the TV minimization")
    parser.add_argument("--add_compression", type=str, required=False, default="false",
                        help="whether to process additional JPEG compression on the input image")
    parser.add_argument("--Q", type=int, required=False,
                        help="quality of the JPEG compression")
    return parser.parse_args()


if __name__ == "__main__":
    # Parse the CLI options and normalize the boolean compression flag.
    args = parse_args()
    add_compression = args.add_compression.lower() == "true"

    # Run the main decompression workflow with the parsed arguments.
    main(input=args.input,
         output=args.output,
         method=args.method,
         K=args.K,
         beta=args.beta,
         add_compression=add_compression,
         Q=args.Q)

# Usage:
#   python main.py --input asset/traffic.jpg --method tv
# or
#   python main.py --input asset/traffic.png --method tv  --add_compression true --Q 80
