from itertools import groupby

from funcy import compose, second

from image_prediction.stitching.utils import make_coord_getter


class CoordGrouper:
    def __init__(self, axis):
        self.c1_getter = make_coord_getter(f"{other_axis(axis)}1")
        self.c2_getter = make_coord_getter(f"{other_axis(axis)}2")

    def group_pairs_by_lesser_coordinate(self, pairs):
        return group_by_coordinate(pairs, self.c1_getter)

    def group_pairs_by_greater_coordinate(self, pairs):
        return group_by_coordinate(pairs, self.c2_getter)


def other_axis(axis):
    return "y" if axis == "x" else "x"


def group_by_coordinate(pairs, coord_getter):
    pairs = sorted(pairs, key=coord_getter)
    return map(compose(list, second), groupby(pairs, coord_getter))