import cv2
import numpy as np

def calculate_homography(in_pts, out_pts):
    """
    Calculates the homography matrix H such that in_pts = H * out_pts.
    :param in_pts: Source points as a numpy array.
    :param out_pts: Destination points as a numpy array.
    :return: Homography matrix H.
    """
    if isinstance(in_pts, list):
        in_pts = np.array(in_pts)

    if isinstance(out_pts, list):
        out_pts = np.array(out_pts)

    mat_A, mat_b = build_sys_equations(in_pts, out_pts)
    H = np.matmul(np.linalg.pinv(mat_A), mat_b)
    H = np.reshape(np.hstack((H, 1)), (3, 3))
    return H

def build_sys_equations(in_pts, out_pts):
    """
    Builds the system of equations for homography calculation.
    :param in_pts: Array of input points.
    :param out_pts: Array of output points.
    :return: Matrix A and vector b.
    """
    mat_A = np.zeros((np.size(in_pts), 8))
    mat_b = in_pts.ravel()

    i = 0
    for x, y in out_pts:
        # x row
        mat_A[i][0:3] = [x, y, 1]
        mat_A[i][-2:] = [-x * mat_b[i], -y * mat_b[i]]

        # y row
        mat_A[i + 1][-5:] = [x, y, 1, -x * mat_b[i + 1], -y * mat_b[i + 1]]

        i += 2

    return mat_A, mat_b

def fit_image_in_target_space(img_src, img_dst, mask, H, offset=np.array([0, 0, 0])):
    """
    Warps img_src into img_dst using the homography matrix H.
    :param img_src: Source image.
    :param img_dst: Target image.
    :param mask: Mask for the destination region.
    :param H: Homography matrix.
    :param offset: Offset correction array [x_offset, y_offset, 0].
    :return: Transformed image.
    """
    pts = get_pixel_coord(mask)  # Get all pixel coordinates in the mask region.
    pts = pts + offset  # Apply offset correction.
    out_src = np.matmul(H, pts.T)
    out_src = out_src / out_src[-1, :]  # Normalize to homogenous coordinates.

    out_src = out_src[0:2, :].T  # Extract x, y coordinates.
    pts = pts[:, 0:2].astype(np.int64)  # Target points in destination image.

    h, w, _ = img_src.shape
    img_dst = get_pixel_val(img_dst, img_src, pts, out_src, offset)
    return img_dst

def get_pixel_coord(mask):
    """
    Extracts x, y coordinates of white pixels in a binary mask.
    :param mask: Binary mask.
    :return: Homogenous coordinates of mask pixels.
    """
    y, x = np.where(mask)
    pts = np.concatenate((x[:, np.newaxis], y[:, np.newaxis], np.ones((x.size, 1))), axis=1)
    return pts

def get_pixel_val(img_dst, img_src, pts, out_src, offset):
    """
    Performs bilinear interpolation to fetch pixel values.
    :param img_dst: Destination image.
    :param img_src: Source image.
    :param pts: Points in the destination image.
    :param out_src: Corresponding points in the source image.
    :param offset: Offset correction.
    :return: Updated destination image.
    """
    h, w, _ = img_src.shape
    tl = np.floor(out_src[:, ::-1]).astype(np.int64)
    br = np.ceil(out_src[:, ::-1]).astype(np.int64)

    pts = pts - offset[:2]
    valid_mask = ~np.logical_or.reduce(
        (np.any(tl < 0, axis=1), np.any(br < 0, axis=1), tl[:, 0] >= h - 1, tl[:, 1] >= w - 1, br[:, 0] >= h - 1, br[:, 1] >= w - 1)
    )
    pts = pts[valid_mask]
    out_src = out_src[valid_mask]
    tl = tl[valid_mask]
    br = br[valid_mask]

    tr = np.concatenate((tl[:, 0:1], br[:, 1:2]), axis=1)
    bl = np.concatenate((br[:, 0:1], tl[:, 1:2]), axis=1)

    weight = np.zeros((out_src.shape[0], 4))
    weight[:, 0] = np.linalg.norm(tl - out_src[:, ::-1], axis=1)
    weight[:, 1] = np.linalg.norm(tr - out_src[:, ::-1], axis=1)
    weight[:, 2] = np.linalg.norm(bl - out_src[:, ::-1], axis=1)
    weight[:, 3] = np.linalg.norm(br - out_src[:, ::-1], axis=1)

    weight[weight == 0] = 1
    weight = 1 / weight
    weight /= np.sum(weight, axis=1, keepdims=True)

    img_dst[pts[:, 1], pts[:, 0], :] = (
        img_src[tl[:, 0], tl[:, 1], :] * weight[:, 0:1]
        + img_src[tr[:, 0], tr[:, 1], :] * weight[:, 1:2]
        + img_src[bl[:, 0], bl[:, 1], :] * weight[:, 2:3]
        + img_src[br[:, 0], br[:, 1], :] * weight[:, 3:4]
    )
    return img_dst