import os
import sys
import argparse
import cv2
import numpy as np

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(ROOT)

from model_data_processing.Crop_And_Fill import process_complete, smart_crop, remove_noise, masked_hist_equalize_color_bgr
from Load_process.file_processing import Process_File


def make_synth(path):
    h, w = 300, 300
    img = np.zeros((h, w, 3), dtype=np.uint8)
    b = 30
    content = np.random.default_rng(7).integers(40, 180, size=(h - 2*b, w - 2*b, 3), dtype=np.uint8)
    img[b:h-b, b:w-b] = content
    cv2.imwrite(path, img)


def hconcat(a, b):
    h = max(a.shape[0], b.shape[0])
    wa = int(a.shape[1] * h / a.shape[0])
    wb = int(b.shape[1] * h / b.shape[0])
    ra = cv2.resize(a, (wa, h))
    rb = cv2.resize(b, (wb, h))
    return cv2.hconcat([ra, rb])


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--input", type=str, default=None)
    p.add_argument("--output", type=str, default=os.path.join(ROOT, "outputs", "crop_fill_single"))
    p.add_argument("--threshold", type=int, default=30)
    p.add_argument("--min_content_ratio", type=float, default=0.3)
    p.add_argument("--fill_method", type=str, default="inpaint_ns")
    p.add_argument("--no_denoise", action="store_true")
    p.add_argument("--denoise_method", type=str, default="median")
    p.add_argument("--denoise_strength", type=str, default="medium")
    p.add_argument("--no_equalize", action="store_true")
    p.add_argument("--ignore_below", type=int, default=None)
    p.add_argument("--equalize_space", type=str, default="LAB")
    args = p.parse_args()

    pf = Process_File()
    pf.JudgeRoot_MakeDir(args.output)

    src = "C1706.jpg"
    img = cv2.imread(src)
    if img is None:
        raise FileNotFoundError(f"無法讀取: {src}")

    denoised = img if args.no_denoise else remove_noise(img, method=args.denoise_method, strength=args.denoise_strength)
    eq = masked_hist_equalize_color_bgr(
        denoised,
        ignore_below=(args.ignore_below if args.ignore_below is not None else args.threshold),
        space=args.equalize_space,
    )
    cropped, borders, success = smart_crop(eq, args.threshold, args.min_content_ratio, keep_margin=0, show_info=False)
    res = process_complete(
        img,
        threshold=args.threshold,
        min_content_ratio=args.min_content_ratio,
        fill_method=args.fill_method,
        denoise=(not args.no_denoise),
        denoise_method=args.denoise_method,
        denoise_strength=args.denoise_strength,
        equalize=(not args.no_equalize),
        equalize_ignore_below=(args.ignore_below if args.ignore_below is not None else None),
        equalize_space=args.equalize_space,
    )

    base = os.path.splitext(os.path.basename(src))[0]
    out_cropped = f"{base}_cropped.png"
    out_equalized = f"{base}_equalized.png"
    out_proc = f"{base}_processed.png"
    out_cmp = f"{base}_compare.png"
    cv2.imwrite(out_cropped, cropped)
    cv2.imwrite(out_equalized, eq)
    cv2.imwrite(out_proc, res)
    cmp = hconcat(img, res)
    cv2.imwrite(out_cmp, cmp)
    print(out_cropped)
    print(out_equalized)
    print(out_proc)
    print(out_cmp)


if __name__ == "__main__":
    main()