#!/usr/bin/env python3
"""Actual bounded2D linear grid integration, independent of the six-amplitude matrix."""
import numpy as np


def evolve_grid(stokes,loading,kx,kz,gas_background,dust_background,seed_vector,basis,cells,end_time,step):
    dx=2*np.pi/kx/cells;dz=2*np.pi/kz/cells
    phase=2*np.pi*(np.arange(cells)[:,None]+np.arange(cells)[None,:])/cells
    full_seed=np.concatenate([basis@seed_vector[:2],seed_vector[2:5],[seed_vector[5]]])
    state=np.real(full_seed[:,None,None]*np.exp(1j*phase)[None,:,:])*1e-6
    frequency=np.fft.fftfreq(cells)*cells
    symbol_x=np.sin(2*np.pi*frequency/cells)/dx
    symbol_z=np.sin(2*np.pi*frequency/cells)/dz
    denominator=symbol_x[:,None]**2+symbol_z[None,:]**2
    inverse=np.zeros_like(denominator);np.divide(1,denominator,out=inverse,where=denominator>1e-20)
    relative=dust_background-gas_background
    def derivative(values):
        radial=(np.roll(values,-1,axis=1)-np.roll(values,1,axis=1))/(2*dx)
        vertical=(np.roll(values,-1,axis=2)-np.roll(values,1,axis=2))/(2*dz)
        gas=values[:3];dust=values[3:6];density=values[6]
        gas_rate=-gas_background[0]*radial[:3]+loading/stokes*(dust-gas+relative[:,None,None]*density)
        gas_rate[0]+=2*gas[1];gas_rate[1]-=.5*gas[0]
        transformed=np.fft.fft2(gas_rate,axes=(1,2))
        divergence=symbol_x[:,None]*transformed[0]+symbol_z[None,:]*transformed[2]
        transformed[0]-=symbol_x[:,None]*divergence*inverse
        transformed[2]-=symbol_z[None,:]*divergence*inverse
        gas_rate=np.fft.ifft2(transformed,axes=(1,2)).real
        dust_rate=-dust_background[0]*radial[3:6]-(dust-gas)/stokes
        dust_rate[0]+=2*dust[1];dust_rate[1]-=.5*dust[0]
        density_rate=-dust_background[0]*radial[6]-radial[3]-vertical[5]
        return np.concatenate([gas_rate,dust_rate,density_rate[None,:,:]])
    history=[];time=0
    initial=np.sqrt(np.mean(state[6]**2))
    while time<end_time-1e-12:
        duration=min(step,end_time-time)
        a=derivative(state);b=derivative(state+.5*duration*a);c=derivative(state+.5*duration*b);d=derivative(state+duration*c)
        state+=duration*(a+2*b+2*c+d)/6;time+=duration
        if len(history)==0 or time-history[-1]['time']>=.19 or time>=end_time-1e-12:history.append({'time':time,'density_rms':float(np.sqrt(np.mean(state[6]**2)))})
    growth=np.log(np.sqrt(np.mean(state[6]**2))/initial)/time
    return {'cells_per_wavelength':cells,'end_time':time,'step':step,'measured_growth':float(growth),'relative_mean_density_perturbation':float(np.mean(state[6])),'history':history,'final_relative_density':state[6].tolist() if cells==32 else None}