'''
Et skamstygt script som løser deler av øving 5
Plis ikke vis til noen andre
Jeg står ikke inne for dette
#If it's stupid and works, it aint stupid
Dette er ikke representativt for meg som person
'''

import numpy as np
import matplotlib.pyplot as plt
import scipy.optimize as opt

def Re_mu(rho, u, d, mu):
    return rho * abs(u) * d / mu

def drag_coeff(Re):
    if Re == 0:
        return 0
    elif Re < 1:
        return 24/Re
    elif Re < 1000:
        return 18.5 * np.power(Re, -3/5)
    else:
        return 0.44

def terminal(u, Re = None):
    if Re is None:
        Re = Re_mu(rho, u, dp, mu)
    Cd = drag_coeff(Re)
    return Vp * g * (rho_p - rho) - 0.5 * rho * abs(u)* u * Ap * Cd

####################################
# oppg. 2b

dp = 20e-6
rho_p = 2.7e3
rho = 7.1e3
mu = 5.5e-3
Vp = (np.pi/6) * dp**3
Ap = (np.pi/4) * dp**2
g = -9.81

u_res = opt.fsolve(terminal, [0.001])
print('2b')
print('u_terminal :', u_res)
print('Re :', Re_mu(rho, u_res, dp, mu))
print('sum(F) :', terminal(u_res), '\n')

u = np.linspace(0,0.001, 1000)
fig, ax = plt.subplots()
twin = ax.twinx()
f_line , = ax.plot(u, [terminal(ui) for ui in u], label=r'$\sum F_i$')
ax.plot([0,0.001], [0,0], color='black')
ax.scatter([u_res], [terminal(u_res)])
re_line , = twin.plot(u, [Re_mu(rho, ui, dp, mu) for ui in u], color= 'red', label = 'Re')
twin.set_ylabel('Re')
ax.set_xlabel('u [m/s]')
plt.legend(handles=[re_line, f_line])
plt.savefig('oving5/oppg2b')
plt.close(fig)

####################################
# oppg. 2c

dp = 5e-3
rho_p = 2.7e3
rho = 1.2
mu = 2e-5
Vp = (np.pi/6) * dp**3
Ap = (np.pi/4) * dp**2
g = -9.81

u_res = opt.fsolve(terminal, [-10])
print('2c')
print('u_terminal :', u_res)
print('Re :', Re_mu(rho, u_res, dp, mu))
print('sum(F) :', terminal(u_res), '\n')

u_max_plot = 0
u_min_plot = -25
u = np.linspace(u_min_plot,u_max_plot, 1000)
fig, ax = plt.subplots()
twin = ax.twinx()
f_line , = ax.plot(u, [terminal(ui) for ui in u], label=r'$\sum F_i$')
ax.plot([u_min_plot,u_max_plot], [0,0], color='black')
ax.scatter([u_res], [terminal(u_res)])
re_line , = twin.plot(u, [Re_mu(rho, ui, dp, mu) for ui in u], color= 'red', label = 'Re')
twin.set_ylabel('Re')
ax.set_xlabel('u [m/s]')
plt.legend(handles=[f_line, re_line])
plt.savefig('oving5/oppg2c')
plt.close(fig)

####################################
# oppg. 2di

dp = 8e-3
rho_p = 1.2
rho = 1e3
mu = 1e-3
Vp = (np.pi/6) * dp**3
Ap = (np.pi/4) * dp**2
g = -9.81

res_u = opt.fsolve(lambda u: Re_mu(rho, u, dp, mu) - 1000, [10])
print('2di')
print('u_terminal :', res_u)
print('Re :', Re_mu(rho, res_u, dp, mu))
print('sum(F) :', terminal(res_u, Re=1000), '\n')

####################################
# oppg. 2dii

dp = 8e-3
rho_p = 1.2
rho = 1e3
mu = 1e-3
Vp = (np.pi/6) * dp**3
Ap = (np.pi/4) * dp**2
g = -9.81

res_u = opt.fsolve(lambda u: Re_mu(rho, u, dp, mu) - 2000, [10])
print('2dii')
print('u_terminal :', res_u)
print('Re :', Re_mu(rho, res_u, dp, mu))
print('sum(F) :', terminal(res_u, Re=2000), '\n')

####################################
# oppg. 3

def root_func(unknowns):
    rho = 1.2
    mu = 1.82e-5

    eps, dp = unknowns

    eq1 = 1.75 * rho * (1 - eps) / (rho * eps**3) - (3.27 - 2.2)*1e3 / 0.25
    eq2 = 150 * mu * (1 - eps)**2 / (dp**2 * eps**3) - 2.2e3

    return [eq1, eq2]

res = opt.root(root_func, [0.5, 1e-3])
print('Oppg. 3')
print('Void fraction :', res.x[0])
print('Particle diameter :', res.x[1], 'm')
print('Root_val :', res.fun)