# CelerisLab/lbm/initializers.py
"""
Flow-field initialization helpers (vortex superposition, etc.).

All functions operate on an LBMField — they download DDF, modify on host,
and re-upload.  Only 2D D2Q9 vortices implemented for now.
"""

import numpy as np
from scipy.special import jv, expi
from typing import Tuple

import pycuda.driver as cuda


def add_vortex(field, center: Tuple[float, float],
               radius: float, strength: float,
               vortex_type: str = "lamb"):
    """Superimpose a vortex onto an existing 2D D2Q9 flow field.

    Supported types: "lamb", "oseen", "taylor".
    """
    if not field.cfg.is_d2q9:
        raise NotImplementedError("Vortex init only for D2Q9")
    if vortex_type not in ("lamb", "oseen", "taylor"):
        raise ValueError(f"Unknown vortex type: {vortex_type}")

    nx, ny = field.nx, field.ny
    x_c, y_c = center
    nu = field.cfg.viscosity

    x = np.linspace(-x_c, nx - 1 - x_c, nx)
    y = np.linspace(-y_c, ny - 1 - y_c, ny)
    X, Y = np.meshgrid(x, y)
    r = np.sqrt(X**2 + Y**2)
    theta = np.arctan2(Y, X)

    if vortex_type == "lamb":
        u_vor, v_vor, p_vor = _lamb_vortex(r, theta, radius, strength)
    elif vortex_type == "oseen":
        u_vor, v_vor, p_vor = _oseen_vortex(r, theta, radius, strength)
    else:
        u_vor, v_vor, p_vor = _taylor_vortex(r, theta, radius, strength, nu)

    # Download current DDF
    field.download_ddf()
    nq = field.nq
    f = field.ddf.reshape(nq, ny, nx)

    # Read lattice descriptors from CUDA module (single source of truth)
    cx, cy = field._read_lattice_vectors()
    w = field._read_lattice_weights()

    # Compute current macroscopic from DDF
    rho_old = np.sum(f, axis=0)
    rho_safe = np.where(np.abs(rho_old) > 1e-12, rho_old, 1.0)
    ux_old = sum(f[i] * cx[i] for i in range(nq)) / rho_safe
    uy_old = sum(f[i] * cy[i] for i in range(nq)) / rho_safe
    p_old = rho_old / 3.0

    # Superimpose
    # Vortex arrays are (ny, nx); u_vor[j, i] adds to node (i, j)
    u_new = ux_old + u_vor
    v_new = uy_old + v_vor
    p_new = p_old + p_vor

    # Reconstruct DDF from new macroscopic (equilibrium)
    for j in range(1, ny - 1):
        for i in range(1, nx - 1):
            u, v, p = float(u_new[j, i]), float(v_new[j, i]), float(p_new[j, i])
            for e in range(nq):
                eu = cx[e] * u + cy[e] * v
                u2 = u * u + v * v
                f[e, j, i] = w[e] * (3*p + 3*eu + 4.5*eu*eu - 1.5*u2)

    field.ddf = f.reshape(-1)
    field.upload_ddf()


# ---------------------------------------------------------------------------
# Vortex models
# ---------------------------------------------------------------------------
def _lamb_vortex(r, theta, radius, strength):
    b = 3.831705970207512
    n = b / radius
    u0 = strength
    inside = r <= radius
    outside = r > radius
    psi = np.zeros_like(r)
    psi[inside] = (
        (2*u0 / n / jv(0, b) * jv(1, n*r[inside]) - u0*r[inside])
        * np.sin(theta[inside])
    )
    psi[outside] = -u0 * radius**2 / r[outside] * np.sin(theta[outside])
    u_vor = np.gradient(psi, axis=0)
    v_vor = -np.gradient(psi, axis=1)
    p_vor = (
        -2*(np.gradient(v_vor, axis=1) - np.gradient(u_vor, axis=0))*psi
        - (u_vor**2 + v_vor**2) / 2
    )
    return u_vor, v_vor, p_vor


def _oseen_vortex(r, theta, radius, strength):
    kappa = 2 * np.pi * radius**2 * strength
    r_safe = np.where(r > 1e-12, r, 1e-12)
    exp_term = 1 - np.exp(-4 * r**2 / radius**2)
    u_vor = -kappa / (2*np.pi*r_safe) * exp_term * np.sin(theta)
    v_vor = kappa / (2*np.pi*r_safe) * exp_term * np.cos(theta)
    zeta = 4 * r**2 / radius**2
    p_vor = (
        -kappa**2 / (8*np.pi**2 * r_safe**2)
        * (-2*zeta*(expi(-zeta) - expi(-2*zeta))
           + (1 - np.exp(-zeta))**2)
    )
    return u_vor, v_vor, p_vor


def _taylor_vortex(r, theta, radius, strength, nu):
    M = strength * np.pi * radius**4 / (8 * nu)
    coeff = M * 4 * nu / radius**4
    exp_term = np.exp(-r**2 / radius**2)
    u_vor = -coeff * r * exp_term * np.sin(theta)
    v_vor = coeff * r * exp_term * np.cos(theta)
    p_vor = (
        -4 * M**2 * nu**2
        * np.exp(-2 * r**2 / radius**2)
        / (np.pi**2 * radius**6)
    )
    return u_vor, v_vor, p_vor