#!/usr/bin/env python3
import numpy as np

# Parameters
Re_tau = 395.0
mu = 1.0 / Re_tau
kappa = 0.42
A = 25.4
# Domain
y_min = 0.0
y_max = 2.0
N = 201  # number of grid points
y = np.linspace(y_min, y_max, N)
dy = y[1] - y[0]

# Compute effective viscosity using the Cess turbulence model
# ues: mu_eff/mu = 0.5*( sqrt[1 + (1/9)*kappa^2*Re_tau^2*(2*y - y**2)**2*(3 - 4*y + 2*y**2)**2*(1 - exp(-y+ / A))**2] - 1)
y_plus = y * Re_tau
extra = (1.0 / 9.0) * (kappa**2) * (Re_tau**2) * ((2*y - y**2)**2) * ((3 - 4*y + 2*y**2)**2) * ((1 - np.exp(-y_plus/A))**2)
ratio = 0.5 * (np.sqrt(1.0 + extra) - 1.0)
mu_eff = mu * ratio

# Compute turbulent viscosity: mu_t = mu_eff - mu (if negative, set to zero)
mu_t = mu_eff - mu
mu_t[mu_t < 0] = 0.0

# Given initial conditions for other variables (they remain constant in this steady problem)
k = np.full_like(y, 0.01)
eps = np.full_like(y, 0.001)
omega = np.full_like(y, 1.0)
nu_SA = np.full_like(y, mu)  # 1/Re_tau

# Solve the steady momentum equation:
# d/dy( mu_eff * du/dy ) = -1, with u(0)=0, u(2)=0.
# Discretize using finite differences:
# For interior i, (mu_{i+1/2}*(u[i+1]-u[i]) - mu_{i-1/2}*(u[i]-u[i-1]))/dy^2 = -1.
# Use mu_{i+1/2} = 0.5*(mu_eff[i]+mu_eff[i+1]).
u = np.zeros_like(y)

# Assemble the coefficient matrix A and right-hand side b
A_matrix = np.zeros((N, N))
b = -np.ones(N)  # right-hand side is -1 at interior nodes

# Boundary conditions u[0]=0 and u[N-1]=0
A_matrix[0,0] = 1.0
b[0] = 0.0
A_matrix[-1,-1] = 1.0
b[-1] = 0.0

# Interior nodes
for i in range(1, N-1):
    mu_ip = 0.5 * (mu_eff[i] + mu_eff[i+1])  # mu_{i+1/2}
    mu_im = 0.5 * (mu_eff[i-1] + mu_eff[i])    # mu_{i-1/2}
    A_matrix[i, i-1] = -mu_im / dy**2
    A_matrix[i, i]   = (mu_im + mu_ip) / dy**2
    A_matrix[i, i+1] = -mu_ip / dy**2

# Solve the linear system
u = np.linalg.solve(A_matrix, b)

# Save the final solution variables in .npy files
np.save("u.npy", u)         # 1D array for velocity u
np.save("mu_t.npy", mu_t)   # 1D array for turbulent viscosity mu_t
np.save("k.npy", k)         # 1D array for turbulent kinetic energy k
np.save("eps.npy", eps)     # 1D array for turbulent dissipation eps
np.save("omega.npy", omega) # 1D array for specific dissipation omega
np.save("nu_SA.npy", nu_SA) # 1D array for Spalart-Allmaras viscosity nu_SA