import matplotlib.pyplot as plt
import numpy as np

from utils.detect_peaks import detect_peaks


def filter_side_sku(
        shelves,
        img_shape,
        empty_space_keys=["emptiness_place"],
        shift_multiplier=0.009,
        max_shelf_size=0.375,
        max_side_cut=8,
        max_total_cut_multiplier=6,
        is_stop_if_tag=False,
        mpd_var=0.04,
        min_sku_num=20,
        debug=False,
):
    new_img = np.zeros(img_shape[:-1])

    left_side = int(max_shelf_size * img_shape[1])
    right_side = img_shape[1] - left_side

    num_sku = 0
    num_empty_space = 0
    for shelf in shelves:
        for place in shelf["places"]:
            for sku in place["skus"]:
                num_sku += 1
                if sku["id_dscr"] in empty_space_keys or "empt" in sku["id_dscr"]:
                    num_empty_space += 1
                    continue
                if left_side > int(sku["bbox"][0]) or right_side < int(sku["bbox"][0]):
                    new_img[
                    int(sku["bbox"][1]): int(sku["bbox"][3]),
                    int(sku["bbox"][0]): int(sku["bbox"][2]),
                    ] = 1

    if num_sku < min_sku_num:
        return shelves, 0, img_shape[1]

    if num_sku - num_empty_space < 10:
        if debug:
            print()
        return shelves, 0, img_shape[1]

    x_sum = new_img.sum(axis=0)

    x = -x_sum + new_img.shape[1]

    height = int(x_sum.shape[0] * 0.98)
    # peaks, _ = ss.find_peaks(x, height=height, prominence=1, width=20)
    peaks = detect_peaks(
        x, mph=height, mpd=new_img.shape[1] * mpd_var, edge="both", show=debug
    )


    peaks = peaks[
        np.where(((left_side > peaks) | (right_side + (right_side * 0.1) < peaks)))[0]
    ]

    for left_x, value in enumerate(x_sum):
        if value != 0:
            break

    left_slice_peaks = peaks[peaks < left_side]
    if len(left_slice_peaks) > 0:
        left_x = left_slice_peaks[-1]

    right_x = img_shape[1]
    for right_x, value in enumerate(reversed(x_sum)):
        if value != 0:
            break

    right_x = x_sum.shape[0] - right_x
    right_slice_peaks = peaks[peaks > right_side]
    if len(right_slice_peaks) > 0:
        right_x = right_slice_peaks[0]

    magic_shift_c = int(img_shape[1] * shift_multiplier)

    right_x += magic_shift_c if right_x != img_shape[1] else img_shape[1]
    left_x -= magic_shift_c if left_x != 0 and left_x - magic_shift_c > 0 else 0

    left_x = 0 if left_x > left_side else left_x
    right_x = x_sum.shape[0] if right_x < right_side else right_x

    new_shelves, status = create_new_shelf(
        shelves, left_x, right_x,
        empty_space_keys,
        max_side_cut=max_side_cut,
        max_total_cut_multiplier=max_total_cut_multiplier,
        is_stop_if_tag=is_stop_if_tag, debug=debug
    )
    if not status:
        right_x = img_shape[1]
        left_x = 0

    if debug:
        print ("Peaks: {0} , left_x: {1}, right_x: {2}".format(peaks, left_x, right_x))
        print ("Left side: {0}, right: {1}".format(left_side, right_side))
        plt.figure(figsize=(10, 10))
        plt.imshow(new_img)
        plt.show()
    return new_shelves, left_x, right_x


def create_new_shelf(
        shelves,
        left_x,
        right_x,
        empty_space_keys,
        max_side_cut=8,
        max_total_cut_multiplier=6,
        is_stop_if_tag=False,
        debug=False,
):
    new_shelves = []
    k = 0
    sku_total = 0
    sku_total_replaced_num = 0
    for shelf_id, shelf in enumerate(shelves):
        new_shelves.append({key: shelf[key] for key in shelf.keys() if key != "places"})
        new_shelves[shelf_id]["places"] = []
        sku_replaced_num_left = 0
        sku_replaced_num_right = 0
        place_deleted = 0
        for place_id_, place in enumerate(shelf["places"]):
            place_id = place_id_ - place_deleted
            new_shelves[shelf_id]["places"].append(
                {key: place[key] for key in place.keys() if key != "skus"}
            )
            new_shelves[shelf_id]["places"][place_id]["id"] = place_id
            new_shelves[shelf_id]["places"][place_id]["skus"] = []

            for sku in place["skus"]:

                if sku["id_dscr"] in empty_space_keys or "empt" in sku["id_dscr"]:
                    is_sku_empty_space = True
                else:
                    is_sku_empty_space = False

                if not is_sku_empty_space:
                    sku_total += 1

                if (left_x < int(sku["bbox"][0]) < right_x) and (
                        left_x < int(sku["bbox"][2]) < right_x
                ):
                    new_sku = sku.copy()
                    new_sku["id_number"] = k
                    k += 1

                    if "skus" not in new_shelves[shelf_id]["places"][place_id].keys():
                        new_shelves[shelf_id]["places"][place_id]["skus"] = [new_sku]
                    else:
                        new_shelves[shelf_id]["places"][place_id]["skus"].append(new_sku)

                elif not is_sku_empty_space:
                    # counting left and right deleted skus
                    if left_x > int(sku["bbox"][0]):
                        sku_replaced_num_left += 1
                    else:
                        sku_replaced_num_right += 1

                    sku_total_replaced_num += 1

                if (
                        sku_replaced_num_left > max_side_cut
                        or sku_replaced_num_right > max_side_cut
                ):
                    if debug:
                        print(
                            "Early stopping",
                            "left",
                            sku_replaced_num_left,
                            "right",
                            sku_replaced_num_right,
                        )
                    return shelves, False

                if is_stop_if_tag:
                    for tag in sku["tag"]:
                        if left_x > int(tag["bbox"][0]) or right_x < int(
                                tag["bbox"][0]
                        ):
                            if debug:
                                print("early stopping")
                            return shelves, False
            # delete places if out of bound
            if not (
                    left_x < int(place["bbox"][0]) < right_x
                    and left_x < int(place["bbox"][2]) < right_x
            ):
                del new_shelves[shelf_id]["places"][place_id]
                place_deleted += 1

        if sku_total_replaced_num > max_total_cut_multiplier * len(shelves):
            if debug:
                print(
                    "Early stopping",
                    "total",
                    sku_total_replaced_num,
                    max_total_cut_multiplier * len(shelves),
                )

            return shelves, False

    return new_shelves, True