# CELL 1 #
import numpy as np
import matplotlib.pyplot as plt

# Initial values

## Constants
g = -9.81              #[m/s^2]

## Pressures
p_ini = 2000000             #[Pa], 20 bar
p_inj = 34000000            #[Pa], 340 bar

## Host rock
k_host = 2.043*10**(-13)  #[m^2], 207mD

## Reservoir and well dimensions
X = 20                 #[m]
Y = 20                 #[m]
X_well = 20                #[m]
Y_well = 20                #[m]
k_well = 1*10**(-13)     #[m^2]

## Wetting (a) & Non-Wetting (b) properties
Rho_a = 700             #[kg/m^3]
V_a = 500
Rho_b = 1400            #[kg/m^3]
V_b = 1000
if Rho_a == Rho_b:
    print("Similair density not supported")

## Initial saturation of (a) and (b)
Saturation_a = 0
Saturation_b = 1
if Saturation_a + Saturation_b is not 1:
    print("Total saturation >1, check input!")

## Brooks-Corey parameters
disp_p = 1325
pore_size_dist = 1.68

n = 2.68      #[-]
m = 1-1/n     #[-]

# CELL 2 #
## Initial absolute pressure matrix
field_p = np.zeros((Y,X))
field_p[:,:] = p_ini
field_p[:Y_well,(X_well-1)] = p_inj

## Initial saturation matrix
Saturation_M_a = np.zeros((Y,X))
Saturation_M_a[:,:] = Saturation_a
Saturation_M_a[:Y_well,(X_well-1)] = Saturation_b

Saturation_M_b = np.subtract(np.ones((Y,X)),Saturation_M_a)

## Initial partial pressure matrix
field_p_a = np.multiply(Saturation_M_a, field_p)
field_p_b = np.multiply(Saturation_M_b, field_p)

## Absolute permeability matrix
k_M_reservoir = np.zeros((Y,X))
k_M_reservoir[:,:] = k_host
k_M_reservoir[:Y_well,(X_well-1)] = k_well

### Permeability matrices directional harmonic average
k_M_left = np.copy(k_M_reservoir)
for i in range(Y):
    for j in range(1,X):
        k_M_left[i,j] = (2*k_M_reservoir[i,j-1]*k_M_reservoir[i,j])/(k_M_reservoir[i,j-1]+k_M_reservoir[i,j])
k_M_right = np.copy(k_M_reservoir)
for i in range(Y):
    for j in range(X-1):
        k_M_right[i,j] = (2*k_M_reservoir[i,j+1]*k_M_reservoir[i,j])/(k_M_reservoir[i,j+1]+k_M_reservoir[i,j])
k_M_top = np.copy(k_M_reservoir)
for i in range(1,Y):
    for j in range(X):
        k_M_top[i,j] = (2*k_M_reservoir[i-1,j]*k_M_reservoir[i,j])/(k_M_reservoir[i-1,j]+k_M_reservoir[i,j])
k_M_bottom = np.copy(k_M_reservoir)
for i in range(Y-1):
    for j in range(X):
        k_M_bottom[i,j] = (2*k_M_reservoir[i+1,j]*k_M_reservoir[i,j])/(k_M_reservoir[i+1,j]+k_M_reservoir[i,j])

## Relative permeability matrices
Sew = np.zeros((Y,X))
for i in range(X):
    for j in range(Y):
        cap_p = np.abs((field_p_b[i,j]-field_p_a[i,j])/(Rho_a*g))
        if cap_p > disp_p:
            Sew[i,j] = (disp_p/cap_p)**pore_size_dist
        elif cap_p <= disp_p:
            Sew[i,j] = 1

k_M_a = np.zeros((Y,X))
for i in range(X):
    for j in range(Y):
        k_M_a[i,j] = (Sew[i,j]**(1/2))*(1-(1-Sew[i,j]**(1/m)**m))**2

k_M_b = np.zeros((Y,X))
for i in range(X):
    for j in range(Y):
        k_M_b[i,j] = ((1-Sew[i,j])**(1/2))*(1-Sew[i,j]**(1/m))**(2*m)
        if np.isnan(k_M_b[i,j]):
            k_M_b[i,j] = 0

# CELL 3 #
# Calculations

## Darcy velocity

### Total velocity
vy_field = np.zeros((Y,X))
for i in range(Y):
    for j in range(X-1):
        if i == Y-1:
            vy_field[i,j+1] = -(k_M_top[i,j]*(k_M_b[i,j]/V_b+k_M_a[i,j]/V_a)*(field_p[i,j]-field_p[i-1,j]))
        else:
            vy_field[i,j+1] = -(k_M_bottom[i,j]*(k_M_b[i,j]/V_b+k_M_a[i,j]/V_a)*(field_p[i+1,j]-field_p[i,j]))

vx_field = np.zeros((Y,X))
for i in range(Y):
    for j in range(X-1):
            vx_field[i,j+1] = -(k_M_right[i,j]*(k_M_b[i,j]/V_b+k_M_a[i,j]/V_a)*(field_p[i,j+1]-field_p[i,j]))

vt_field = np.hypot(vx_field, vy_field)

### Fluid a
vy_field_a = np.zeros((Y,X))
for i in range(Y):
    for j in range(X):
        vy_field_a[i,j] = ((k_M_a[i,j]/V_a)/(k_M_a[i,j]/V_a+k_M_b[i,j]/V_b))*vy_field[i,j] \
                          -((k_M_a[i,j]/V_a)/(k_M_a[i,j]/V_a+k_M_b[i,j]/V_b))*(k_M_b[i,j]/V_b) \
                          *(Rho_b-Rho_a)*k_M_reservoir[i,j]*g

vx_field_a = np.zeros((Y,X))
for i in range(Y):
    for j in range(X):
        vx_field_a[i,j] = ((k_M_a[i,j]/V_a)/(k_M_a[i,j]/V_a+k_M_b[i,j]/V_b))*vx_field[i,j]

v_field_a = np.hypot(vx_field_a, vy_field_a)

### Fluid b
vy_field_b = np.zeros((Y,X))
for i in range(Y):
    for j in range(X):
        vy_field_b[i,j] = ((k_M_b[i,j]/V_b)/(k_M_b[i,j]/V_b+k_M_a[i,j]/V_a))*vy_field[i,j] \
                          -((k_M_b[i,j]/V_b)/(k_M_b[i,j]/V_b+k_M_a[i,j]/V_a))*(k_M_a[i,j]/V_a) \
                          *(Rho_a-Rho_b)*k_M_reservoir[i,j]*g

vx_field_b = np.zeros((Y,X))
for i in range(Y):
    for j in range(X):
        vx_field_b[i,j] = ((k_M_b[i,j]/V_b)/(k_M_b[i,j]/V_b+k_M_a[i,j]/V_a))*vx_field[i,j]

v_field_b = np.hypot(vx_field_b, vy_field_b)

## Pressure matrices

### Total pressure matrix t+dt
field_pn = np.copy(field_p)
for i in range(Y):
    for j in range(1,X-1):
        if i == 0:
            field_pn[i,j] = (k_M_left[i,j]*(k_M_a[i,j-1]/V_a+k_M_b[i,j-1]/V_b)*field_p[i,j-1] \
                            +k_M_right[i,j]*(k_M_a[i,j+1]/V_a+k_M_b[i,j+1]/V_b)*field_p[i,j+1] \
                            +k_M_bottom[i,j]*(k_M_a[i+1,j]/V_a+k_M_b[i+1,j]/V_b)*field_p[i+1,j] \
                            -k_M_bottom[i,j]*(Rho_a*(k_M_a[i+1,j]/V_a)+Rho_b*(k_M_b[i+1,j]/V_b))*g \
                            )/(k_M_left[i,j-1]*(k_M_a[i,j]/V_a+k_M_b[i,j-1]/V_b) \
                            +k_M_right[i,j]*(k_M_a[i,j+1]/V_a+k_M_b[i,j+1]/V_b) \
                            +k_M_bottom[i,j]*(k_M_a[i+1,j]/V_a+k_M_b[i+1,j]/V_b))
        elif i == (Y-1):
            field_pn[i,j] = (k_M_left[i,j]*(k_M_a[i,j-1]/V_a+k_M_b[i,j-1]/V_b)*field_p[i,j-1] \
                            +k_M_right[i,j]*(k_M_a[i,j+1]/V_a+k_M_b[i,j+1]/V_b)*field_p[i,j+1] \
                            +k_M_top[i,j]*(k_M_a[i-1,j]/V_a+k_M_b[i-1,j]/V_b)*field_p[i-1,j] \
                            +k_M_top[i,j]*(Rho_a*(k_M_a[i-1,j]/V_a)+Rho_b*(k_M_b[i-1,j]/V_b))*g \
                            )/(k_M_left[i,j]*(k_M_a[i,j-1]/V_a+k_M_b[i,j-1]/V_b) \
                            +k_M_right[i,j]*(k_M_a[i,j+1]/V_a+k_M_b[i,j+1]/V_b) \
                            +k_M_top[i,j]*(k_M_a[i-1,j]/V_a+k_M_b[i-1,j]/V_b))
        else:
            field_pn[i,j] = (k_M_left[i,j]*(k_M_a[i,j-1]/V_a+k_M_b[i,j-1]/V_b)*field_p[i,j-1] \
                            +k_M_right[i,j]*(k_M_a[i,j+1]/V_a+k_M_b[i,j+1]/V_b)*field_p[i,j+1] \
                            +k_M_bottom[i,j]*(k_M_a[i+1,j]/V_a+k_M_b[i+1,j]/V_b)*field_p[i+1,j] \
                            -k_M_bottom[i,j]*(Rho_a*(k_M_a[i+1,j]/V_a)+Rho_b*(k_M_b[i+1,j]/V_b))*g \
                            +k_M_top[i,j]*(k_M_a[i-1,j]/V_a+k_M_b[i-1,j]/V_b)*field_p[i-1,j] \
                            +k_M_top[i,j]*(Rho_a*(k_M_a[i-1,j]/V_a)+Rho_b*(k_M_b[i-1,j]/V_b))*g \
                            )/(k_M_left[i,j]*(k_M_a[i,j-1]/V_a+k_M_b[i,j-1]/V_b) \
                            +k_M_right[i,j]*(k_M_a[i,j+1]/V_a+k_M_b[i,j+1]/V_b) \
                            +k_M_bottom[i,j]*(k_M_a[i+1,j]/V_a+k_M_b[i+1,j]/V_b) \
                            +k_M_top[i,j]*(k_M_a[i-1,j]/V_a+k_M_b[i-1,j]/V_b))

### Partial pressure change matrix fluid b
field_p_by = np.zeros((Y,X))
for i in range(Y):
    for j in range(X-1):
        A = vy_field_b[i,j+1]
        if i == Y-1:
            B = -(k_M_top[i,j]*(k_M_b[i,j]/V_b))
            if A == 0:
                field_p_by[i,j] = 0
            elif B == 0:
                field_p_by[i,j] = 0
            else:
                field_p_by[i,j] = A/B
        else:
            B = -(k_M_bottom[i,j]*(k_M_b[i,j]/V_b))
            if A == 0:
                field_p_by[i,j] = 0
            elif B == 0:
                field_p_by[i,j] = 0
            else:
                field_p_by[i,j] = A/B

field_p_bx = np.zeros((Y,X))
for i in range(Y):
    for j in range(X-1):
        A = vx_field_b[i,j+1]
        B = -(k_M_right[i,j]*(k_M_b[i,j]/V_b))
        if A == 0:
            field_p_bx[i,j] = 0
        elif B == 0:
            field_p_bx[i,j] = 0
        else:
            field_p_bx[i,j] = A/B

field_pcomb_b = np.add(field_p_by, field_p_bx)
field_pn_b = np.add(field_pcomb_b, field_p_b)

### Partial pressure change matrix fluid a
field_p_ay = np.zeros((Y,X))
for i in range(Y):
    for j in range(X-1):
        A = vy_field_a[i,j+1]
        if i == Y-1:
            B = -(k_M_top[i,j]*(k_M_a[i,j]/V_a))
            if A == 0:
                field_p_ay[i,j] = 0
            elif B == 0:
                field_p_ay[i,j] = 0
            else:
                field_p_ay[i,j] = A/B
        else:
            B = -(k_M_bottom[i,j]*(k_M_a[i,j]/V_a))
            if A == 0:
                field_p_ay[i,j] = 0
            elif B == 0:
                field_p_ay[i,j] = 0
            else:
                field_p_ay[i,j] = A/B

field_p_ax = np.zeros((Y,X))
for i in range(Y):
    for j in range(X-1):
        A = vx_field_a[i,j+1]
        B = -(k_M_right[i,j]*(k_M_a[i,j]/V_a))
        if A == 0:
            field_p_ax[i,j] = 0
        elif B == 0:
            field_p_ax[i,j] = field_pn[i,j]-field_p[i,j]
        else:
            field_p_ax[i,j] = A/B

field_pcomb_a = np.add(field_p_ay, field_p_ax)
field_pn_a = np.add(field_pcomb_a, field_p_a)

### Apply changes to pressure fields
field_p = np.copy(field_pn)
field_p_a = np.copy(field_pn_a)
field_p_b = np.copy(field_pn_b)

## Change in relative permeability-loop
Sew = np.zeros((Y,X))
for i in range(X):
    for j in range(Y):
        cap_p = np.abs((field_p_b[i,j]-field_p_a[i,j])/(Rho_a*g))
        if cap_p > disp_p:
            Sew[i,j] = (disp_p/cap_p)**pore_size_dist
        elif cap_p <= disp_p:
            Sew[i,j] = 1

knew_M_a = np.zeros((Y,X))
for i in range(X):
    for j in range(Y):
        knew_M_a[i,j] = (Sew[i,j]**(1/2))*(1-(1-Sew[i,j]**(1/m)**m))**2
k_M_a = np.copy(knew_M_a)

knew_M_b = np.zeros((Y,X))
for i in range(X):
    for j in range(Y):
        knew_M_b[i,j] = ((1-Sew[i,j])**(1/2))*(1-Sew[i,j]**(1/m))**(2*m)
        if np.isnan(knew_M_b[i,j]):
            knew_M_b[i,j] = 0
k_M_b = np.copy(knew_M_b)

# Plots
fig, (ax1,ax2) = plt.subplots(2,1,figsize=(15,15))
ax1.set_title("Total velocity vector field (m/s)")
Mr = np.around((vt_field*3600),decimals=2)
color = np.copy(vt_field)
Q = ax1.quiver(vx_field, vy_field, color, units='x', pivot='tail')
ax1.set(xlabel="Width [m]", ylabel="Height [m]")
ax1.set_aspect('equal')

plot1 = ax2.matshow(field_p,cmap=plt.get_cmap('plasma'))
ax2.set_title("Pressure field")
ax2.set(xlabel="Width [m]", ylabel="Depth [m]")
plt.show()