from typing import List, Tuple, Optional
from towhee.operator.base import Operator
from towhee import register


@register(output_schema=['vec'])
class FilterTinySegments(Operator):
    """
    FilterTinySegments
    """

    def __init__(self, filter_s_thresh: float, segment_len_rate: float = 0.8):
        super().__init__()
        self.filter_s_thresh = filter_s_thresh
        self.segment_len_rate = segment_len_rate

    def __call__(self, pred_value: List[Optional[List]], sim_hw: Tuple = None) -> List:
        res_pred_list = []
        for time_info in pred_value:
            filterable = True
            if ((time_info[2] - time_info[0]) + (time_info[3] - time_info[1])) / 2.0 > self.filter_s_thresh:
                filterable = False
            if sim_hw is not None:
                sim_h, sim_w = sim_hw
                if (time_info[2] - time_info[0]) > sim_h * self.segment_len_rate or (
                        time_info[3] - time_info[1]) > sim_w * self.segment_len_rate:
                    filterable = False
            if not filterable:
                res_pred_list.append(time_info)
        return res_pred_list