import numpy as np

def solve_channel_flow():
    # Problem parameters
    Re_tau = 395.0
    kappa = 0.42
    A = 25.4
    mu = 1.0 / Re_tau

    # Domain
    y_min = 0.0
    y_max = 2.0
    ny = 101
    y = np.linspace(y_min, y_max, ny)
    dy = y[1] - y[0]

    # Initial conditions
    u = np.zeros(ny)
    mu_t = np.zeros(ny)

    # Iteration parameters
    max_iter = 1000
    tolerance = 1e-6

    # Finite difference coefficients
    diag = np.zeros(ny)
    upper_diag = np.zeros(ny - 1)
    lower_diag = np.zeros(ny - 1)
    rhs = np.zeros(ny)

    # Iteration loop
    for iter in range(max_iter):
        # Compute effective viscosity
        y_plus = y * Re_tau
        mut_over_mu = 0.5 * (1 + (1/9) * kappa**2 * Re_tau**2 * (2*y - y**2)**2 * (3 - 4*y + 2*y**2)**2 * (1 - np.exp(-y_plus/A))**2)**0.5 - 0.5
        mu_eff = mu * (1 + mut_over_mu)

        # Assemble finite difference matrix
        for i in range(1, ny - 1):
            dmu_eff_dy = (mu_eff[i+1] - mu_eff[i-1]) / (2*dy)
            diag[i] = mu_eff[i] * (-2.0 / dy**2)
            upper_diag[i-1] = mu_eff[i] / dy**2 + dmu_eff_dy / (2*dy)
            lower_diag[i-1] = mu_eff[i] / dy**2 - dmu_eff_dy / (2*dy)
            rhs[i] = -1.0

        # Boundary conditions
        diag[0] = 1.0
        upper_diag[0] = 0.0
        rhs[0] = 0.0

        diag[ny-1] = 1.0
        lower_diag[ny-2] = 0.0
        rhs[ny-1] = 0.0

        # Solve tridiagonal system
        diag_modified = np.copy(diag)
        rhs_modified = np.copy(rhs)

        for i in range(1, ny):
            factor = lower_diag[i-1] / diag_modified[i-1]
            diag_modified[i] -= factor * upper_diag[i-1]
            rhs_modified[i] -= factor * rhs_modified[i-1]

        u_new = np.zeros(ny)
        u_new[ny-1] = rhs_modified[ny-1] / diag_modified[ny-1]
        for i in range(ny-2, -1, -1):
            u_new[i] = (rhs_modified[i] - upper_diag[i] * u_new[i+1]) / diag_modified[i]

        # Check convergence
        residual = np.max(np.abs(u_new - u))
        if residual < tolerance:
            break

        # Update solution
        u = np.copy(u_new)

    # Save the solution
    np.save('/opt/CFD-Benchmark/PDE_Benchmark/results/prediction/gemini/prompts/u_Fully_Developed_Turbulent_Channel_Flow.npy', u)

solve_channel_flow()