from tarfile import SUPPORTED_TYPES

import numpy as np

import glog

import bisect



class KdeGearbox:

    def __init__(self, maxTooth=200):

        self.minTooth = 6

        self.maxTooth = maxTooth

    

    

    def get_validToothPair(self, ratio):

        arToothPair = []

        for toothNumber1 in np.arange(self.minTooth, self.maxTooth+1):

            for toothNumber2 in np.arange(toothNumber1, self.maxTooth+1): #toothNumber2 always greater than toothNumber1

                if toothNumber2/toothNumber1<ratio:

                    arToothPair.append([toothNumber1, toothNumber2, toothNumber2*1.0/toothNumber1])

        return sorted(arToothPair, key=lambda x: x[-1], reverse=False) #升序排列的可能性组合



    def get_gearbox_tooth_numer(self, target_ratio, tolerance=0.001):

        arResult = []

        glog.info(f"start to get all ratios of per stage..")

        arPairs = self.get_validToothPair(target_ratio)

        glog.info(f"total 1 step space cnts:{len(arPairs)}")

        valid_pairs = arPairs

        sorted_pairs = valid_pairs

        sorted_ratios = [p[2] for p in sorted_pairs]  # 提取传动比列表用于二分查找

        

        results = []

        # 遍历每个可能的第一级齿轮对

        for p1 in valid_pairs:

            tooth1, tooth2, ratio1 = p1

            # 计算第二级需要的目标传动比

            target_p2_ratio = target_ratio / ratio1

            

            # 二分查找最接近target_p2_ratio的传动比位置

            idx = bisect.bisect_left(sorted_ratios, target_p2_ratio)

            

            # 检查索引附近的几个候选（避免遗漏最优解）

            for i in [idx-2, idx-1, idx, idx+1, idx+2]:

                if 0 <= i < len(sorted_pairs):

                    p2 = sorted_pairs[i]

                    tooth3, tooth4, ratio2 = p2

                    total_ratio = ratio1 * ratio2

                    error = abs(total_ratio - target_ratio)

                    

                    # 误差在容忍范围内则记录结果

                    if error <= tolerance:

                        results.append({

                            'tooth1': tooth1, 'tooth2': tooth2,

                            'tooth3': tooth3, 'tooth4': tooth4,

                            'total_ratio': round(total_ratio, 4),

                            'error': round(error, 6)

                        })

        

        # 按误差升序排序，相同误差按总传动比接近目标值排序

        results.sort(key=lambda x: (x['error'], abs(x['total_ratio'] - target_ratio)))

        # 转换为列表形式返回（保留原始数据结构习惯）

        return [[r['tooth1'], r['tooth2'], r['tooth3'], r['tooth4'], r['total_ratio'], r['error']] for r in results]

    



kgb = KdeGearbox()

r = kgb.get_gearbox_tooth_numer(20.267)

for i in r:

    print(i)