#!/usr/bin/env python3

import math, argparse

# Use the scipy library for numerical solutions.
# https://docs.scipy.org/doc/scipy/reference/optimize.html
import scipy.optimize
import numpy as np

#================================================================

# Each belt drive system is described by a set of properties.  This allows the
# same code to work with different systems.

GT2 = { 'units'             : 'mm',
        'pitch'             : 2.00,
        'tooth_height'      : 0.76,
        'pitch_line_offset' : 0.25,
        'belt_thickness'    : 1.52,
       }

MXL = { 'units'             : 'inch',
        'pitch'             : 0.080,
        'tooth_height'      : 0.020,
        'pitch_line_offset' : 0.010,
        'belt_thickness'    : 0.045,
       }

XL =  { 'units'             : 'inch',
        'pitch'             : 0.200,
        'tooth_height'      : 0.050,
        'pitch_line_offset' : 0.010,
        'belt_thickness'    : 0.090,
       }

L  =  { 'units'             : 'inch',
        'pitch'             : 0.375,
        'tooth_height'      : 0.075,
        'pitch_line_offset' : 0.015,
        'belt_thickness'    : 0.140,
       }

#================================================================ 
# Timing pulley center distance formulas.
# See section 22 in https://www.sdp-si.com/PDFS/Technical-Section-Timing.pdf

# Given:
#  C    center distance
#  L    belt length
#  R1   larger pulley pitch radius
#  R2   smaller pulley pitch radius
#  phi  One half angle of wrap on smaller pulley (radians)

#  C * cos(phi) = R1 - R2
#  2 * C * sin(phi) = L - pi*(R1+R2) - (pi - 2*phi)*(R1-R2)

# Rewriting:
#  phi = acos((R1 - R2) / C)
#  L   = 2 * C * sin(phi) + pi * (R1 + R2) + (pi - 2 * phi) * (R1 - R2)

def timing_belt_length(N1, N2, C, system=GT2):
    """Calculate the belt length in teeth for a pair of pulleys with N1 and N2 teeth
    separated by center distance C.  The result may be fractional and thus
    unrealizable; endless belts have an integer number of teeth chosen from a
    specific set of available sizes.
    """

    # Make sure N1 >= N2:
    if N2 > N1:
        N1, N2 = N2, N1

    # Look up belt properties.
    pitch = system['pitch']
    
    # Calculate the radius of each pulley in the system units.  The circumference
    # at the pitch diameter is the number of teeth multiplied by the belt tooth pitch.
    R1 = N1 * pitch / (2 * math.pi)
    R2 = N2 * pitch / (2 * math.pi)

    # Check the pulley separation:
    if R1 + R2 >= C:
        print("Warning: pulleys collide, solution not feasible.")
        
    # The essential calculation does not depend on units:
    phi = math.acos((R1 - R2) / C)
    L   = 2 * C * math.sin(phi) + math.pi * (R1 + R2) + (math.pi - 2 * phi) * (R1 - R2)

    # Convert the length to a tooth count (possibly non-integer).
    belt_teeth = L / pitch
    return belt_teeth

#================================================================
def timing_belt_center_distance(N1, N2, T, system=GT2):
    """Calculate the center to center distance for a pair of pulleys with N1 and N2
    teeth driving a belt with T teeth.
    """

    # Look up belt properties.
    pitch = system['pitch']
    
    # Use the scipy fmin algorithm to calculate an inverse solution using the
    # length function.  The initial guess is based on the pulley sizes.
    x0 = np.array(((N1+N2) * pitch))
    
    result = scipy.optimize.fmin(lambda x: abs(timing_belt_length(N1, N2, x[0]) - T), x0, disp=False)
    return result[0]

#================================================================
def timing_pulley_dimensions(N, system=GT2):
    """Calculate properties for a timing pulley of a given size and system.

    :param N: integer number of teeth
    :param system: dictionary of belt system properties.
    :return: dictionary of pulley properties
    """

    # Look up belt properties.
    pitch = system['pitch']
    tooth_height = system['tooth_height']

    # Calculate pulley properties.

    # The pitch diameter is the effective diameter of the pulley acting as a
    # wheel; it falls inside the belt, so outside the actual pulley.
    pitch_diameter =  N * pitch / math.pi

    # The stock diameter is the maximum diameter of the pulley teeth,
    # corresponding to the stock circular diameter before cutting grooves.
    stock_diameter = pitch_diameter - tooth_height
    
    return { 'teeth' : N,
             'pitch_diameter' : pitch_diameter,
             'stock_diameter' : stock_diameter
             }
    

#================================================================
# Main script follows.  This sequence is executed when the script is initiated from the command line.

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Timing belt calculator.")
    parser.add_argument('--debug', action='store_true', help='Enable debugging logging to console.')
    parser.add_argument('-N1', type=int, help="Number of teeth on first pulley.")
    parser.add_argument('-N2', type=int, help="Number of teeth on second pulley.")
    parser.add_argument('-T', type=int, help="Number of teeth on timing belt.")
    parser.add_argument('-C', type=float,help="Distance between pulley centers in default units.")
    parser.add_argument('-sys', type=str, default='GT2', help="Belt system: GT2, MXL, XL, or L (default %(default)s).")
    
    args = parser.parse_args()

    system = {'GT2' : GT2, 'MXL' : MXL, 'XL' : XL, 'L' : L}.get(args.sys)
    if system is None:
        print("Warning: unrecognized belting system type, defaulting to GT2.")
        system = GT2
    
    # Choose a calculation depending on the combination of arguments.
    if args.N1 is not None and args.N2 is not None and args.C is not None and args.T is None:
        belt_teeth = timing_belt_length(args.N1, args.N2, args.C, system=system)
        print(f"Timing belt has {belt_teeth} teeth.")

    elif args.N1 is not None and args.N2 is not None and args.C is None and args.T is not None:
        center_distance = timing_belt_center_distance(args.N1, args.N2, args.T, system=system)
        print(f"Pulley center distance is %f %s." % (center_distance, system['units']))
        
    else:
        print("No calculation defined for this combination of arguments.")
