import sys
import gmsh
import shutil
import numpy as np
from migflow import fluid, advdiff, scontact, time_integration
from gen_geometry import *
from tools import *

temperature = np.array([100, 150, 200, 250, 300, 350, 400, 450, 500], dtype=float)

np.random.seed(0)
p2p1 = False
v_element = b"triangle_p2" if p2p1 else b"triangle_p1"
thickness_condition = False
T0 = 0

# ========= Numerical calibration =========
volume_ratio = 0.5

g = np.array([0.0, -9.81])  # [m/s^2]
# ========= Geometry =========
height = 20e-3  # [m]
p_height = height - 4e-3  # [m]
width = 7e-3  # [m]
depth = 3e-3  # [m]
origin = np.array([-width / 2, 0])
# ========= Air properties =========
nitrogen_density = 1.204  # [kg/m**3]
nitrogen_dynamic_viscosity = 2.0e-5  # [kg/(ms)]
nitrogen_kinematic_viscosity = nitrogen_dynamic_viscosity / nitrogen_density  # [m**2/s]
nitrogen_heat_capacity = 1040  # [J/(kg K)]
nitrogen_conductivity = 0.026  # [W/(m K)]
nitrogen_thermal_expansion = 2.0e-3  # [1/K]
# ========= Grains properties =========
particle_density = volume_ratio * 3200  # [kg/m**3]
particle_heat_capacity = 290  # [J/(kg K)]
particle_conductivity = 234  # [W/(m K)]
# particle_diameter = 28e-6                                                       # [m]
particle_diameter = 55e-6  # [m]
particle_radii = particle_diameter / 2  # [m]
young_modulus = 1e11  # [Pa]
poisson_ratio = 0.35  # [\]
mu = 0.4  # [\]
T_grains = temperature[int(sys.argv[1])] + T0 if len(sys.argv) > 1 else 400 + T0

# ========= Droplet properties =========
droplet_radii = 1.4e-3  # [m]
droplet_density = 958  # [kg/m**3]
droplet_mass = droplet_density * np.pi * droplet_radii**2  # [kg/m]
droplet_heat_capacity = 3768  # [J/(kg K)]
droplet_conductivity = 0.677  # [W/(m K)]
droplet_volume_2d = np.pi * droplet_radii**2  # [m**2]
droplet_dynamic_viscosity = 2.8e-4  # [kg/(ms)]
vapor_thickness = (
    0.1 * droplet_radii
)  # 1e-4                                                          # [m]
latent_heat = 2.26e6  # [J/kg]
evaporation_velocity = 0.0  # [m/s]
droplet_velocity = np.array([0.0, -0.0])  # [m/s]
mass_rate = 0.0  # [kg/(m**2 s)]
T_vap = 100 + T0
# ========= Numerical parameters =========
t = 0  # [s]
dt = 5e-5  # [s]
tEnd = 10  # [s]
i = 0
outf = 50
n_remeshing = 100


print(f"================ STARTING SIMULATION ================")
print(f"................ {T_grains=:g} ")
print(f"................ {droplet_radii=:g} ")
print(f"................ {particle_diameter=:g} ")
print(f"................ {particle_density=:g} ")
print(f"................ {t=:g} ")
# ========= Particle Problem =========
p = scontact.ParticleProblem(2)
p.set_friction_coefficient(mu)
gen_box(p, width, height, origin)
gen_particles(p, width, p_height, origin, particle_radii, particle_density, int(1e5))

droplet_initial_height = np.minimum(
    np.max(p.position()[:, 1]) - origin[1] + particle_radii + 1.2 * droplet_radii,
    height - origin[1] - particle_radii - 1.05 * droplet_radii,
)
droplet_position = np.array([0, droplet_initial_height])
# droplet_position = np.array([0, 3*p_height/4])

body_droplet, particles_droplet = get_particles_inside_droplet(
    p, 1.05 * droplet_radii, droplet_position
)
p.remove_bodies_flag(~body_droplet)

# ========= Mesh Generation =========
mesh_size = 40 * particle_diameter  # 60
gen_mesh(width, height, origin)
gen_mesh_cut(width, height, droplet_position, droplet_radii, origin)
get_physical_entities(width, height, droplet_position, droplet_radii, origin)
gen_mesh_callback(droplet_position, droplet_radii, mesh_size)

# ========= Fluid Problem =========
volume_drag = 0 * 12 * nitrogen_dynamic_viscosity / (depth**2)
f = fluid.FluidProblem(
    2,
    g,
    nitrogen_dynamic_viscosity,
    nitrogen_density,
    alpha=nitrogen_thermal_expansion,
    p2p1=p2p1,
    volume_drag=volume_drag,
)
f.load_msh(None)


def droplet_velocity_x_callback(x):
    return (
        evaporation_velocity * (x[..., 0] - droplet_position[0]) / droplet_radii
        + droplet_velocity[0]
    )


def droplet_velocity_y_callback(x):
    return (
        evaporation_velocity * (x[..., 1] - droplet_position[1]) / droplet_radii
        + droplet_velocity[1]
    )


f.set_open_boundary("Top", pressure=0)
f.set_wall_boundary("Bottom", velocity=[0, 0])
f.set_wall_boundary("Left", velocity=[0, 0])
f.set_wall_boundary("Right", velocity=[0, 0])
f.set_open_boundary(
    "Droplet", velocity=[droplet_velocity_x_callback, droplet_velocity_y_callback]
)

# ========= Heat Problem =========
c = advdiff.AdvDiffProblem(
    2, 0, nitrogen_conductivity, nitrogen_heat_capacity, velocity_element=v_element
)
c.load_msh(None)
c.set_mu_rho(nitrogen_kinematic_viscosity, nitrogen_density)
c.set_weak_boundary("Top", "Neumann")
c.set_weak_boundary("Bottom", "Neumann")
c.set_weak_boundary("Left", "Neumann")
c.set_weak_boundary("Right", "Neumann")
c.set_weak_boundary("Droplet", "Dirichlet", solution=T_vap)

# ========= Initial Conditions =========
droplet_nodes = np.unique(f.mesh_boundaries()[b"Droplet"])
particle_temperature = np.full(p.n_particles(), T_grains, dtype=np.float64)
c.solution()[:] = T_vap
c.solution()[droplet_nodes] = T_vap

# ========= FEM-DEM coupling =========
f.temperature()[:] = c.solution()[:, None]
c.velocity()[:] = f.velocity().copy()
particle_mass = particle_density * np.pi * p.r()[:, 0] ** 2
f.set_particles(
    volume_ratio * p.delassus(),
    volume_ratio * p.volume(),
    p.position(),
    p.velocity(),
    p.omega(),
    p.contact_forces(),
)
c.set_particles(
    volume_ratio * p.volume(),
    p.position(),
    p.velocity(),
    particle_heat_capacity * particle_mass,
    particle_conductivity,
    particle_temperature,
)

# ========= Initial output =========
outputdir = f"output-article/{droplet_radii=:1.6f}/{T_grains=:1.0f}/{particle_diameter=:1.6f}/{particle_density=:1.2f}"
shutil.rmtree(outputdir, ignore_errors=True)
f.write_mig(outputdir, t)
c.write_mig(outputdir, t)
p.write_mig(outputdir, t, extra_fields={"temperature": particle_temperature})
droplet_write_mig(
    outputdir, droplet_position, droplet_velocity, droplet_mass, droplet_radii
)
with open(f"{outputdir}/script.py", "w") as file:
    file.write(open(__file__, "r").read())

# ========= Time integration =========
while t < tEnd:
    print(f"============= {i=} : {t:1.4f}/{tEnd} =============")
    b_outside, p_outside = get_particles_outside_domain(p, origin, height, width)
    p.remove_bodies_flag(~b_outside)

    particle_temperature = particle_temperature[~p_outside]
    particle_mass = particle_density * np.pi * p.r()[:, 0] ** 2
    particle_volume = volume_ratio * p.volume().reshape(-1)

    _, p_inside_droplet = get_particles_inside_droplet(
        p, droplet_radii, droplet_position
    )
    # ========= FEM-DEM coupling =========
    f.temperature()[:] = c.solution()[:, None]
    c.velocity()[:] = f.velocity().copy()
    f.set_particles(
        volume_ratio * p.delassus(),
        particle_volume,
        p.position(),
        p.velocity(),
        p.omega(),
        p.contact_forces(),
    )
    c.set_particles(
        volume_ratio * p.volume(),
        p.position(),
        p.velocity(),
        particle_heat_capacity * particle_mass,
        particle_conductivity,
        particle_temperature,
    )
    alpha = np.clip(f.particle_volume_intersected() / particle_volume, 0, 1)

    p_volume_inside_droplet = np.sum((1 - alpha) * particle_volume)
    droplet_volume = droplet_mass / droplet_density + p_volume_inside_droplet

    Re, Cd_Re, Nu = get_dimensional_number(
        1 - alpha,
        particle_volume,
        p.velocity(),
        droplet_volume,
        droplet_density,
        droplet_dynamic_viscosity,
        droplet_heat_capacity,
        droplet_conductivity,
        droplet_velocity,
    )
    # ========= Energy conservation =========
    c.implicit_euler(dt)

    q_droplet_particles = get_fluid_particle_heat_transfer(
        droplet_conductivity,
        T_vap,
        Nu,
        particle_heat_capacity,
        particle_mass,
        particle_temperature,
        dt,
    )
    q_air_particles = c.particles_heat()[:, 0]
    q_particles_particles = get_particle_particle_heat_transfer(
        p, particle_conductivity, young_modulus, poisson_ratio, particle_temperature
    )
    particle_temperature += (
        (
            q_air_particles * alpha
            + (1 - alpha) * q_droplet_particles
            + q_particles_particles
        )
        * dt
        / (particle_heat_capacity * particle_mass)
    )

    # ========= Evaporation =========
    q_droplet_nitrogen, x_bnd, l_bnd = get_flux(c, b"Droplet")
    w = get_weights_to_edges(
        p.position(), np.mean(x_bnd, axis=1)[:, :2], vapor_thickness
    )
    q_droplet_particles_at_bnd = -np.einsum(
        "pb, p -> b", w, (1 - alpha) * q_droplet_particles
    )
    q_droplet = q_droplet_particles_at_bnd + q_droplet_nitrogen
    print(
        f"{np.sum(q_droplet) = :g} ------- {np.sum(q_droplet_particles_at_bnd) = :g} ------- {np.sum(q_droplet_nitrogen) = :g}"
    )
    evaporation_velocity = get_evaporation_velocity(
        q_droplet, latent_heat, nitrogen_density, l_bnd
    )
    mass_rate = np.sum(q_droplet / latent_heat)
    droplet_mass = droplet_mass - mass_rate * dt

    # ========= Momentum conservation =========
    f.implicit_euler(dt)
    force_droplet_particles = (
        get_fluid_particle_momentum_transfer(
            droplet_dynamic_viscosity,
            droplet_velocity,
            particle_mass,
            p.velocity(),
            Cd_Re,
            dt,
        )
        + g[None, :] * droplet_density * particle_volume[:, None]
    )
    # force_droplet_particles = get_fluid_particle_momentum_transfer(droplet_dynamic_viscosity, droplet_velocity, particle_mass, particle_velocity_pred, Cd_Re, dt) \
    force_air_particles = f.compute_node_force(dt)
    drag = (
        alpha[:, None] * force_air_particles
        + (1 - alpha[:, None]) * force_droplet_particles
    )

    # ========= Particle Resolution =========
    time_integration._advance_particles(
        p,
        particle_mass[:, None] * g[None, :] + drag,
        dt,
        min_nsub=1,
        contact_tol=5e-4 * particle_radii,
        max_nsub=5,
    )

    # ========= Droplet velocity =========
    f_droplet_nitrogen, x_bnd, l_bnd = get_forces(f, b"Droplet")
    # f_droplet_nitrogen_bis = np.sum(f.boundary_force_by_contributions("Droplet").reshape(-1,2), axis=0)
    # f_droplet_nitrogen_by_contributions = f.boundary_force_by_contributions("Droplet").reshape(-1,2)
    # f_dp_dropplet_nitrogen = f_droplet_nitrogen_by_contributions[0]
    # f_viscous_nitrogen = f_droplet_nitrogen_by_contributions[1]
    f_bnd_droplet = -np.sum(f_droplet_nitrogen, axis=0)
    f_buoyancy = droplet_mass * g
    f_particle_droplet = -np.sum((1 - alpha[:, None]) * force_droplet_particles, axis=0)
    f_wall = -volume_drag * droplet_velocity * droplet_volume

    resulting_force = f_bnd_droplet + f_buoyancy + f_particle_droplet + f_wall

    droplet_velocity += resulting_force * dt / droplet_mass
    print(f"    resulting_force        = {resulting_force[1]:g}")
    print(f"    droplet_velocity  = ({droplet_velocity[0]:g}, {droplet_velocity[1]:g})")

    # ========= Mesh regularisation =========
    droplet_nodes = np.unique(f.mesh_boundaries()[b"Droplet"])
    bnd_nodes = np.setdiff1d(
        np.unique(np.concatenate(list(f.mesh_boundaries().values()))), droplet_nodes
    )
    p_volume_inside_droplet = volume_ratio * np.sum(
        np.pi * p.r()[p_inside_droplet, 0] ** 2
    )

    new_radius = np.sqrt(droplet_volume / np.pi)
    front_velocity = (new_radius - droplet_radii) / dt
    droplet_radii = new_radius
    print(f"........ compacity = {100*p_volume_inside_droplet/droplet_volume:3.1f}%")
    print(f"........ {front_velocity = }")
    xold = f.coordinates()[:, :2].copy()
    coord = move_mesh(
        f,
        droplet_position,
        droplet_velocity,
        droplet_radii,
        droplet_nodes,
        front_velocity,
        dt,
        fixed_nodes=bnd_nodes,
    )
    coord = np.insert(coord, 2, 0, axis=1)
    f.set_coordinates(coord)
    c.set_coordinates(coord)
    f.mesh_velocity()[:, :] = (f.coordinates()[:, :2] - xold[:, :2]) / dt
    c.u_mesh()[:] = f.mesh_velocity().copy()
    f.update_node_volume()
    c.update_node_volume()

    droplet_position += droplet_velocity * dt

    # ========= Output =========
    t += dt
    i += 1
    if i % outf == 0:
        print("================= Writing output =================")
        f.write_mig(outputdir, t)
        c.write_mig(outputdir, t)
        p.write_mig(outputdir, t, extra_fields={"temperature": particle_temperature})
        droplet_write_mig(
            outputdir, droplet_position, droplet_velocity, droplet_mass, droplet_radii
        )

    if i % n_remeshing == 0:
        boundaries = f.mesh_boundaries()
        gmsh.clear()
        generate_mesh2d(coord, boundaries, mesh_size)
        f.adapt_mesh()
        c.adapt_mesh()

    exit(0)