import matplotlib
matplotlib.use('Agg')
import numpy as np
import scipy.ndimage as ndimage
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.patches import Circle
from matplotlib.colors import LinearSegmentedColormap
import gradio as gr
import tempfile

# ============================================================
# ПАРАМЕТРЫ СЕТКИ LBM И ЦВЕТОВАЯ ПАЛИТРА
# ============================================================
N = 250
X, Y = np.meshgrid(np.arange(N), np.arange(N))
c = np.array([[0,0],[1,0],[0,1],[-1,0],[0,-1],[1,1],[-1,1],[-1,-1],[1,-1]], dtype=float)
w = np.array([4/9] + [1/9]*4 + [1/36]*4, dtype=float)
opp = np.array([0, 3, 4, 1, 2, 7, 8, 5, 6], dtype=int)

cfd_multicolor_cmap = LinearSegmentedColormap.from_list("cfd_multicolor",
    [(0.00,"#004d99"), (0.10, "#00b7ff"), (0.25, "#00ffd0"), (0.40, "#00ff66"),  
    (0.60, "#dfff00"), (0.85, "#ff9a00"), (1.00, "#ff2200")], N=512)
cfd_multicolor_cmap.set_bad(color='#1c1c1c')

def equilibrium(rho, u):
    feq = np.zeros((9, N, N), dtype=float)
    usqr = u[0]**2 + u[1]**2
    for i in range(9):
        cu = 3.0 * (c[i,0] * u[0] + c[i,1] * u[1])
        feq[i] = rho * w[i] * (1.0 + cu + 0.5*cu**2 - 1.5*usqr)
    return feq

def equilibrium_1d(rho_1d, u_1d):
    feq = np.zeros((9, len(rho_1d)), dtype=float)
    usqr = u_1d[0]**2 + u_1d[1]**2
    for i in range(9):
        cu = 3.0 * (c[i,0] * u_1d[0] + c[i,1] * u_1d[1])
        feq[i] = rho_1d * w[i] * (1.0 + cu + 0.5*cu**2 - 1.5*usqr)
    return feq

# ============================================================
# ГЕОМЕТРИЯ: ЛАГРАНЖЕВЫЕ ЯЧЕЙКИ + ЭЙЛЕРОВА МАСКА
# ============================================================
def get_fish_geometry(cx, cy, theta, phase, L=65.0, d_max=9.0, n_pts=50):
    s = np.linspace(0, L, n_pts)
    envelope = (0.04 + 0.22 * (s / L)**2) * L
    y_mid = envelope * np.sin(2.0 * np.pi * (s / L - phase))
    
    ds = L / (n_pts - 1)
    dy_mid_ds = np.gradient(y_mid, ds)
    dx_mid_ds = np.ones_like(dy_mid_ds)
    
    mag = np.hypot(dx_mid_ds, dy_mid_ds) + 1e-10
    nx_loc = -dy_mid_ds / mag
    ny_loc = dx_mid_ds / mag
    
    s_norm = s / L
    h = d_max * np.sin(np.pi * (s_norm**0.6)) * (1.0 - 0.65 * s_norm)
    h = np.maximum(h, 1.0)
    
    x_loc = s - L / 2.0
    
    xl_loc = x_loc + h * nx_loc
    yl_loc = y_mid + h * ny_loc
    xr_loc = x_loc - h * nx_loc
    yr_loc = y_mid - h * ny_loc
    
    cos_t, sin_t = np.cos(theta), np.sin(theta)
    
    def rot_trans(xl, yl):
        return cx + xl * cos_t - yl * sin_t, cy + xl * sin_t + yl * cos_t
        
    def rot_vec(vx, vy):
        return vx * cos_t - vy * sin_t, vx * sin_t + vy * cos_t
        
    x_c_glob, y_c_glob = rot_trans(x_loc, y_mid)
    
    X_left, Y_left = rot_trans(xl_loc, yl_loc)
    X_right, Y_right = rot_trans(xr_loc, yr_loc)
    
    NX_left, NY_left = rot_vec(nx_loc, ny_loc)
    NX_right, NY_right = rot_vec(-nx_loc, -ny_loc)
    
    X_bnd = np.concatenate([X_left, X_right[::-1]])
    Y_bnd = np.concatenate([Y_left, Y_right[::-1]])
    NX_bnd = np.concatenate([NX_left, NX_right[::-1]])
    NY_bnd = np.concatenate([NY_left, NY_right[::-1]])
    
    dx_grid = X[:, :, None] - x_c_glob
    dy_grid = Y[:, :, None] - y_c_glob
    dist_sq = dx_grid**2 + dy_grid**2
    solid = np.any(dist_sq < (h**2), axis=-1)
    idx_nearest = np.argmin(dist_sq, axis=-1)
    
    eroded = ndimage.binary_erosion(solid, structure=np.ones((3, 3), dtype=bool))
    boundary = solid & ~eroded
    
    return solid, boundary, x_c_glob, y_c_glob, idx_nearest, X_bnd, Y_bnd, NX_bnd, NY_bnd

# ============================================================
# ОСНОВНОЙ ДВИЖОК
# ============================================================
def simulate(frames=400, frequency=3.0, relative_density=1.0, tau=0.58, vis_mode="Pressure"):
    
    omega_lbm = 1.0 / tau
    rho = np.ones((N, N), dtype=float)
    u = np.zeros((2, N, N), dtype=float)
    f = equilibrium(rho, u)
    
    L_fish = 65.0
    cm_x, cm_y = float(N) * 0.25, float(N) * 0.5 
    v_cm_x, v_cm_y = 0.0, 0.0
    theta, omega_h = 0.0, 0.0
    phase = 0.0
    
    dt_lbm = 1.0
    LBM_DT = 0.003 
    lbm_frequency = frequency * LBM_DT
    
    Ff_x, Ff_y, Tf = 0.0, 0.0, 0.0
    
    solid_old, _, _, _, _, X_bnd_old, Y_bnd_old, _, _ = get_fish_geometry(cm_x, cm_y, theta, phase, L=L_fish)
    
    sponge_thickness = 20
    sponge_sigma = np.zeros((N, N), dtype=float)
    for i in range(sponge_thickness):
        damping = ((sponge_thickness - i) / sponge_thickness)**2 * 0.15 
        sponge_sigma[i, :] = np.maximum(sponge_sigma[i, :], damping)
        sponge_sigma[-1-i, :] = np.maximum(sponge_sigma[-1-i, :], damping)
        sponge_sigma[:, i] = np.maximum(sponge_sigma[:, i], damping)
        sponge_sigma[:, -1-i] = np.maximum(sponge_sigma[:, -1-i], damping)
        
    f_rest = equilibrium(np.ones((N, N)), np.zeros((2, N, N)))
    
    fig, ax = plt.subplots(figsize=(8, 8), facecolor='#121212')
    ax.set_facecolor('#121212')
    ax.axis('off')
    fig.tight_layout()
    imgs = []

    # --------------------------------------------------------
    # ИНИЦИАЛИЗАЦИЯ: СКАЛЯРНАЯ ЭНЕРГИЯ И УСРЕДНЕННЫЙ ВЕКТОР НАПРАВЛЕНИЯ
    # --------------------------------------------------------
    N_bnd = len(X_bnd_old) 
    stored_energy = np.zeros(N_bnd, dtype=float)
    stored_dir_x = np.zeros(N_bnd, dtype=float)  # Хранит усредненное направление X
    stored_dir_y = np.zeros(N_bnd, dtype=float)  # Хранит усредненное направление Y
    P_old = np.zeros(N_bnd, dtype=float) 
    
    for t in range(frames):
        phase_new = phase + lbm_frequency
        solid_new, boundary_new, x_pts, y_pts, idx_nearest, X_bnd, Y_bnd, NX_bnd, NY_bnd = get_fish_geometry(
            cm_x, cm_y, theta, phase_new, L=L_fish
        )
        
        v_skin_x = (X_bnd - X_bnd_old) / dt_lbm
        v_skin_y = (Y_bnd - Y_bnd_old) / dt_lbm
        X_bnd_old, Y_bnd_old = np.copy(X_bnd), np.copy(Y_bnd)
        
        area_fish = np.sum(solid_new)
        m_fish = relative_density * area_fish
        M_eff = m_fish * 1.2 
        I_eff = 0.5 * m_fish * (L_fish / 3.5)**2
        
        mask_fluid = ~solid_old
        rho_sum = np.sum(f, axis=0)
        u[0] = np.sum(f * c[:,0].reshape(9,1,1), axis=0) / (rho_sum + 1e-5)
        u[1] = np.sum(f * c[:,1].reshape(9,1,1), axis=0) / (rho_sum + 1e-5)
        
        max_fluid_u = 0.25 
        u[0] = np.clip(u[0], -max_fluid_u, max_fluid_u)
        u[1] = np.clip(u[1], -max_fluid_u, max_fluid_u)
        u[0, solid_old] = 0.0
        u[1, solid_old] = 0.0
        
        feq = equilibrium(rho_sum, u)
        f[:, mask_fluid] = f[:, mask_fluid] * (1 - omega_lbm) + omega_lbm * feq[:, mask_fluid]
        f = f * (1.0 - sponge_sigma) + f_rest * sponge_sigma
        
        for i in range(9):
            f[i] = np.roll(f[i], c[i].astype(int), axis=(1,0))
            
        uncovered = solid_old & ~solid_new 
        r_vec_x, r_vec_y = X - cm_x, Y - cm_y
        
        if np.any(uncovered):
            alpha = 0.25
            ux_unc = alpha * v_cm_x + (1-alpha) * u[0][uncovered]
            uy_unc = alpha * v_cm_y + (1-alpha) * u[1][uncovered]
            rho_unc = np.ones_like(ux_unc)
            feq_refill = equilibrium_1d(rho_unc, np.array([ux_unc, uy_unc]))
            for i in range(9):
                f[i, uncovered] = feq_refill[i]
                
        F_fluid_x_raw, F_fluid_y_raw, tau_fluid_raw = 0.0, 0.0, 0.0
        f_new = np.copy(f)
        y_coords, x_coords = np.where(boundary_new)
        
        for i in range(1, 9):
            o = opp[i]
            src_y = np.clip(y_coords - int(c[i, 1]), 0, N-1)
            src_x = np.clip(x_coords - int(c[i, 0]), 0, N-1)
            is_fluid_source = (~solid_old[src_y, src_x]) | (~solid_new[src_y, src_x])
            
            valid_y = y_coords[is_fluid_source]
            valid_x = x_coords[is_fluid_source]
            
            if len(valid_y) == 0: continue
                
            rx_bb, ry_bb = r_vec_x[valid_y, valid_x], r_vec_y[valid_y, valid_x]
            vw_rot_x, vw_rot_y = -omega_h * ry_bb, omega_h * rx_bb
            
            near_idx = idx_nearest[valid_y, valid_x]
            vw_bend_x = v_skin_x[near_idx]
            vw_bend_y = v_skin_y[near_idx]
            
            v_wall_x = v_cm_x + vw_rot_x + vw_bend_x * 0.5
            v_wall_y = v_cm_y + vw_rot_y + vw_bend_y * 0.5
            
            cu_wall = 3.0 * (c[i,0] * v_wall_x + c[i,1] * v_wall_y)
            rho_wall = rho_sum[valid_y, valid_x]
            f_in = f[i, valid_y, valid_x]
            
            f_out = f_in - 2.0 * w[i] * rho_wall * cu_wall
            f_new[o, valid_y, valid_x] = f_out
            
            dp_x = (f_in + f_out) * c[i,0]
            dp_y = (f_in + f_out) * c[i,1]
            
            F_fluid_x_raw += np.sum(dp_x)
            F_fluid_y_raw += np.sum(dp_y)
            tau_fluid_raw += np.sum(rx_bb * dp_y - ry_bb * dp_x)
            
        f = f_new
        solid_old = np.copy(solid_new)
        phase = phase_new
        
        rho_sum = np.sum(f, axis=0)
        pressure = rho_sum - 1.0

        # ============================================================
        # ФИЗИЧЕСКИЙ РАСЧЕТ РАБОТЫ, ЭНЕРГИИ И АДВЕКЦИИ
        # ============================================================
        offset = 1.5 
        X_sample = np.clip(X_bnd + offset * NX_bnd, 0, N-1)
        Y_sample = np.clip(Y_bnd + offset * NY_bnd, 0, N-1)
        
        P_bnd = ndimage.map_coordinates(pressure, [Y_sample, X_sample], order=1, mode='nearest')
        u_f_x = ndimage.map_coordinates(u[0], [Y_sample, X_sample], order=1, mode='nearest')
        u_f_y = ndimage.map_coordinates(u[1], [Y_sample, X_sample], order=1, mode='nearest')
        
        v_rel_x = u_f_x - v_skin_x
        v_rel_y = u_f_y - v_skin_y
        
        T_x = np.roll(X_bnd, -1) - np.roll(X_bnd, 1)
        T_y = np.roll(Y_bnd, -1) - np.roll(Y_bnd, 1)
        T_mag = np.hypot(T_x, T_y) + 1e-10
        T_x /= T_mag
        T_y /= T_mag
        
        v_tangent = v_rel_x * T_x + v_rel_y * T_y
        
        ds_approx = (2.0 * L_fish) / N_bnd
        v_idx = v_tangent / ds_approx  
        
        # Перенос энергии и направления работы вдоль контура
        mobility = 1.0 
        source_indices = (np.arange(N_bnd) - v_idx * mobility * dt_lbm) % N_bnd
        
        stored_energy = ndimage.map_coordinates(stored_energy, [source_indices], order=1, mode='wrap')
        stored_dir_x = ndimage.map_coordinates(stored_dir_x, [source_indices], order=1, mode='wrap')
        stored_dir_y = ndimage.map_coordinates(stored_dir_y, [source_indices], order=1, mode='wrap')
        P_old = ndimage.map_coordinates(P_old, [source_indices], order=1, mode='wrap')
        
        # Нормализуем направление после интерполяции, чтобы оно не затухало
        dir_norm = np.hypot(stored_dir_x, stored_dir_y) + 1e-12
        stored_dir_x /= dir_norm
        stored_dir_y /= dir_norm
        
        # Сила давления жидкости НА рыбу
        Fx_p = -P_bnd * NX_bnd
        Fy_p = -P_bnd * NY_bnd
        
        power = Fx_p * v_skin_x + Fy_p * v_skin_y
        dW = power * dt_lbm 
        
        # Направление воздействия рыбы НА воду (противоположно движению кожи)
        vmag_skin = np.hypot(v_skin_x, v_skin_y) + 1e-12
        dir_x = -v_skin_x / vmag_skin
        dir_y = -v_skin_y / vmag_skin
        
        # Накопление: только когда P < 0 (формируется зона разрежения/вихрь) и работа положительна
        acc_mask = (P_bnd < 0.0) & (dW > 0.0)
        
        # Сначала запоминаем старую энергию для корректного усреднения вектора
        old_energy = stored_energy[acc_mask]
        stored_energy[acc_mask] += dW[acc_mask]
        
        # Взвешенное усреднение направления: не уничтожает скалярную энергию при противоположных импульсах
        new_vec_x = stored_dir_x[acc_mask] * old_energy + dir_x[acc_mask] * dW[acc_mask]
        new_vec_y = stored_dir_y[acc_mask] * old_energy + dir_y[acc_mask] * dW[acc_mask]
        
        norm_vec = np.hypot(new_vec_x, new_vec_y) + 1e-12
        stored_dir_x[acc_mask] = new_vec_x / norm_vec
        stored_dir_y[acc_mask] = new_vec_y / norm_vec
        
        # Отдача: происходит только там, где есть энергия и давление начало расти (схлопывание)
        dP = P_bnd - P_old
        release_mask = (stored_energy > 1e-12) & (dP > 0.0)
        
        release_fraction = np.zeros_like(stored_energy)
        release_fraction[release_mask] = np.clip(dP[release_mask] / (np.abs(P_old[release_mask]) + 1e-8), 0.0, 1.0)
        
        released_energy = stored_energy * release_fraction
        
        # Отдаваемая сила направлена туда же, куда направлено усредненное воздействие на воду
        Fx_release = stored_dir_x * released_energy
        Fy_release = stored_dir_y * released_energy
        
        # Уменьшаем только скалярную энергию, вектор усредненного направления остается (он просто перестает действовать, когда E=0)
        stored_energy -= released_energy
        
        P_old = P_bnd.copy()
        
        # Превращение высвобожденной энергии в силу толчка (через F = E / dt)
        thrust_multiplier = 10.5
        F_vortex_x = (np.sum(Fx_release) / dt_lbm) * thrust_multiplier
        F_vortex_y = (np.sum(Fy_release) / dt_lbm) * thrust_multiplier
                                        
        # ============================================================
        
        alpha_force = 0.3 
        Ff_x += alpha_force * (F_fluid_x_raw - Ff_x)
        Ff_y += alpha_force * (F_fluid_y_raw - Ff_y)
        Tf   += alpha_force * (tau_fluid_raw - Tf)
        
        Total_F_x = Ff_x + F_vortex_x
        Total_F_y = Ff_y + F_vortex_y
        
        v_cm_x += (Total_F_x / M_eff) * dt_lbm
        v_cm_y += (Total_F_y / M_eff) * dt_lbm
            
        max_v = 0.15 
        v_cm_x = np.clip(v_cm_x, -max_v, max_v)
        v_cm_y = np.clip(v_cm_y, -max_v, max_v)
        
        cm_x += v_cm_x * dt_lbm
        cm_y += v_cm_y * dt_lbm
        
        margin = L_fish * 0.6
        cm_x = np.clip(cm_x, margin, N - margin)
        cm_y = np.clip(cm_y, margin, N - margin)
        
        omega_h += (Tf / I_eff) * dt_lbm
        theta += omega_h * dt_lbm
        
        # ============================================================
        # ВИЗУАЛИЗАЦИЯ
        # ============================================================
        if t % 2 == 0: 
            if vis_mode == "Pressure":
                data = np.nan_to_num(pressure, nan=0.0)
                vabs = float(np.percentile(np.abs(data), 99.5)) or 1e-5
                vmin, vmax = -vabs, vabs
                cmap = cfd_multicolor_cmap
            elif vis_mode == "Pressure Deviation":
                data = np.nan_to_num(np.abs(pressure), nan=0.0)
                vmax = float(np.percentile(data, 99.7)) or 1e-5
                vmin = 0.0
                cmap = cfd_multicolor_cmap
                
            data[solid_new] = np.nan
            im = ax.imshow(data, cmap=cmap, animated=True, vmin=vmin, vmax=vmax, interpolation='bilinear')
            frame_artists = [im]
            
            fish_line, = ax.plot(x_pts, y_pts, color='white', lw=1.8, animated=True, zorder=5)
            eye_x, eye_y = x_pts[7], y_pts[7]
            eye_patch = Circle((eye_x, eye_y), 1.8, color='white', animated=True, zorder=6)
            pupil_patch = Circle((eye_x, eye_y), 0.8, color='black', animated=True, zorder=7)
            
            energy_pts = ax.scatter(X_bnd, Y_bnd, c=stored_energy, cmap='Wistia',
                                   s=25, zorder=8, animated=True, vmin=0, vmax=0.05)
                                   
            ax.add_patch(eye_patch)
            ax.add_patch(pupil_patch)
            frame_artists.extend([fish_line, eye_patch, pupil_patch, energy_pts])
            
            imgs.append(frame_artists)
            
    ani = animation.ArtistAnimation(fig, imgs, interval=40, blit=True)
    tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    ani.save(tmp.name, fps=25, extra_args=['-vcodec', 'libx264'])
    plt.close(fig)
    return tmp.name

# ============================================================
# GRADIO ИНТЕРФЕЙС
# ============================================================
with gr.Blocks() as demo:
    with gr.Column(elem_classes="border-class"):
        output = gr.Video(label="Fish Hydrodynamics: Physically Accurate Vector Advection", elem_classes="square-video-box")
        with gr.Row():
            with gr.Column():
                gr.Markdown("#### Global Params")
                frames = gr.Slider(minimum=50, maximum=2000, value=600, step=50, label="Frames")
                frequency = gr.Slider(minimum=0.1, maximum=10.0, value=3.0, step=0.1, label="Swimming Freq (Hz)")
            with gr.Column():
                gr.Markdown("#### Hydrodynamics")
                relative_density = gr.Slider(minimum=0.1, maximum=5.0, value=1.0, step=0.1, label="Fish Density")
                tau = gr.Slider(minimum=0.51, maximum=0.95, value=0.58, step=0.01, label="Viscosity (Tau)")
            with gr.Column():
                gr.Markdown("#### Visualization")
                vis_mode = gr.Radio(
                    choices=["Pressure", "Pressure Deviation"],
                    value="Pressure", label="Visualization Mode"
                )
                btn = gr.Button("RUN SIMULATION", variant="primary")
                btn.click(simulate, inputs=[frames, frequency, relative_density, tau, vis_mode], outputs=output)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)