from gekko import GEKKO
m = GEKKO()

# Constants
pi = 3.14159
a = 1.67   # feet
p = 168.5  # density lb/ft^3
u = 0.8
s = 75000  # stress lbf/ft^2
t = 0.0092 # ft

# Variables
w = m.Var(lb=0.01,ub=0.5)
d = [m.Var(lb=1/5) for i in range(5)]
N_o = [950, 650, 450, 250, 150]
N_in = 350

weight = m.Var()

# Intermediates
# Diameter for the ith step of the input pulley
d_in = [m.Intermediate(d[i]*(N_o[i]/N_in)) for i in range(5)] 

# Intermediate for objective function
Inter = [m.Intermediate(d[i]**2 + d_in[i]**2) for i in range(5)]

# Belt lengths
C = [m.Intermediate((pi*d[i])/2*(1 + N_o[i]/N_in) \
                    + ((((N_o[i]/N_in) - 1)**2)*d[i]**2)/(4*a) \
                    + 2*a) for i in range(5)]

# Angles of lap of the belt over the ith pulley step
O = [m.Intermediate(pi - 2*m.asin((((N_o[i]/N_in) \
        - 1)*d[i])/(2*a))) for i in range(5)]

# Tensions on the tight side of the ith step (limitation on max tension)
T1 = [m.Intermediate(s*t*w) for i in range(5)]

# Tensions on the slack side of the ith step
T2 = [m.Intermediate(T1[i]/(m.exp(u*O[i]))) for i in range(5)]

# Equations
# Weight, Objective Function
m.Equation(weight == p*w*(pi/4)*(sum(Inter[0:5])))

# Constraints
belt_length = [m.Equation(C[0] - C[i + 1] == 0) for i in range (4)]
tension_ratio = [m.Equation(m.exp(u*O[i]) >= 2) for i in range(5)]
power_transmitted = [m.Equation(((T1[i]-T2[i])*pi*d_in[i]*350)/33000 \
                                >= 0.65) for i in range(5)]

# Objective and Solve
m.Minimize(weight)

m.options.IMODE = 3
m.options.SOLVER = 1
m.solve()

print('Optimal weight: ' + str(weight[0]))
print('Optimal diameter: ' + str(d))
print('Optimal width: ' + str(w[0]))