# Pkg.add(["Images", "TestImages", "ImageFeatures", "ImageDraw", "ImageFiltering", "CoordinateTransformations", "Rotations", "StaticArrays"])
using Images, TestImages, CoordinateTransformations, Rotations, StaticArrays, LinearAlgebra, Statistics

function draw_line!(img, p1, p2, color)
    x1, y1 = round(Int, p1[1]), round(Int, p1[2])
    x2, y2 = round(Int, p2[1]), round(Int, p2[2])
    dx = abs(x2 - x1)
    dy = abs(y2 - y1)
    sx = x1 < x2 ? 1 : -1
    sy = y1 < y2 ? 1 : -1
    err = dx - dy
    x, y = x1, y1
    while true
        if 1 <= x <= size(img,2) && 1 <= y <= size(img,1)
            img[y, x] = color
        end
        if x == x2 && y == y2
            break
        end
        e2 = 2 * err
        if e2 > -dy
            err -= dy
            x += sx
        end
        if e2 < dx
            err += dx
            y += sy
        end
    end
end

function draw_circle!(img, center, radius, color)
    cx, cy = round(Int, center[1]), round(Int, center[2])
    r = radius
    x = 0
    y = r
    d = 3 - 2r
    while x <= y
        for (dx, dy) in [(x, y), (y, x), (-x, y), (-y, x), (x, -y), (y, -x), (-x, -y), (-y, -x)]
            px = cx + dx
            py = cy + dy
            if 1 <= px <= size(img,2) && 1 <= py <= size(img,1)
                img[py, px] = color
            end
        end
        if d < 0
            d += 4x + 6
        else
            d += 4(x - y) + 10
            y -= 1
        end
        x += 1
    end
end

function harris_corners(img; k=0.04, threshold=0.01)
    img_f = Float64.(img)
    img_f = replace(img_f, NaN => 0.0, Inf => 0.0, -Inf => 0.0)
    Ix = imfilter(img_f, Kernel.sobel()[1])
    Iy = imfilter(img_f, Kernel.sobel()[2])
    Ixx = Ix.^2
    Iyy = Iy.^2
    Ixy = Ix .* Iy
    window = ones(3,3) ./ 9
    Sxx = imfilter(Ixx, window)
    Syy = imfilter(Iyy, window)
    Sxy = imfilter(Ixy, window)
    det = Sxx .* Syy - Sxy.^2
    trace = Sxx + Syy
    response = det - k * trace.^2
    max_resp = maximum(response)
    corners = findall(response .> threshold * max_resp)
    local_max = []
    for c in corners
        i, j = c.I
        val = response[i,j]
        r_min = max(1, i-2)
        r_max = min(size(img,1), i+2)
        c_min = max(1, j-2)
        c_max = min(size(img,2), j+2)
        window_vals = response[r_min:r_max, c_min:c_max]
        if all(val .>= window_vals)
            push!(local_max, (j, i))  
        end
    end
    return local_max
end

function filter_valid_keypoints(img, keypoints; margin=10)
    h, w = size(img)
    valid = []
    for (x, y) in keypoints
        if margin < x <= w-margin && margin < y <= h-margin
            if img[round(Int,y), round(Int,x)] != 0
                push!(valid, (x, y))
            end
        end
    end
    return valid
end

function extract_patch_descriptors(img, keypoints; patch_size=7)
    img_f = Float64.(img)
    img_f = replace(img_f, NaN => 0.0, Inf => 0.0, -Inf => 0.0)
    descriptors = []
    offset = div(patch_size, 2)
    valid_indices = []
    
    for (idx, (x,y)) in enumerate(keypoints)
        x_int = clamp(round(Int, x), offset+1, size(img,2)-offset)
        y_int = clamp(round(Int, y), offset+1, size(img,1)-offset)
        patch = img_f[y_int-offset:y_int+offset, x_int-offset:x_int+offset]
        if std(patch) < 1e-6 || count(patch .== 0) / length(patch) > 0.3
            continue
        end
        desc = patch[:]
        desc = (desc .- mean(desc)) / (std(desc) + 1e-8)
        push!(descriptors, desc)
        push!(valid_indices, idx)
    end
    
    if length(descriptors) == 0
        return zeros(0, 0), Int[]
    end
    
    return hcat(descriptors...), valid_indices
end

function match_descriptors_euclidean(desc1, desc2, idx1, idx2; ratio_thresh=0.7)
    matches = []
    n1 = size(desc1, 2)
    n2 = size(desc2, 2)
    if n1 == 0 || n2 == 0
        return matches, Int[], Int[]
    end
    for i in 1:n1
        dists = [norm(desc1[:,i] - desc2[:,j]) for j in 1:n2]
        best_idx = argmin(dists)
        sorted_dists = sort(dists)
        if length(sorted_dists) >= 2 && sorted_dists[1] < ratio_thresh * sorted_dists[2]
            push!(matches, (idx1[i], idx2[best_idx]))
        end
    end
    return matches
end

function ransac_homography(kp1, kp2, matches; iter=500, thresh=3.0)
    best_inliers = []
    best_H = nothing
    n = length(matches)
    if n < 4
        return matches, nothing
    end
    for _ in 1:iter
        sample_idx = rand(1:n, 4)
        sample = matches[sample_idx]
        A = zeros(8, 8)
        b = zeros(8)
        for (row, (idx1, idx2)) in enumerate(sample)
            x1, y1 = kp1[idx1]
            x2, y2 = kp2[idx2]
            A[2*row-1, :] = [x1, y1, 1, 0, 0, 0, -x2*x1, -x2*y1]
            b[2*row-1] = x2
            A[2*row, :] = [0, 0, 0, x1, y1, 1, -y2*x1, -y2*y1]
            b[2*row] = y2
        end
        try
            h = A \ b
            H = [h[1] h[2] h[3]; h[4] h[5] h[6]; h[7] h[8] 1.0]
            inliers = []
            for (idx1, idx2) in matches
                x1, y1 = kp1[idx1]
                x2, y2 = kp2[idx2]
                p = [x1, y1, 1.0]
                pred = H * p
                pred = pred[1:2] ./ pred[3]
                if norm(pred - [x2, y2]) < thresh
                    push!(inliers, (idx1, idx2))
                end
            end
            if length(inliers) > length(best_inliers)
                best_inliers = inliers
                best_H = H
            end
        catch
            continue
        end
    end
    return best_inliers, best_H
end

println("\N[1]上传图片。..")
img_original = testimage("lena_gray_512")
img_original = Float64.(img_original)  
display(Gray.(img_original))

println("\n[2]变换（旋转+缩放）。..")
angle = π/15
scale = 1.0
center = SVector(size(img_original, 2)/2, size(img_original, 1)/2)
tfm = Translation(center) ∘ LinearMap(RotMatrix(angle)) ∘ LinearMap(SDiagonal(scale, scale)) ∘ Translation(-center)

img_transformed = warp(img_original, tfm, axes(img_original))
img_transformed = replace(img_transformed, NaN => 0.0, Inf => 0.0, -Inf => 0.0)
img_transformed_gray = Gray.(clamp01.(img_transformed))
display(img_transformed_gray)

println("\n[3]关键点检测（Harris角点）。..")
keypoints_original = harris_corners(img_original)
keypoints_transformed = harris_corners(img_transformed)
keypoints_transformed = filter_valid_keypoints(img_transformed, keypoints_transformed)

println("   原文上的要点：$(length(keypoints_original))")
println("   转换（有效）上的关键点：$（长度（keypoints_transformed））")

img_kp_original = RGB.(img_original)
for (x,y) in keypoints_original
    draw_circle!(img_kp_original, (x, y), 3, RGB(1,0,0))
end
display(img_kp_original)

img_kp_transformed = RGB.(img_transformed)
for (x,y) in keypoints_transformed
    draw_circle!(img_kp_transformed, (x, y), 3, RGB(0,0,1))
end
display(img_kp_transformed)

println("\n[4]提取描述符（基于补丁）。..")
desc_orig, valid_idx_orig = extract_patch_descriptors(img_original, keypoints_original)
desc_trans, valid_idx_trans = extract_patch_descriptors(img_transformed, keypoints_transformed)

println("   原文中的有效描述符：$(length(valid_idx_orig))")
println("   转换后的有效描述符：$(length(valid_idx_trans))")

println("\n[5]用比率检验匹配描述符。..")
matches = match_descriptors_euclidean(desc_orig, desc_trans, valid_idx_orig, valid_idx_trans, ratio_thresh=0.7)
println("   比率测试后的匹配:$(长度(匹配))")

h, w = size(img_original)
combined = fill(Gray(0), h, 2w)
combined[:, 1:w] = Gray.(img_original)
combined[:, w+1:2w] = img_transformed_gray
combined_rgb = RGB.(combined)

for (i1, i2) in matches
    x1, y1 = keypoints_original[i1]
    x2, y2 = keypoints_transformed[i2]
    p1 = (x1, y1)
    p2 = (x2 + w, y2)  
    draw_line!(combined_rgb, p1, p2, RGB(0,1,0))
end
display(combined_rgb)

println("\n[6] RANSAC...")
inliers, H = ransac_homography(keypoints_original, keypoints_transformed, matches, iter=1000, thresh=3.0)
println("   RANSAC后的匹配项(inliers):$(length(inliers))")

combined_ransac = RGB.(combined)
for (i1, i2) in inliers
    x1, y1 = keypoints_original[i1]
    x2, y2 = keypoints_transformed[i2]
    p1 = (x1, y1)
    p2 = (x2 + w, y2)
    draw_line!(combined_ransac, p1, p2, RGB(0,1,0))
    draw_circle!(combined_ransac, p1, 5, RGB(1,0,0))
    draw_circle!(combined_ransac, p2, 5, RGB(1,0,0))
end
display(combined_ransac)