#!/usr/bin env python3
from formulary import *
from scipy.constants import gas_constant, Avogadro, elementary_charge

Faraday = Avogadro * elementary_charge

# BUTLER VOLMER / TAFEL

@np.vectorize
def fbutler_volmer_anode(ac, z, eta, T):
    return np.exp((1-ac)*z*Faraday*eta/(gas_constant*T))

@np.vectorize
def fbutler_volmer_cathode(ac, z, eta, T):
    return -np.exp(-ac*z*Faraday*eta/(gas_constant*T))

def fbutler_volmer(ac, z, eta, T):
    return fbutler_volmer_anode(ac, z, eta, T) + fbutler_volmer_cathode(ac, z, eta, T)

def butler_volmer():
    fig, ax = plt.subplots(figsize=size_formula_fill_default)
    ax.set_xlabel("$\\eta$ [V]")
    ax.set_ylabel("$j/j_0$")
    etas = np.linspace(-0.1, 0.1, 400)
    T = 300
    z = 1.0
    # other a
    ac2, ac3 = 0.2, 0.8
    i2 = fbutler_volmer(0.2, z, etas, T)
    i3 = fbutler_volmer(0.8, z, etas, T)
    ax.plot(etas, i2, color="blue", linestyle="dashed", label=f"$\\alpha_\\text{{C}}={ac2}$")
    ax.plot(etas, i3, color="green", linestyle="dashed", label=f"$\\alpha_\\text{{C}}={ac3}$")
    # 0.5
    ac = 0.5
    irel_anode = fbutler_volmer_anode(ac, z, etas, T)
    irel_cathode = fbutler_volmer_cathode(ac, z, etas, T)
    ax.plot(etas, irel_anode,  color="gray")
    ax.plot(etas, irel_cathode, color="gray")
    ax.plot(etas, irel_cathode + irel_anode, color="black", label=f"$\\alpha_\\text{{C}}=0.5$")
    ax.grid()
    ax.legend()
    ylim = 6
    ax.set_ylim(-ylim, ylim)
    return fig

@np.vectorize
def ftafel_anode(ac, z, eta, T):
    return 10**((1-ac)*z*Faraday*eta/(gas_constant*T*np.log(10)))

@np.vectorize
def ftafel_cathode(ac, z, eta, T):
    return -10**(-ac*z*Faraday*eta/(gas_constant*T*np.log(10)))

def tafel():
    i0 = 1
    ac = 0.2
    z = 1
    T = 300
    eta_max = 0.2
    etas = np.linspace(-eta_max, eta_max, 400)
    i = np.abs(fbutler_volmer(ac, z, etas ,T))
    iright = i0 * np.abs(ftafel_cathode(ac, z, etas, T))
    ileft = i0 * ftafel_anode(ac, z, etas, T)

    fig, ax = plt.subplots(figsize=size_formula_normal_default)
    ax.set_xlabel("$\\eta$ [V]")
    ax.set_ylabel("$\\log_{10}\\left(\\frac{|j|}{j_0}\\right)$")
    # ax.set_ylabel("$\\log_{10}\\left(|j|/j_0\\right)$")
    ax.set_yscale("log")
    # ax.plot(etas, linear, label="Tafel slope")
    ax.plot(etas[etas >= 0], ileft[etas >= 0],  linestyle="dashed", color="gray", label="Tafel Approximation")
    ax.plot(etas[etas <= 0], iright[etas <= 0], linestyle="dashed", color="gray")
    ax.plot(etas, i, label=f"Butler-Volmer $\\alpha_\\text{{C}}={ac:.1f}$")
    ax.legend()
    ax.grid()
    return fig


# NYQUIST
@np.vectorize
def fZ_ohm(R, omega):
    return R

@np.vectorize
def fZ_cap(C, omega):
    return 1/(1j*omega*C)

@np.vectorize
def fZ_ind(L, omega):
    return (1j*omega*L)


def nyquist():
    fig, ax = plt.subplots(figsize=size_formula_fill_default)
    split_z = lambda Z: (Z.real, -Z.imag)
    ax.grid()
    ax.set_xlabel("$\\text{Re}(Z)$ [\\si{\\ohm}]")
    ax.set_ylabel("$-\\text{Im}(Z)$ [\\si{\\ohm}]")
    # ax.scatter(*split_z(Z_series), label="series")


    R1 = 20
    R2 = 5
    RS = 7.5
    C1 = 1e-4
    C2 = 1e-6
    ws1 = np.power(10, np.linspace(1, 8, 1000))


    Z_ohm1 = fZ_ohm(R1, ws1)
    Z_ohm2 = fZ_ohm(R2, ws1)
    Z_ohmS = fZ_ohm(RS, ws1)
    Z_cap1 = fZ_cap(C1, ws1)
    Z_cap2 = fZ_cap(C2, ws1)
    Z_parallel1 = 1/(1/Z_ohm1 + 1/Z_cap1)
    Z_parallel2 = 1/(1/Z_ohm2 + 1/Z_cap2)
    Z_cell = Z_parallel1 + Z_parallel2 + Z_ohmS
    ax.scatter(*split_z(Z_parallel1), label="Parallel $C_1,R_1$")
    ax.scatter(*split_z(Z_parallel2), label="Parallel $C_2,R_2$")
    ax.scatter(*split_z(Z_cell), label="P1 + $R_3$ + P2")
    ax.scatter(*split_z(Z_cap1), label=f"$C_1=\\SI{{{C1:.0e}}}{{\\farad}}$")
    ax.scatter(*split_z(Z_ohm1), label=f"$R_1 = \\SI{{{R1}}}{{\\ohm}}$")

    # wmax1 = 1/(R1 * C1)
    # ZatWmax1 = Z_parallel1[np.argmin(ws1 - wmax1)]
    # print(ws1[0], ws1[-1])
    # print(wmax1, ZatWmax1)
    # ax.scatter(*split_z(ZatWmax1), color="red")
    # ax.scatter(*split_z(Z_cell1), label="cell")
    # ax.scatter(*split_z(Z_ohm2), label="ohmic")
    # ax.scatter(*split_z(Z_cell2), label="cell")
    ax.axis('equal')
    ax.set_ylim(0,R1*1.1)
    ax.legend()
    return fig

def fZ_tlm(Rel, Rion, Rct, Cct, ws, N):
    Zion = fZ_ohm(Rion, ws)
    Zel = fZ_ohm(Rel, ws)
    Zct = 1/(1/fZ_ohm(Rct, ws) + 1/fZ_cap(Cct, ws))

    Z = Zct
    for _ in range(N):
        Z = Zion + 1/(1/Zct + 1/Z)
    Z += Zel
    return Z

def nyquist_tlm():
    fig, ax = plt.subplots(figsize=(width_formula, width_formula*0.5))
    split_z = lambda Z: (Z.real, -Z.imag)
    ax.grid()
    ax.set_xlabel("$\\text{Re}(Z)$ [\\si{\\ohm}]")
    ax.set_ylabel("$-\\text{Im}(Z)$ [\\si{\\ohm}]")
    Rct1 = 300
    Rct2 = 100
    Rion = 10
    ws = np.power(10, np.linspace(1e-6, 5, 1000))
    Z1 = fZ_tlm(0, Rion, Rct1, 1e-4, ws, 100)
    Z2 = fZ_tlm(0, Rion, Rct2, 1e-4, ws, 100)
    ax.scatter(*split_z(Z1), label=f"$R_\\text{{ct}} = \\SI{{{Rct1}}}{{\\ohm}}$", marker=".")
    ax.scatter(*split_z(Z2), label=f"$R_\\text{{ct}} = \\SI{{{Rct2}}}{{\\ohm}}$", marker=".")
    ax.axis('equal')
    # ax.set_ylim(0,R1*1.1)
    ax.legend()
    return fig

def fkohlrausch(L0, K, c):
    return L0 - K*np.sqrt(c)

def kohlrausch():
    fig, ax = plt.subplots(figsize=size_formula_small_quadratic)
    ax.grid()
    ax.set_xlabel("$c_\\text{salt}$")
    ax.set_ylabel("$\\Lambda_\\text{M}$")
    L0 = 10
    K1 = 1
    K2 = 2
    cs = np.linspace(0, 10)
    L1 = fkohlrausch(L0, K1, cs)
    L2 = fkohlrausch(L0, K2, cs)
    ax.plot(cs, L1, label=f"$K={K1}$")
    ax.plot(cs, L2, label=f"$K={K2}$")
    ax.legend()
    return fig

if __name__ == '__main__':
    export(butler_volmer(), "ch_butler_volmer")
    export(tafel(), "ch_tafel")
    export(nyquist(), "ch_nyquist")
    export(nyquist_tlm(), "ch_nyquist_tlm")
    export(kohlrausch(), "ch_kohlrausch")