import cv2
import numpy as np
import os
import json

CONFIG_PATH = "config.json"

def load_config():
    if os.path.exists(CONFIG_PATH):
        with open(CONFIG_PATH) as f:
            return json.load(f)
    return {
        "preprocessing": {
            "clahe_clip": 3.0,
            "clahe_grid": 8,
            "bilateral_d": 9,
            "bilateral_sigma": 75,
            "adaptive_blocksize": 21,
            "adaptive_c": 8
        }
    }


def correct_rotation(img):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) if len(img.shape) == 3 else img
    edges = cv2.Canny(gray, 50, 150)
    lines = cv2.HoughLinesP(edges, 1, np.pi/180, 50, minLineLength=30, maxLineGap=10)
    if lines is None:
        return img
    angles = [np.degrees(np.arctan2(l[0][3]-l[0][1], l[0][2]-l[0][0])) for l in lines]
    if abs(np.median(angles)) > 45:
        img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
    return img


def preprocess_chassis(image_path, save_comparison=False, output_dir="results"):
    config = load_config()
    p = config["preprocessing"]

    img = cv2.imread(image_path)
    if img is None:
        raise ValueError(f"Could not load image: {image_path}")

    original = img.copy()
    img = correct_rotation(img)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    clahe = cv2.createCLAHE(
        clipLimit=p["clahe_clip"],
        tileGridSize=(p["clahe_grid"], p["clahe_grid"])
    )
    enhanced = clahe.apply(gray)

    v0 = cv2.cvtColor(enhanced, cv2.COLOR_GRAY2BGR)

    filtered = cv2.bilateralFilter(enhanced, p["bilateral_d"],
                                   p["bilateral_sigma"], p["bilateral_sigma"])
    v1 = cv2.cvtColor(filtered, cv2.COLOR_GRAY2BGR)

    _, otsu = cv2.threshold(filtered, 0, 255,
                            cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    v2 = cv2.cvtColor(otsu, cv2.COLOR_GRAY2BGR)

    bs = p["adaptive_blocksize"]
    bs = bs if bs % 2 == 1 else bs + 1
    adaptive = cv2.adaptiveThreshold(
        filtered, 255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY,
        blockSize=bs, C=p["adaptive_c"]
    )
    v3 = cv2.cvtColor(adaptive, cv2.COLOR_GRAY2BGR)

    pad = p.get("padding", 20)
    def add_padding(im):
        return cv2.copyMakeBorder(im, pad, pad, pad, pad,
                                  cv2.BORDER_CONSTANT, value=(255, 255, 255))

    variations = [add_padding(v) for v in [v0, v1, v2, v3]]

    if save_comparison:
        os.makedirs(output_dir, exist_ok=True)
        fname = os.path.splitext(os.path.basename(image_path))[0]
        h = 200

        def resize_h(im, height):
            r = height / im.shape[0]
            return cv2.resize(im, (int(im.shape[1] * r), height))

        def add_label(im, label):
            out = im.copy() if len(im.shape) == 3 else cv2.cvtColor(im, cv2.COLOR_GRAY2BGR)
            cv2.putText(out, label, (10, 25), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
            return out

        def to_bgr(p):
            return p if len(p.shape) == 3 else cv2.cvtColor(p, cv2.COLOR_GRAY2BGR)

        panels = [add_label(resize_h(cv2.cvtColor(original, cv2.COLOR_BGR2GRAY), h), "Original")]
        labels = ["CLAHE", "Bilateral", "Otsu", "Adaptive"]
        for i, var in enumerate(variations):
            g = cv2.cvtColor(var, cv2.COLOR_BGR2GRAY)
            panels.append(add_label(resize_h(g, h), labels[i]))

        comparison = np.hstack([to_bgr(p) for p in panels])
        cv2.imwrite(os.path.join(output_dir, f"{fname}_comparison.jpg"), comparison)
        print(f"  Saved comparison -> {output_dir}/{fname}_comparison.jpg")

    return variations


if __name__ == "__main__":
    import sys
    if len(sys.argv) < 2:
        print("Usage: python preprocess.py <image_path>")
    else:
        variations = preprocess_chassis(sys.argv[1], save_comparison=True)
        print(f"Generated {len(variations)} variations — check results/")