"""
Definition of Fluid, IncompressibleFlow as well as fluid-related functions.
"""
from typing import Tuple

from phi import math, field
from phi.field import SoftGeometryMask, AngularVelocity, Grid, divergence, spatial_gradient, where, HardGeometryMask, CenteredGrid
from phi.geom import union
from ..math import extrapolation
from ..math._tensors import copy_with
from ..math.extrapolation import combine_sides


def make_incompressible(velocity: Grid,
                        obstacles: tuple or list = (),
                        solve=math.Solve('auto', 1e-5, 1e-5, gradient_solve=math.Solve('auto', 1e-5, 1e-5))) -> Tuple[Grid, CenteredGrid]:
    """
    Projects the given velocity field by solving for the pressure and subtracting its spatial_gradient.
    
    This method is similar to :func:`field.divergence_free()` but differs in how the boundary conditions are specified.

    Args:
        velocity: Vector field sampled on a grid
        obstacles: List of Obstacles to specify boundary conditions inside the domain (Default value = ())
        solve: Parameters for the pressure solve as.

    Returns:
        velocity: divergence-free velocity of type `type(velocity)`
        pressure: solved pressure field, `CenteredGrid`
    """
    assert isinstance(obstacles, (tuple, list)), f"obstacles must be a tuple or list but got {type(obstacles)}"
    input_velocity = velocity
    accessible_extrapolation = _accessible_extrapolation(input_velocity.extrapolation)
    active = CenteredGrid(HardGeometryMask(~union(*[obstacle.geometry for obstacle in obstacles])), resolution=velocity.resolution, bounds=velocity.bounds, extrapolation=extrapolation.NONE)
    accessible = active.with_extrapolation(accessible_extrapolation)
    hard_bcs = field.stagger(accessible, math.minimum, input_velocity.extrapolation, type=type(velocity))
    velocity = apply_boundary_conditions(velocity, obstacles)
    div = divergence(velocity) * active
    if not input_velocity.extrapolation.connects_to_outside:
        assert solve.preprocess_y is None, "fluid.make_incompressible() does not support custom preprocessing"
        solve = copy_with(solve, preprocess_y=lambda y: _balance_divergence(y, active))
    if solve.x0 is None:
        pressure_extrapolation = _pressure_extrapolation(input_velocity.extrapolation)
        solve = copy_with(solve, x0=CenteredGrid(0, resolution=div.resolution, bounds=div.bounds, extrapolation=pressure_extrapolation))
    pressure = math.solve_linear(masked_laplace, f_args=[hard_bcs, active], y=div, solve=solve)
    grad_pressure = field.spatial_gradient(pressure, input_velocity.extrapolation, type=type(velocity)) * hard_bcs
    velocity = velocity - grad_pressure
    return velocity, pressure


@math.jit_compile_linear
def masked_laplace(pressure: CenteredGrid, hard_bcs: Grid, active: CenteredGrid):
    grad = spatial_gradient(pressure, hard_bcs.extrapolation, type=type(hard_bcs))
    grad *= hard_bcs
    div = divergence(grad)
    lap = where(active, div, pressure)
    return lap


def _balance_divergence(div, active):
    return div - active * (field.mean(div) / field.mean(active))


def apply_boundary_conditions(velocity: Grid, obstacles: tuple or list):
    """
    Enforces velocities boundary conditions on a velocity grid.
    Cells inside obstacles will get their velocity from the obstacle movement.
    Cells outside far away will be unaffected.

    Args:
      velocity: Velocity `Grid`.
      obstacles: Obstacles as `tuple` or `list`

    Returns:
        Velocity of same type as `velocity`
    """
    # velocity = field.bake_extrapolation(velocity)  # TODO we should bake only for divergence but keep correct extrapolation for velocity. However, obstacles should override extrapolation.
    for obstacle in obstacles:
        obs_mask = SoftGeometryMask(obstacle.geometry, balance=1) @ velocity
        if obstacle.is_stationary:
            velocity = (1 - obs_mask) * velocity
        else:
            angular_velocity = AngularVelocity(location=obstacle.geometry.center, strength=obstacle.angular_velocity, falloff=None) @ velocity
            velocity = (1 - obs_mask) * velocity + obs_mask * (angular_velocity + obstacle.velocity)
    return velocity


def _pressure_extrapolation(vext: math.Extrapolation):
    if vext == extrapolation.PERIODIC:
        return extrapolation.PERIODIC
    elif vext == extrapolation.BOUNDARY:
        return extrapolation.ZERO
    elif isinstance(vext, extrapolation.ConstantExtrapolation):
        return extrapolation.BOUNDARY
    elif isinstance(vext, extrapolation._MixedExtrapolation):
        return combine_sides(**{dim: (_pressure_extrapolation(lo), _pressure_extrapolation(hi)) for dim, (lo, hi) in vext.ext.items()})
    else:
        raise ValueError(f"Unsupported extrapolation: {type(vext)}")


def _accessible_extrapolation(vext: math.Extrapolation):
    if vext == extrapolation.PERIODIC:
        return extrapolation.PERIODIC
    elif vext == extrapolation.BOUNDARY:
        return extrapolation.ONE
    elif isinstance(vext, extrapolation.ConstantExtrapolation):
        return extrapolation.ZERO
    elif isinstance(vext, extrapolation._MixedExtrapolation):
        return combine_sides(**{dim: (_accessible_extrapolation(lo), _accessible_extrapolation(hi)) for dim, (lo, hi) in vext.ext.items()})
    else:
        raise ValueError(f"Unsupported extrapolation: {type(vext)}")

@@ -285,24 +285,59 @@
 mport union%0A
+from ..field._grid import GridType%0A
 from ..math 
@@ -353,24 +353,24 @@
 trapolation%0A
-
 from ..math.
@@ -482,16 +482,20 @@
 ty: Grid
+Type
 ,%0A      
@@ -543,16 +543,16 @@
 t = (),%0A
-
         
@@ -665,16 +665,20 @@
 ple%5BGrid
+Type
 , Center

Fix type hints in make_incompressible()