from unittest import TestCase

from phi.flow import *
from phi.math.backend import Backend
from phi.physics._boundaries import Domain, STICKY

BACKENDS = phi.detect_backends()


class ColabNotebookTest(TestCase):

    def test_graph_gradients(self):
        for backend in BACKENDS:
            if backend.supports(Backend.record_gradients):
                with backend:
                    INFLOW_LOCATION = math.tensor([(4., 5), (8., 5), (12., 5), (16., 5)], batch('inflow_loc'), channel('vector'))
                    INFLOW = CenteredGrid(Sphere(center=INFLOW_LOCATION, radius=3), extrapolation.BOUNDARY, x=32, y=40) * 0.6

                    smoke = CenteredGrid(math.zeros(batch(inflow_loc=4)), extrapolation.BOUNDARY, x=32, y=40)
                    velocity = initial_velocity = StaggeredGrid(0, 0, x=32, y=40) * math.ones(batch(inflow_loc=4))

                    with math.record_gradients(velocity.values):
                        for _ in range(3):
                            smoke = advect.mac_cormack(smoke, velocity, dt=1) + INFLOW
                            buoyancy_force = smoke * (0, 0.5) @ velocity
                            velocity = advect.semi_lagrangian(velocity, velocity, dt=1) + buoyancy_force
                            velocity, _ = fluid.make_incompressible(velocity)
                        loss = math.mean(field.l2_loss(smoke - field.stop_gradient(smoke.inflow_loc[-1])))
                        grad = math.gradients(loss, initial_velocity.values)

    def test_functional_gradient(self):
        for backend in BACKENDS:
            if backend.supports(Backend.functional_gradient):
                with backend:
                    INFLOW_LOCATION = math.tensor([(4., 5), (8., 5), (12., 5), (16., 5)], batch('inflow_loc'), channel('vector'))
                    INFLOW = CenteredGrid(Sphere(center=INFLOW_LOCATION, radius=3), extrapolation.BOUNDARY, x=32, y=40) * 0.6

                    def simulate(velocity: StaggeredGrid, smoke: CenteredGrid):
                        for _ in range(3):
                            smoke = advect.mac_cormack(smoke, velocity, dt=1) + INFLOW
                            buoyancy_force = smoke * (0, 0.5) @ velocity
                            velocity = advect.semi_lagrangian(velocity, velocity, dt=1) + buoyancy_force
                            velocity, _ = fluid.make_incompressible(velocity)
                        loss = field.l2_loss(diffuse.explicit(smoke - field.stop_gradient(smoke.inflow_loc[-1]), 1, 1, 10))
                        return loss, smoke, velocity

                    initial_smoke = CenteredGrid(math.zeros(batch(inflow_loc=4)), extrapolation.BOUNDARY, x=32, y=40)
                    initial_velocity = StaggeredGrid(0, 0, x=32, y=40) * math.ones(batch(inflow_loc=4))

                    sim_grad = field.functional_gradient(simulate, wrt=[0], get_output=False)

                    for _ in range(2):
                        velocity_grad, = sim_grad(initial_velocity, initial_smoke)
                        initial_velocity = initial_velocity - 0.01 * velocity_grad

    def test_functional_gradient_jit(self):
        for backend in BACKENDS:
            if backend.supports(Backend.functional_gradient):
                with backend:
                    DOMAIN = Domain(x=32, y=40, boundaries=STICKY, bounds=Box[0:32, 0:40])
                    INFLOW_LOCATION = math.tensor([(4., 5), (8., 5), (12., 5), (16., 5)], batch('inflow_loc'), channel('vector'))
                    INFLOW = DOMAIN.scalar_grid(Sphere(center=INFLOW_LOCATION, radius=3)) * 0.6

                    @math.jit_compile
                    def simulate(velocity: StaggeredGrid, smoke: CenteredGrid):
                        for _ in range(3):
                            smoke = advect.mac_cormack(smoke, velocity, dt=1) + INFLOW
                            buoyancy_force = smoke * (0, 0.5) @ velocity
                            velocity = advect.semi_lagrangian(velocity, velocity, dt=1) + buoyancy_force
                            velocity, _ = fluid.make_incompressible(velocity)
                        loss = field.l2_loss(diffuse.explicit(smoke - field.stop_gradient(smoke.inflow_loc[-1]), 1, 1, 10))
                        return loss, smoke, velocity

                    initial_smoke = DOMAIN.scalar_grid(math.zeros(batch(inflow_loc=4)))
                    initial_velocity = DOMAIN.staggered_grid(0) * math.ones(batch(inflow_loc=4))

                    sim_grad = field.functional_gradient(simulate, wrt=[0], get_output=False)

                    for _ in range(2):
                        velocity_grad, = sim_grad(initial_velocity, initial_smoke)
                        initial_velocity = initial_velocity - 0.01 * velocity_grad

@@ -87,59 +87,8 @@
 kend
-%0Afrom phi.physics._boundaries import Domain, STICKY
 %0A%0ABA
@@ -3196,99 +3196,8 @@
 nd:%0A
-                    DOMAIN = Domain(x=32, y=40, boundaries=STICKY, bounds=Box%5B0:32, 0:40%5D)%0A
     
@@ -3351,31 +3351,25 @@
 NFLOW = 
-DOMAIN.scalar_g
+CenteredG
 rid(Sphe
@@ -3404,16 +3404,52 @@
 adius=3)
+, extrapolation.BOUNDARY, x=32, y=40
 ) * 0.6%0A
@@ -4171,23 +4171,17 @@
 e = 
-DOMAIN.scalar_g
+CenteredG
 rid(
@@ -4211,16 +4211,52 @@
 _loc=4))
+, extrapolation.BOUNDARY, x=32, y=40
 )%0A      
@@ -4292,24 +4292,17 @@
 y = 
-DOMAIN.s
+S
 taggered
 _gri
@@ -4301,15 +4301,29 @@
 ered
-_g
+G
 rid(0
+, 0, x=32, y=40
 ) * 

Update test_colab_fluids_tutorial.py