import os
import cv2
import numpy as np
from flask import Flask, request, jsonify, send_file
from werkzeug.utils import secure_filename

app = Flask(__name__)

UPLOAD_FOLDER = os.path.join(os.getcwd(), "storage", "inputs")
OUTPUT_FOLDER = os.path.join(os.getcwd(), "storage", "outputs")

app.config["UPLOAD_FOLDER"] = UPLOAD_FOLDER
app.config["OUTPUT_FOLDER"] = OUTPUT_FOLDER

os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(OUTPUT_FOLDER, exist_ok=True)


def align_image_to_reference(ref_path: str, target_path: str, output_path: str):
    """Aligns target_path image to match ref_path image geometry and saves to output_path.
    
    Uses multiple fallback strategies for robustness:
    1. SIFT with scale normalization and contrast enhancement
    2. ORB (good for low-texture / small images)
    3. AKAZE 
    4. Phase correlation (frequency domain, good when features fail)
    5. Image moments (center-of-mass + principal axis)
    6. Ultimate fallback: center-crop + resize
    """
    ref_img = cv2.imread(ref_path)
    targ_img = cv2.imread(target_path)

    if ref_img is None or targ_img is None:
        raise ValueError("Could not read input images from storage.")

    h, w = ref_img.shape[:2]

    # Try strategies in order of preference
    aligned = None
    method_used = "unknown"

    # --- Strategy 1: SIFT with scale normalization ---
    try:
        aligned = _align_sift(ref_img, targ_img, w, h)
        if aligned is not None:
            method_used = "sift"
    except Exception:
        pass

    # --- Strategy 2: ORB (better for small/blurry images) ---
    if aligned is None:
        try:
            aligned = _align_orb(ref_img, targ_img, w, h)
            if aligned is not None:
                method_used = "orb"
        except Exception:
            pass

    # --- Strategy 3: AKAZE ---
    if aligned is None:
        try:
            aligned = _align_akaze(ref_img, targ_img, w, h)
            if aligned is not None:
                method_used = "akaze"
        except Exception:
            pass

    # --- Strategy 4: Phase correlation (rotation + translation in frequency domain) ---
    if aligned is None:
        try:
            aligned = _align_phase_correlation(ref_img, targ_img, w, h)
            if aligned is not None:
                method_used = "phase"
        except Exception:
            pass

    # --- Strategy 5: Image moments (centroid + principal axis) ---
    if aligned is None:
        try:
            aligned = _align_moments(ref_img, targ_img, w, h)
            if aligned is not None:
                method_used = "moments"
        except Exception:
            pass

    # --- Strategy 6: Ultimate fallback ---
    if aligned is None:
        aligned = _fallback_resize_center(targ_img, w, h)
        method_used = "fallback"

    cv2.imwrite(output_path, aligned)
    return method_used


def _preprocess(gray):
    """Enhance contrast to improve feature detection on blurry/low-contrast images."""
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    return clahe.apply(gray)


def _normalize_scales(ref_img, targ_img, max_dim=1024):
    """Resize images to similar scales before feature matching."""
    # Scale target to reference scale if they differ too much
    ref_max = max(ref_img.shape[:2])
    targ_max = max(targ_img.shape[:2])
    
    if ref_max / targ_max > 2.0 or targ_max / ref_max > 2.0:
        scale = ref_max / targ_max
        new_h = int(targ_img.shape[0] * scale)
        new_w = int(targ_img.shape[1] * scale)
        targ_img = cv2.resize(targ_img, (new_w, new_h), interpolation=cv2.INTER_LANCZOS4)
    
    # Cap at max_dim for performance
    if ref_max > max_dim:
        s = max_dim / ref_max
        ref_img = cv2.resize(ref_img, None, fx=s, fy=s, interpolation=cv2.INTER_AREA)
    if max(targ_img.shape[:2]) > max_dim:
        s = max_dim / max(targ_img.shape[:2])
        targ_img = cv2.resize(targ_img, None, fx=s, fy=s, interpolation=cv2.INTER_AREA)
    
    return ref_img, targ_img


def _get_affine_matrix(src_pts, dst_pts):
    """Estimate affine matrix with relaxed RANSAC for difficult cases."""
    if len(src_pts) < 3 or len(dst_pts) < 3:
        return None
    matrix, inliers = cv2.estimateAffinePartial2D(
        src_pts, dst_pts,
        method=cv2.RANSAC,
        ransacReprojThreshold=5.0,
        maxIters=5000,
        confidence=0.99
    )
    if matrix is None:
        return None
    if inliers is not None and np.sum(inliers) < 3:
        # Very few inliers - try without RANSAC as last resort
        matrix, _ = cv2.estimateAffinePartial2D(src_pts, dst_pts, method=cv2.LMEDS)
    return matrix


def _align_sift(ref_img, targ_img, w, h):
    ref_gray = cv2.cvtColor(ref_img, cv2.COLOR_BGR2GRAY)
    targ_gray = cv2.cvtColor(targ_img, cv2.COLOR_BGR2GRAY)
    
    # Normalize scales
    ref_norm, targ_norm = _normalize_scales(ref_img, targ_img)
    ref_g = _preprocess(cv2.cvtColor(ref_norm, cv2.COLOR_BGR2GRAY))
    targ_g = _preprocess(cv2.cvtColor(targ_norm, cv2.COLOR_BGR2GRAY))
    
    sift = cv2.SIFT_create(nfeatures=5000)
    kp1, des1 = sift.detectAndCompute(ref_g, None)
    kp2, des2 = sift.detectAndCompute(targ_g, None)
    
    if des1 is None or des2 is None or len(kp1) < 6 or len(kp2) < 6:
        return None
    
    # Use BFMatcher instead of FLANN - more stable across scale differences
    bf = cv2.BFMatcher(cv2.NORM_L2)
    matches = bf.knnMatch(des1, des2, k=2)
    
    good = [m for m, n in matches if m.distance < 0.75 * n.distance]
    if len(good) < 6:
        return None
    
    # Scale keypoints back to original image coordinates
    scale_ref = ref_img.shape[1] / ref_norm.shape[1]
    scale_targ = targ_img.shape[1] / targ_norm.shape[1]
    
    src_pts = np.float32([kp1[m.queryIdx].pt for m in good]).reshape(-1, 1, 2) * scale_ref
    dst_pts = np.float32([kp2[m.trainIdx].pt for m in good]).reshape(-1, 1, 2) * scale_targ
    
    matrix = _get_affine_matrix(dst_pts, src_pts)
    if matrix is None:
        return None
    
    return cv2.warpAffine(targ_img, matrix, (w, h),
                          flags=cv2.INTER_LANCZOS4,
                          borderMode=cv2.BORDER_CONSTANT,
                          borderValue=(0, 0, 0))


def _align_orb(ref_img, targ_img, w, h):
    ref_gray = cv2.cvtColor(ref_img, cv2.COLOR_BGR2GRAY)
    targ_gray = cv2.cvtColor(targ_img, cv2.COLOR_BGR2GRAY)
    ref_gray = _preprocess(ref_gray)
    targ_gray = _preprocess(targ_gray)
    
    orb = cv2.ORB_create(nfeatures=5000, scaleFactor=1.2, nlevels=8)
    kp1, des1 = orb.detectAndCompute(ref_gray, None)
    kp2, des2 = orb.detectAndCompute(targ_gray, None)
    
    if des1 is None or des2 is None or len(kp1) < 6 or len(kp2) < 6:
        return None
    
    bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=False)
    matches = bf.knnMatch(des1, des2, k=2)
    
    good = []
    for pair in matches:
        if len(pair) == 2:
            m, n = pair
            if m.distance < 0.8 * n.distance:
                good.append(m)
    
    if len(good) < 6:
        return None
    
    src_pts = np.float32([kp1[m.queryIdx].pt for m in good]).reshape(-1, 1, 2)
    dst_pts = np.float32([kp2[m.trainIdx].pt for m in good]).reshape(-1, 1, 2)
    
    matrix = _get_affine_matrix(dst_pts, src_pts)
    if matrix is None:
        return None
    
    return cv2.warpAffine(targ_img, matrix, (w, h),
                          flags=cv2.INTER_LANCZOS4,
                          borderMode=cv2.BORDER_CONSTANT,
                          borderValue=(0, 0, 0))


def _align_akaze(ref_img, targ_img, w, h):
    ref_gray = cv2.cvtColor(ref_img, cv2.COLOR_BGR2GRAY)
    targ_gray = cv2.cvtColor(targ_img, cv2.COLOR_BGR2GRAY)
    
    akaze = cv2.AKAZE_create()
    kp1, des1 = akaze.detectAndCompute(ref_gray, None)
    kp2, des2 = akaze.detectAndCompute(targ_gray, None)
    
    if des1 is None or des2 is None or len(kp1) < 6 or len(kp2) < 6:
        return None
    
    bf = cv2.BFMatcher(cv2.NORM_HAMMING)
    matches = bf.knnMatch(des1, des2, k=2)
    
    good = [m for m, n in matches if m.distance < 0.8 * n.distance]
    if len(good) < 6:
        return None
    
    src_pts = np.float32([kp1[m.queryIdx].pt for m in good]).reshape(-1, 1, 2)
    dst_pts = np.float32([kp2[m.trainIdx].pt for m in good]).reshape(-1, 1, 2)
    
    matrix = _get_affine_matrix(dst_pts, src_pts)
    if matrix is None:
        return None
    
    return cv2.warpAffine(targ_img, matrix, (w, h),
                          flags=cv2.INTER_LANCZOS4,
                          borderMode=cv2.BORDER_CONSTANT,
                          borderValue=(0, 0, 0))


def _align_phase_correlation(ref_img, targ_img, w, h):
    """Frequency-domain alignment for translation/rotation."""
    ref_gray = cv2.cvtColor(ref_img, cv2.COLOR_BGR2GRAY).astype(np.float32)
    targ_gray = cv2.cvtColor(targ_img, cv2.COLOR_BGR2GRAY)
    
    # Resize target to reference size
    targ_resized = cv2.resize(targ_gray, (w, h)).astype(np.float32)
    
    # Hanning window to reduce edge artifacts
    window = cv2.createHanningWindow((w, h), cv2.CV_32F)
    
    shift, response = cv2.phaseCorrelate(ref_gray * window, targ_resized * window)
    
    matrix = np.array([[1, 0, shift[0]], [0, 1, shift[1]]], dtype=np.float32)
    return cv2.warpAffine(targ_img, matrix, (w, h),
                          flags=cv2.INTER_LANCZOS4,
                          borderMode=cv2.BORDER_CONSTANT,
                          borderValue=(0, 0, 0))


def _align_moments(ref_img, targ_img, w, h):
    """Align using centroid and principal axis - works even with almost no texture."""
    ref_gray = cv2.cvtColor(ref_img, cv2.COLOR_BGR2GRAY)
    targ_gray = cv2.cvtColor(targ_img, cv2.COLOR_BGR2GRAY)
    
    # Otsu threshold to isolate component from cyan background
    _, ref_thresh = cv2.threshold(ref_gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    _, targ_thresh = cv2.threshold(targ_gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    
    # Clean up noise
    kernel = np.ones((5, 5), np.uint8)
    ref_thresh = cv2.morphologyEx(ref_thresh, cv2.MORPH_CLOSE, kernel)
    targ_thresh = cv2.morphologyEx(targ_thresh, cv2.MORPH_CLOSE, kernel)
    
    ref_m = cv2.moments(ref_thresh)
    targ_m = cv2.moments(targ_thresh)
    
    if ref_m["m00"] == 0 or targ_m["m00"] == 0:
        return None
    
    # Centroids
    rcx, rcy = ref_m["m10"] / ref_m["m00"], ref_m["m01"] / ref_m["m00"]
    tcx, tcy = targ_m["m10"] / targ_m["m00"], targ_m["m01"] / targ_m["m00"]
    
    # Principal axis angles
    def principal_angle(m):
        return 0.5 * np.arctan2(2 * m["mu11"], m["mu20"] - m["mu02"])
    
    r_angle = principal_angle(ref_m)
    t_angle = principal_angle(targ_m)
    rotation = r_angle - t_angle
    
    # Scale from area ratio
    scale = np.sqrt(ref_m["m00"] / targ_m["m00"]) if targ_m["m00"] > 0 else 1.0
    
    cos_r = np.cos(rotation) * scale
    sin_r = np.sin(rotation) * scale
    
    tx = rcx - (cos_r * tcx - sin_r * tcy)
    ty = rcy - (sin_r * tcx + cos_r * tcy)
    
    matrix = np.array([[cos_r, -sin_r, tx],
                       [sin_r,  cos_r, ty]], dtype=np.float32)
    
    return cv2.warpAffine(targ_img, matrix, (w, h),
                          flags=cv2.INTER_LANCZOS4,
                          borderMode=cv2.BORDER_CONSTANT,
                          borderValue=(0, 0, 0))


def _fallback_resize_center(targ_img, w, h):
    """Last resort: center the target in a canvas of reference size."""
    th, tw = targ_img.shape[:2]
    
    # Scale to fit within reference while preserving aspect ratio
    scale = min(w / tw, h / th) * 0.9  # 90% fill
    new_w, new_h = int(tw * scale), int(th * scale)
    resized = cv2.resize(targ_img, (new_w, new_h), interpolation=cv2.INTER_LANCZOS4)
    
    # Create black canvas and center the image
    canvas = np.zeros((h, w, 3), dtype=np.uint8)
    y_off = (h - new_h) // 2
    x_off = (w - new_w) // 2
    canvas[y_off:y_off+new_h, x_off:x_off+new_w] = resized
    return canvas


@app.route("/align", methods=["POST"])
def align_endpoint():
    if "reference" not in request.files or "target" not in request.files:
        return jsonify({"error": "Missing 'reference' or 'target' file in request form-data."}), 400

    ref_file = request.files["reference"]
    target_file = request.files["target"]

    if ref_file.filename == "" or target_file.filename == "":
        return jsonify({"error": "No file selected."}), 400

    ref_name = secure_filename(ref_file.filename)
    target_name = secure_filename(target_file.filename)

    input_ref_path = os.path.join(app.config["UPLOAD_FOLDER"], f"ref_{ref_name}")
    input_target_path = os.path.join(app.config["UPLOAD_FOLDER"], f"target_{target_name}")

    ref_file.save(input_ref_path)
    target_file.save(input_target_path)

    output_aligned_path = os.path.join(app.config["OUTPUT_FOLDER"], f"aligned_{target_name}")

    try:
        method_used = align_image_to_reference(input_ref_path, input_target_path, output_aligned_path)
        return send_file(output_aligned_path, mimetype="image/png")

    except Exception as e:
        return jsonify({"status": "error", "message": str(e)}), 500


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000, debug=True)