import numpy as np
import cv2

def assert_image_format(image, fcn_name: str, arg_name: str, force_alpha: bool = True):
    if not isinstance(image, np.ndarray):
        err_msg = 'The blend_modes function "{fcn_name}" received a value of type "{var_type}" for its argument ' \
                  '"{arg_name}". The function however expects a value of type "np.ndarray" for this argument. Please ' \
                  'supply a variable of type np.ndarray to the "{arg_name}" argument.' \
            .format(fcn_name=fcn_name, arg_name=arg_name, var_type=str(type(image).__name__))
        raise TypeError(err_msg)

    if not image.dtype.kind == 'f':
        err_msg = 'The blend_modes function "{fcn_name}" received a numpy array of dtype (data type) kind ' \
                  '"{var_kind}" for its argument "{arg_name}". The function however expects a numpy array of the ' \
                  'data type kind "f" (floating-point) for this argument. Please supply a numpy array with the data ' \
                  'type kind "f" (floating-point) to the "{arg_name}" argument.' \
            .format(fcn_name=fcn_name, arg_name=arg_name, var_kind=str(image.dtype.kind))
        raise TypeError(err_msg)

    if not len(image.shape) == 3:
        err_msg = 'The blend_modes function "{fcn_name}" received a {n_dim}-dimensional numpy array for its argument ' \
                  '"{arg_name}". The function however expects a 3-dimensional array for this argument in the shape ' \
                  '(height x width x R/G/B/A layers). Please supply a 3-dimensional numpy array with that shape to ' \
                  'the "{arg_name}" argument.' \
            .format(fcn_name=fcn_name, arg_name=arg_name, n_dim=str(len(image.shape)))
        raise TypeError(err_msg)

    if force_alpha and not image.shape[2] == 4:
        err_msg = 'The blend_modes function "{fcn_name}" received a numpy array with {n_layers} layers for its ' \
                  'argument "{arg_name}". The function however expects a 4-layer array representing red, green, ' \
                  'blue, and alpha channel for this argument. Please supply a numpy array that includes all 4 layers ' \
                  'to the "{arg_name}" argument.' \
            .format(fcn_name=fcn_name, arg_name=arg_name, n_layers=str(image.shape[2]))
        raise TypeError(err_msg)


def assert_opacity(opacity, fcn_name: str, arg_name: str = 'opacity'):
    if not isinstance(opacity, float) and not isinstance(opacity, int):
        err_msg = 'The blend_modes function "{fcn_name}" received a variable of type "{var_type}" for its argument ' \
                  '"{arg_name}". The function however expects the value passed to "{arg_name}" to be of type ' \
                  '"float". Please pass a variable of type "float" to the "{arg_name}" argument of function ' \
                  '"{fcn_name}".' \
            .format(fcn_name=fcn_name, arg_name=arg_name, var_type=str(type(opacity).__name__))
        raise TypeError(err_msg)

    if not 0.0 <= opacity <= 1.0:
        err_msg = 'The blend_modes function "{fcn_name}" received the value "{val}" for its argument "{arg_name}". ' \
                  'The function however expects that the value for "{arg_name}" is inside the range 0.0 <= x <= 1.0. ' \
                  'Please pass a variable in that range to the "{arg_name}" argument of function "{fcn_name}".' \
            .format(fcn_name=fcn_name, arg_name=arg_name, val=str(opacity))
        raise ValueError(err_msg)


def _compose_alpha(img_in, img_layer, opacity):
    comp_alpha = np.minimum(img_in[:, :, 3], img_layer[:, :, 3]) * opacity
    new_alpha = img_in[:, :, 3] + (1.0 - img_in[:, :, 3]) * comp_alpha
    np.seterr(divide='ignore', invalid='ignore')
    ratio = comp_alpha / new_alpha
    ratio[ratio == np.nan] = 0.0
    return ratio


def create_hard_light_layover(img_in, img_layer, opacity, disable_type_checks: bool = False):
    if not disable_type_checks:
        _fcn_name = 'hard_light'
        assert_image_format(img_in, _fcn_name, 'img_in')
        assert_image_format(img_layer, _fcn_name, 'img_layer')
        assert_opacity(opacity, _fcn_name)
    img_in_norm = img_in / 255.0
    img_layer_norm = img_layer / 255.0
    ratio = _compose_alpha(img_in_norm, img_layer_norm, opacity)
    comp = np.greater(img_layer_norm[:, :, :3], 0.5) \
           * np.minimum(1.0 - ((1.0 - img_in_norm[:, :, :3])
                               * (1.0 - (img_layer_norm[:, :, :3] - 0.5) * 2.0)), 1.0) \
           + np.logical_not(np.greater(img_layer_norm[:, :, :3], 0.5)) \
           * np.minimum(img_in_norm[:, :, :3] * (img_layer_norm[:, :, :3] * 2.0), 1.0)
    ratio_rs = np.reshape(np.repeat(ratio, 3), [comp.shape[0], comp.shape[1], comp.shape[2]])
    img_out = comp * ratio_rs + img_in_norm[:, :, :3] * (1.0 - ratio_rs)
    img_out = np.nan_to_num(np.dstack((img_out, img_in_norm[:, :, 3])))  # add alpha channel and replace nans
    return img_out * 255.0

def create_soft_light_layover(img_in, img_layer, opacity, disable_type_checks: bool = False):
    if not disable_type_checks:
        _fcn_name = 'soft_light'
        assert_image_format(img_in, _fcn_name, 'img_in')
        assert_image_format(img_layer, _fcn_name, 'img_layer')
        assert_opacity(opacity, _fcn_name)

    img_in_norm = img_in / 255.0
    img_layer_norm = img_layer / 255.0

    ratio = _compose_alpha(img_in_norm, img_layer_norm, opacity)

    # The following code does this:
    #   multiply = img_in_norm[:, :, :3]*img_layer[:, :, :3]
    #   screen = 1.0 - (1.0-img_in_norm[:, :, :3])*(1.0-img_layer[:, :, :3])
    #   comp = (1.0 - img_in_norm[:, :, :3]) * multiply + img_in_norm[:, :, :3] * screen
    #   ratio_rs = np.reshape(np.repeat(ratio,3),comp.shape)
    #   img_out = comp*ratio_rs + img_in_norm[:, :, :3] * (1.0-ratio_rs)

    comp = (1.0 - img_in_norm[:, :, :3]) * img_in_norm[:, :, :3] * img_layer_norm[:, :, :3] \
           + img_in_norm[:, :, :3] * (1.0 - (1.0 - img_in_norm[:, :, :3]) * (1.0 - img_layer_norm[:, :, :3]))

    ratio_rs = np.reshape(np.repeat(ratio, 3), [comp.shape[0], comp.shape[1], comp.shape[2]])
    img_out = comp * ratio_rs + img_in_norm[:, :, :3] * (1.0 - ratio_rs)
    img_out = np.nan_to_num(np.dstack((img_out, img_in_norm[:, :, 3])))  # add alpha channel and replace nans
    return img_out * 255.0

def stitch_images(image1, image2, overlap_width):
    """Stitch two images side by side with overlapping edges."""
    height = min(image1.shape[0], image2.shape[0])
    image1 = cv2.resize(image1, (image1.shape[1], height))
    image2 = cv2.resize(image2, (image2.shape[1], height))

    mask = np.zeros((height, overlap_width), dtype=np.float32)
    mask[:, :overlap_width] = np.linspace(1, 0, overlap_width)
    mask = np.dstack([mask] * 3)

    overlap1 = image1[:, -overlap_width:]
    overlap2 = image2[:, :overlap_width]
    blended_overlap = overlap1 * mask + overlap2 * (1 - mask)
    blended_overlap = blended_overlap.astype(np.uint8)

    stitched_image = np.hstack((image1[:, :-overlap_width], blended_overlap, image2[:, overlap_width:]))
    return stitched_image


def tile_image_to_dimensions(image, target_width, target_height, overlap_width):
    times_width = target_width // (image.shape[1] - overlap_width) + 1
    times_height = target_height // (image.shape[0] - overlap_width) + 1
    row_image = image
    for _ in range(times_width - 1):
        row_image = stitch_images(row_image, image, overlap_width)
    final_image = row_image
    for _ in range(times_height - 1):
        final_image = np.vstack((final_image, row_image))
    final_image = final_image[:target_height, :target_width]
    return final_image


def create_image_with_feather_tile(tile_image, width, height, overlap_width):
    from PIL import Image
    tile_image = Image.open(tile_image)
    tile_image.save("tiled_image_pil_converted.png")
    tile_cv2 = cv2.imread("tiled_image_pil_converted.png")
    tiled_image = tile_image_to_dimensions(tile_cv2, width, height, overlap_width)
    cv2.imwrite('tiled_image.png', tiled_image)


def stitch_images_with_control(image1, image2, overlap_width, direction='horizontal'):
    """Stitch two images side by side or top and bottom with overlapping edges."""
    if direction == 'horizontal':
        height = min(image1.shape[0], image2.shape[0])
        image1 = cv2.resize(image1, (image1.shape[1], height))
        image2 = cv2.resize(image2, (image2.shape[1], height))
        mask = np.zeros((height, overlap_width), dtype=np.float32)
        mask[:, :overlap_width] = np.linspace(1, 0, overlap_width)
        mask = np.dstack([mask] * 3)
        overlap1 = image1[:, -overlap_width:]
        overlap2 = image2[:, :overlap_width]
        blended_overlap = overlap1 * mask + overlap2 * (1 - mask)
        blended_overlap = blended_overlap.astype(np.uint8)
        stitched_image = np.hstack((image1[:, :-overlap_width], blended_overlap, image2[:, overlap_width:]))
    elif direction == 'vertical':
        width = min(image1.shape[1], image2.shape[1])
        image1 = cv2.resize(image1, (width, image1.shape[0]))
        image2 = cv2.resize(image2, (width, image2.shape[0]))
        mask = np.zeros((overlap_width, width), dtype=np.float32)
        mask[:overlap_width, :] = np.linspace(1, 0, overlap_width).reshape(-1, 1)
        mask = np.dstack([mask] * 3)
        overlap1 = image1[-overlap_width:, :]
        overlap2 = image2[:overlap_width, :]
        blended_overlap = overlap1 * mask + overlap2 * (1 - mask)
        blended_overlap = blended_overlap.astype(np.uint8)
        stitched_image = np.vstack((image1[:-overlap_width, :], blended_overlap, image2[overlap_width:, :]))
    else:
        raise ValueError("Direction must be 'horizontal' or 'vertical'")
    return stitched_image


def color_extract(image_path, new_width, new_height):
    from PIL import Image
    format_image = Image.open(image_path)
    format_image.save("pil_image.png")
    image = cv2.imread("pil_image.png")
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    height, width, _ = image_rgb.shape
    center_pixel = image_rgb[height // 2, width // 2]
    color = center_pixel
    new_image = np.full((new_height, new_width, 3), color, dtype=np.uint8)
    output_path = 'color_image.jpg'
    cv2.imwrite(output_path, cv2.cvtColor(new_image, cv2.COLOR_RGB2BGR))
    print(f'Color image saved at {output_path}')


def control_texture(texture_image, direction, overlap, width, height):
    import os
    import cv2
    color_extract(texture_image, width, height)
    create_image_with_feather_tile(texture_image, width, height, overlap)
    img1 = cv2.imread('color_image.jpg')
    img2 = cv2.imread('tiled_image.png')
    control_tile_image = stitch_images_with_control(img1,
                                                    img2, overlap, direction)
    os.remove('tiled_image.png')
    cv2.imwrite('tiled_image_2.png', control_tile_image)