from pathlib import Path
import os
from collections import defaultdict

from IPython.display import HTML, Image
from phi.torch.flow import *
from tqdm.notebook import trange
import matplotlib.pyplot as plt
import torch
import numpy as np
import warnings

warnings.filterwarnings('ignore')

ROOT = Path("./assets/img/")

if not os.path.exists(ROOT):
    os.makedirs(ROOT)

def clamp_field(f, lo=0.0, hi=1.0):
    v = f.values                     # backend tensor
    v = math.maximum(lo, math.minimum(hi, v))
    return f.with_values(v)

def smoke_plume():

    domain = Box(x=100, y=100)
    
    inflow = Sphere(x=50, y=9.5, radius=5)
    inflow_rate = 0.2

    def step(v, s, p, dt):
        s = advect.mac_cormack(s, v, dt) + inflow_rate * resample(inflow, to=s, soft=True)
        
        buoyancy = resample(s * (0, 0.1), to=v)
        v = advect.semi_lagrangian(v, v, dt) + buoyancy * dt
        
        v, p = fluid.make_incompressible(v, (), Solve('auto', 1e-3, x0=p))
        return v, s, p

    v0 = StaggeredGrid(0, 0, domain, x=64, y=64)
    smoke0 = CenteredGrid(0, ZERO_GRADIENT, domain, x=200, y=200)

    v_trj, s_trj, p_trj = iterate(step, batch(time=150), v0, smoke0, None, dt=.5, substeps=3, range=trange)

    anim = plot(s_trj, animate='time', frame_time=80,)

    gif_path = ROOT / "smoke_plume.gif"  
    anim.save(gif_path, writer="pillow", fps=30)
    return Image(url=gif_path)

smoke_plume()

def smoke_phi():
    
    domain = Box(x=100, y=100)
    
    def step(v, s, p, dt):
        s = advect.mac_cormack(s, v, dt)
        
        buoyancy = resample(s * (0, 0.1), to=v)
        v = advect.semi_lagrangian(v, v, dt) + buoyancy * dt
        
        v, p = fluid.make_incompressible(v, (), Solve('auto', 1e-3, x0=p))
        return v, s, p

    v0 = StaggeredGrid(0, 0, domain, x=64, y=64)
    smoke0 = CenteredGrid(0, ZERO_GRADIENT, domain, x=256, y=256)

    center = vec(x=50, y=20)
    R_out, R_in = 15., 10.
    
    outer = Sphere(center, radius=R_out)
    inner = Sphere(center, radius=R_in)
    
    outer_mask = resample(outer, to=smoke0, soft=True)
    inner_mask = resample(inner, to=smoke0, soft=True)
    
    ring_mask = outer_mask * (1.0 - inner_mask) 
    
    smoke0 = smoke0 + ring_mask

    v_trj, s_trj, p_trj = iterate(step, batch(time=150), v0, smoke0, None, dt=.5, range=trange, substeps=3)

    anim = plot(s_trj, animate='time', frame_time=80)
    plt.close()

    gif_path = ROOT / "smoke_phi.gif"  
    anim.save(gif_path, writer="pillow", fps=30)
    return Image(url=gif_path)

smoke_phi()

def smoke_puff():
    
    domain = Box(x=100, y=100)

    inflow = Box(x=(0, 10), y=(40, 60))
    inflow_rate = 0.1
    in_velocity = vec(x=1., y=0.)
    
    def step(v, s, p, dt):
        s = advect.mac_cormack(s, v, dt)

        s = s + inflow_rate*resample(inflow, to=s, soft=True)
        
        buoyancy = resample(s * (0, 0.1), to=v)

        v = v + resample(inflow, to=v, soft=True)*in_velocity
        v = advect.semi_lagrangian(v, v, dt) + buoyancy * dt
        
        v, p = fluid.make_incompressible(v, (), Solve('auto', 1e-2, x0=p))
        return v, s, p

    v0 = StaggeredGrid(0, 0, domain, x=64, y=64)
    smoke0 = CenteredGrid(0, ZERO_GRADIENT, domain, x=256, y=256)

    v_trj, s_trj, p_trj = iterate(step, batch(time=150), v0, smoke0, None, dt=.5, range=trange, substeps=3)

    anim = plot(s_trj, animate='time', frame_time=80)
    plt.close()

    gif_path = ROOT / "smoke_puff.gif"  
    anim.save(gif_path, writer="pillow", fps=30)
    return Image(url=gif_path)

smoke_puff()

def smoke_obstacle():
    
    domain = Box(x=100, y=100)

    inflow = Box(x=(0, 10), y=(40, 60))
    inflow_rate = 0.1
    in_velocity = vec(x=1., y=0.)

    geometry = Box(x=(20, 30), y=(45, 55))
    
    def step(v, s, p, dt):
        s = advect.mac_cormack(s, v, dt)

        s = s + inflow_rate*resample(inflow, to=s, soft=True)
        
        buoyancy = resample(s * (0, 0.1), to=v)

        v = v + resample(inflow, to=v, soft=True)*in_velocity
        v = advect.semi_lagrangian(v, v, dt) + buoyancy * dt
        
        v, p = fluid.make_incompressible(v, geometry, Solve('auto', 1e-2, x0=p))
        return v, s, p

    v0 = StaggeredGrid(0, 0, domain, x=64, y=64)
    smoke0 = CenteredGrid(0, ZERO_GRADIENT, domain, x=256, y=256)

    v_trj, s_trj, p_trj = iterate(step, batch(time=150), v0, smoke0, None, dt=.5, range=trange, substeps=3)
    
    anim = plot(s_trj, animate='time', frame_time=80,)
    plt.close()
    
    gif_path = ROOT / "smoke_obstacle.gif"  
    anim.save(gif_path, writer="pillow", fps=30)

    return Image(url=gif_path)

smoke_obstacle()

def smoke_collide():
    
    domain = Box(x=100, y=100)

    inflow_left = Box(x=(0, 10), y=(40, 60))
    inflow_right = Box(x=(90, 100), y=(40, 60))
    inflow_rate = 0.1
    in_velocity_left = vec(x=1., y=0.)
    in_velocity_right = vec(x=-1., y=0.)
    
    def step(v, s, p, dt):
        s = advect.mac_cormack(s, v, dt)

        s = s + inflow_rate*resample(inflow_left, to=s, soft=True) + \
            inflow_rate*resample(inflow_right, to=s, soft=True)
        
        buoyancy = resample(s * (0, 0.1), to=v)

        v = v + resample(inflow_left, to=v, soft=True)*in_velocity_left + \
            resample(inflow_right, to=v, soft=True)*in_velocity_right
        
        v = advect.semi_lagrangian(v, v, dt) + buoyancy * dt
        
        v, p = fluid.make_incompressible(v, (), Solve('auto', 1e-2, x0=p))
        return v, s, p

    v0 = StaggeredGrid(0, 0, domain, x=64, y=64)
    smoke0 = CenteredGrid(0, ZERO_GRADIENT, domain, x=256, y=256)

    v_trj, s_trj, p_trj = iterate(step, batch(time=150), v0, smoke0, None, dt=.5, range=trange, substeps=3)
    
    anim = plot(s_trj, animate='time', frame_time=80,)
    plt.close()
    
    gif_path = ROOT / "smoke_collide.gif"  
    anim.save(gif_path, writer="pillow", fps=30)

    return Image(url=gif_path)

smoke_collide()