from scipy.optimize import root
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def model(x_list):
    d_ck, d_ch = x_list

    c0_kl = 0.1
    c0_hl = 0.1
    c0_cl = 0.2

    c0_kr = 0.1
    c0_hr = 0
    c0_cr = 0.1

    mu_dh = 0.727
    F = 96485
    eps = 80.2 * 8.81e-12
    R = 8.314
    T = 298

    c_kl = c0_kl - d_ck
    c_hl = c0_hl - d_ch

    c_kr = c0_kr + d_ck
    c_hr = c0_hr + d_ch

    d_phi = (F / eps) * (c_kl - c_kr + c_hl - c_hr - c0_cl + c0_cr)

    mu_l = 0.5 * (c_kl + c_hl + c0_cl)
    mu_r = 0.5 * (c_kr + c_hr + c0_cr)

    g_l = np.exp(-np.sqrt(mu_l/mu_dh))
    g_r = np.exp(-np.sqrt(mu_r/mu_dh))

    eq1 = np.log((g_l * c_kl)/(g_r * c_kr)) + (F/(R*T)) * d_phi
    eq2 = np.log((g_l * c_hl)/(g_r * c_hr)) + (F/(R*T)) * d_phi

    print(d_phi, '\n')

    return [eq1, eq2]



guess = [-0.05, 0.05]
ans = root(model, guess)

print(ans)

ck_l = 0.1 - ans.x[0]
ck_r = 0.1 + ans.x[0]
ch_l = 0.1 - ans.x[1]
ch_r = ans.x[1]

F = 96485
eps = 80.2 * 8.81e-12

d_phi = (F / eps) * (ck_l - ck_r + ch_l - ch_r - 0.2 + 0.1)

print('\nC_l =', round(ck_l + ch_l, 2))
print('C_r =', round(ck_r + ch_r, 2))

print('End concentrations: ')
print('|       L       |       R       |')
print('| C_K+ =', round(0.1 - ans.x[0], 4), '| C_K+ =', round(0.1 + ans.x[0], 4), '|')
print('| C_H+ =', round(0.1 - ans.x[1], 4), '| C_H+ =', round(ans.x[1], 4), '|')