%% py
from builtins import range
import random
import autograd.numpy as np
from autograd import value_and_grad

from scipy.optimize import minimize

# Fluid simulation code based on
# "Real-Time Fluid Dynamics for Games" by Jos Stam
# and taken from examples of autograd
# http://www.intpowertechcorp.com/GDC03.pdf




def project(vx, vy):
    """Project the velocity field to be approximately mass-conserving,
       using a few iterations of Gauss-Seidel."""
    p = np.zeros(vx.shape)
    h = 1.0/vx.shape[0]
    div = -0.5 * h * (np.roll(vx, -1, axis=0) - np.roll(vx, 1, axis=0)
                    + np.roll(vy, -1, axis=1) - np.roll(vy, 1, axis=1))

    for k in range(10):
        p = (div + np.roll(p, 1, axis=0) + np.roll(p, -1, axis=0)
                 + np.roll(p, 1, axis=1) + np.roll(p, -1, axis=1))/4.0

    vx -= 0.5*(np.roll(p, -1, axis=0) - np.roll(p, 1, axis=0))/h
    vy -= 0.5*(np.roll(p, -1, axis=1) - np.roll(p, 1, axis=1))/h
    return vx, vy  

def advect(f, vx, vy):
    """Move field f according to x and y velocities (u and v)
       using an implicit Euler integrator."""
    rows, cols = f.shape
    cell_ys, cell_xs = np.meshgrid(np.arange(rows), np.arange(cols))
    center_xs = (cell_xs - vx).ravel()
    center_ys = (cell_ys - vy).ravel()

    # Compute indices of source cells.
    left_ix = np.floor(center_xs).astype(int)
    top_ix  = np.floor(center_ys).astype(int)
    rw = center_xs - left_ix              # Relative weight of right-hand cells.
    bw = center_ys - top_ix               # Relative weight of bottom cells.
    left_ix  = np.mod(left_ix,     rows)  # Wrap around edges of simulation.
    right_ix = np.mod(left_ix + 1, rows)
    top_ix   = np.mod(top_ix,      cols)
    bot_ix   = np.mod(top_ix  + 1, cols)

    # A linearly-weighted sum of the 4 surrounding cells.
    flat_f = (1 - rw) * ((1 - bw)*f[left_ix,  top_ix] + bw*f[left_ix,  bot_ix]) \
                 + rw * ((1 - bw)*f[right_ix, top_ix] + bw*f[right_ix, bot_ix])
    return np.reshape(flat_f, (rows, cols))

def simulate(vx, vy, smoke, num_time_steps):
    print("Running simulation...", str(random.random()) )
    out = []
    for t in range(num_time_steps):
        
        vx_updated = advect(vx, vx, vy)
        
        vy_updated = advect(vy, vx, vy)
        
        vx, vy = project(vx_updated, vy_updated)
        
        smoke = advect(smoke, vx, vy)
        
        out.append(smoke)
    return smoke, out

simulation_timesteps = 100

init_smoke = np.zeros([16,16]);
target     = np.ones([16,16]);

init_smoke[4:6,5:8] = 1.;

rows, cols = target.shape

init_dx_and_dy = np.zeros((2, rows, cols)).ravel()

def distance_from_target_image(smoke):
    return np.mean((target-smoke)**2)

def convert_param_vector_to_matrices(params):
    vx = np.reshape(params[:(rows*cols)], (rows, cols))
    vy = np.reshape(params[(rows*cols):], (rows, cols))
    return vx, vy

def objective(params):
    init_vx, init_vy = convert_param_vector_to_matrices(params)
    final_smoke, _ = simulate(init_vx, init_vy, init_smoke, simulation_timesteps)
    return distance_from_target_image(final_smoke)

objective_with_grad = value_and_grad(objective)


def callback(params):
    
    init_vx, init_vy = convert_param_vector_to_matrices(params)
    
    simulate(init_vx, init_vy, init_smoke, simulation_timesteps)
    

result = minimize( objective_with_grad, init_dx_and_dy, jac = True, 
    options = {'maxiter': 1, 'disp': True}, callback=callback )

init_vx, init_vy = convert_param_vector_to_matrices(result.x)

_, out = simulate(init_vx, init_vy, init_smoke, simulation_timesteps)

out = np.reshape(out, [-1])

print(out.shape);

TypeError: Cannot read property 'apply' of undefined
    at env.(anonymous function) (/home/mgasmallah/DCP/pyodide/build/pyodide.asm.js:8:9688)
    at _dcsrch_ (wasm-function[80]:17)
    at byn$fpcast-emu$_dcsrch_ (wasm-function[101]:37)
    at _f2py_rout_minpack2_dcsrch (wasm-function[62]:2366)
    at byn$fpcast-emu$_f2py_rout_minpack2_dcsrch (wasm-function[102]:13)
    at _fortran_call (wasm-function[69]:251)
    at byn$fpcast-emu$_fortran_call (wasm-function[109]:10)
    at wasm-function[841]:1420
    at wasm-function[2755]:1110
    at wasm-function[2750]:44565
TypeError: Cannot read property 'apply' of undefined
    at env.(anonymous function) (/home/mgasmallah/DCP/pyodide/build/pyodide.asm.js:8:9688)
    at _dcsrch_ (wasm-function[80]:17)
    at byn$fpcast-emu$_dcsrch_ (wasm-function[101]:37)
    at _f2py_rout_minpack2_dcsrch (wasm-function[62]:2366)
    at byn$fpcast-emu$_f2py_rout_minpack2_dcsrch (wasm-function[102]:13)
    at _fortran_call (wasm-function[69]:251)
    at byn$fpcast-emu$_fortran_call (wasm-function[109]:10)
    at wasm-function[841]:1420
    at wasm-function[2755]:1110
    at wasm-function[2750]:44565