"""
Triangle, pentagonal, and hexagonal numbers are generated by the following formulae:
    Triangle: T_n = n(n+1)/2
    Pentagonal: P_n = n(3n-1)/2
    Hexagonal: H_n = n(2n-1)

It can be verified that T_285 = P_165 = H_143 = 40755.

Find the next triangle number that is also pentagonal and hexagonal.
"""

import time


def timer(func):
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        print(f"{func.__name__} taken: {end_time - start_time:.6f} seconds")
        return result

    return wrapper


def is_pentagonal(x):
    """检查一个数是否为五边形数"""
    # 解方程 n(3n-1)/2 = x => 3n² - n - 2x = 0
    # 判别式 Δ = 1 + 24x 必须是完全平方数
    delta = 1 + 24 * x
    sqrt_delta = int(delta**0.5)
    if sqrt_delta * sqrt_delta != delta:
        return False
    # 并且 (1 + sqrt(Δ)) ≡ 0 mod 6
    return (1 + sqrt_delta) % 6 == 0


def is_hexagonal(x):
    """检查一个数是否为六边形数"""
    # 解方程 n(2n-1) = x => 2n² - n - x = 0
    # 判别式 Δ = 1 + 8x 必须是完全平方数
    delta = 1 + 8 * x
    sqrt_delta = int(delta**0.5)
    if sqrt_delta * sqrt_delta != delta:
        return False
    # 并且 (1 + sqrt(Δ)) ≡ 0 mod 4
    return (1 + sqrt_delta) % 4 == 0


@timer
def main() -> None:
    n = 286
    while True:
        triangle = n * (n + 1) // 2

        # 直接使用已有的验证函数，避免重复计算
        if is_pentagonal(triangle) and is_hexagonal(triangle):
            # 找到后计算对应的索引
            pent_index = (1 + int((1 + 24 * triangle) ** 0.5)) // 6
            hex_index = (1 + int((1 + 8 * triangle) ** 0.5)) // 4
            print(f"Answer: T({n})=P({pent_index})=H({hex_index})={triangle}")
            break

        n += 1


@timer
def main_quick_math(nth: int = 2) -> None:
    u = 5
    v = 3
    sn = 0
    while True:
        u, v = (2 * u + 3 * v), (u + 2 * v)
        if u % 6 == 5 and v % 4 == 3:
            k = (u + 1) // 6
            j = (v + 1) // 4
            n = 2 * j - 1
            sn += 1
            if sn == nth:
                print(f"Answer: T({n})=P({k})=H({j})={n * (n + 1) // 2}")
                break


if __name__ == "__main__":
    main()
    main_quick_math()