import subprocess
import os

import skimage
import jpeglib
import numpy as np
import magic
from skimage.io import imsave

from tv_decompress import decompress, coef2img, psnr


ROOT = os.path.dirname(os.path.realpath(__file__))


# Convert a floating point image array to uint8 format.
def float2uint8(x):
    return np.uint8(np.round(np.clip(x, 0, 255)))


# Normalize a residual image for visualization.
def normalize_residual(img):
    vmax = max(np.max(img), -np.min(img))
    img = img / vmax * 0.5 + 0.5
    img = img * 255
    return img


# Restore a file extension from its MIME type when needed.
def restore_file_ext(fname):
    metadata = magic.from_file(fname, mime=True)
    ext = metadata.split("/")[-1]
    
    # fix some extension
    if ext == "jpeg":
        ext = "jpg"

    if not fname.endswith(ext):
        fname_root, _ = os.path.splitext(fname)
        new_fname = fname_root + "." + ext
        new_fname = fname + "." + ext
        command = f"mv {fname} {new_fname}"
        subprocess.run(command, shell=True, capture_output=True, text=True)
        return new_fname
    else:
        return fname
    

# Run one decoding branch and save its output artifacts.
def decode_with_method(dct_coeffs: np.ndarray,
                       q_table: np.ndarray,
                       height: int,
                       width: int,
                       output: str,
                       method: str,
                       K: int,
                       beta: float,
                       print_name: str,
                       u_ref=None,
                       output_diff=None):
    # Decode the image with the selected variational method.
    u_decoded = decompress(
        dct_coeffs,
        q_table,
        K=K,
        method=method,
        beta=beta,
        height=height,
        width=width,
    )
    u_decoded = np.clip(u_decoded, 0, 255)
    imsave(output, u_decoded.astype(np.uint8))

    # Compare against the reference image when available.
    if u_ref is not None:
        psnr_decoded = psnr(u_decoded, u_ref)
        print(f"PSNR ({print_name}): {psnr_decoded:.2f}\n")

        # Save the residual image against the reference, matching main.py.
        if output_diff is not None:
            diff_decoded = u_decoded - u_ref
            diff_decoded = normalize_residual(diff_decoded)
            imsave(output_diff, diff_decoded.astype(np.uint8))


# Decode a JPEG image and save all IPOL outputs.
def main_with_jpeg(input: str,
                   output_baseline: str,
                   output_TV: str,
                   output_diff_TV: str,
                   output_sTV: str,
                   output_diff_sTV: str,
                   output_aTV: str,
                   output_diff_aTV: str,
                   K: int,
                   beta: float,
                   u_ref=None):
    # 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 JPEG coefficients and the original JPEG size.
    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

    # Save the standard JPEG decoding baseline.
    u_baseline = coef2img(dct_coeffs, q_table)[:height, :width]
    u_baseline = np.clip(u_baseline, 0, 255)

    imsave(output_baseline, u_baseline.astype(np.uint8))

    # Report the baseline PSNR when a reference image is available.
    if u_ref is not None:
        psnr_baseline = psnr(u_baseline, u_ref)
        print(f"PSNR (standard JPEG): {psnr_baseline:.2f}\n")

    # Save the TV decoding result.
    decode_with_method(
        dct_coeffs,
        q_table,
        height,
        width,
        output_TV,
        method="tv",
        K=K,
        beta=beta,
        print_name="classical TV",
        u_ref=u_ref,
        output_diff=output_diff_TV,
    )

    # Save the adapted TV decoding result.
    decode_with_method(
        dct_coeffs,
        q_table,
        height,
        width,
        output_aTV,
        method="atv",
        K=K,
        beta=beta,
        print_name="adapted TV",
        u_ref=u_ref,
        output_diff=output_diff_aTV,
    )

    # Save the dirichlet-based decoding result under the sTV label.
    decode_with_method(
        dct_coeffs,
        q_table,
        height,
        width,
        output_sTV,
        method="dirichlet",
        K=K,
        beta=beta,
        print_name="Dirichlet integral",
        u_ref=u_ref,
        output_diff=output_diff_sTV,
    )

    print("JPEG decoding successful.")


# Main entry point for optional JPEG compression and IPOL decoding outputs.
def main(input:str,
         output_ref:str,
         output_baseline:str,
         output_TV:str,
         output_diff_TV:str,
         output_sTV:str,
         output_diff_sTV:str,
         output_aTV:str,
         output_diff_aTV:str,
         K:int,
         beta: float,
         add_compression:bool,
         Q:int):

    # Restore the input extension before branching on compression mode.
    input = restore_file_ext(input)

    if add_compression:
        # Build a JPEG input from the reference grayscale image.
        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(output_ref, float2uint8(u_ref))
        skimage.io.imsave("tmp.jpg", float2uint8(u_ref), quality=Q)

        # Reuse the JPEG-only decoding path once the compressed input is ready.
        main_with_jpeg(
            "tmp.jpg",
            output_baseline,
            output_TV,
            output_diff_TV,
            output_sTV,
            output_diff_sTV,
            output_aTV,
            output_diff_aTV,
            K,
            beta,
            u_ref=u_ref,
        )
    else:
        # Decode an existing JPEG input without a reference image.
        main_with_jpeg(
            input,
            output_baseline,
            output_TV,
            output_diff_TV,
            output_sTV,
            output_diff_sTV,
            output_aTV,
            output_diff_aTV,
            K,
            beta,
            u_ref=None,
        )


if __name__ == "__main__":
    # Parse the command line interface for the IPOL workflow.
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", type=str, required=True,
                        help="filename of the jpeg image")
    parser.add_argument("--output_ref", type=str, required=True,
                        help="filename of the non-compressed reference image")
    parser.add_argument("--output_baseline", type=str, required=True,
                        help="filename of the decoded image using standard JPEG decoding")
    parser.add_argument("--output_TV", type=str, required=True,
                        help="filename of the decoded image using TV decoding, with uniform weights for computing TV")
    parser.add_argument("--output_diff_TV", type=str, required=True,
                        help="filename of the difference image between TV decoding and standard decoding, with uniform weights for computing TV")
    parser.add_argument("--output_sTV", type=str, required=True,
                        help="filename of the decoded image using squared TV decoding, with uniform weights for computing TV")
    parser.add_argument("--output_diff_sTV", type=str, required=True,
                        help="filename of the difference image between squared TV decoding and standard decoding, with uniform weights for computing TV")
    parser.add_argument("--output_aTV", type=str, required=True,
                        help="filename of the decoded image using TV decoding, with adapted weights for computing TV")
    parser.add_argument("--output_diff_aTV", type=str, required=True,
                        help="filename of the difference image between TV decoding and standard decoding, with adapted weights for computing TV")
    parser.add_argument("--K", type=int, required=True,
                        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=True,
                        help="whether to process additional JPEG compression on the input image")
    parser.add_argument("--Q", type=int, required=True,
                        help="quality of the JPEG compression")
    
    args = parser.parse_args()

    # Normalize the compression flag to a boolean.
    if args.add_compression == "true":
        add_compression = True
    elif args.add_compression == "false":
        add_compression = False
    else:
        raise ValueError("`--add_compression` must be either `true` or `false`.")

    # Run the full IPOL decoding workflow.
    main(args.input,
         args.output_ref,
         args.output_baseline,
         args.output_TV,
         args.output_diff_TV,
         args.output_sTV,
         args.output_diff_sTV,
         args.output_aTV,
         args.output_diff_aTV,
         args.K,
         args.beta,
         add_compression=add_compression,
         Q=args.Q)

    # command on IPOL: 
    # python $bin/main_ipol.py --input $input_0 --output_ref ref.png --output_baseline baseline.png --output_TV decompress_TV.png --output_diff_TV diff_TV.png --output_sTV decompress_sTV.png --output_diff_sTV diff_sTV.png --output_aTV decompress_aTV.png --output_diff_aTV diff_aTV.png --K $K --add_compression $add_compression --Q $Q
    
    # command on local machine: 
    # python ./main_ipol.py --input traffic.png --output_ref ref.png --output_baseline baseline.png --output_TV decompress_TV.png --output_diff_TV diff_TV.png --output_sTV decompress_sTV.png --output_diff_sTV diff_sTV.png --output_aTV decompress_aTV.png --output_diff_aTV diff_aTV.png --K 15 --add_compression true --Q 90
