"""
World / look track, part 3 (module of the import runner): PP_Water, the post-process material that turns distance into water
(DESIGN_UE.md section 7.2, port of v2 GLSL_WATER / FOG_FRAG). Domain Post Process, blendable location "Before DOF".

  absorption = exp(-d * rho * (3.2, 1.05, 0.8))      red dies first
  fog        = 1 - exp(-(d * rho)^1.35)
  water      = lerp(WaterDeep, WaterTop, smoothstep(0.2, 1.0, up)) * (0.03 + 0.97 L) + sun glow + the Floor's magenta glow below
  result     = lerp(scene * absorption, water, fog)
  pixels whose custom stencil is 1 AND whose custom depth is the depth of the visible surface are the dry interior and pass through
  untouched (section 7.4), but only while the camera itself is in dry air (MPC DryCamera = 1, ADescentEnvironment's Dry state): seen
  from the sea, the pool rim, a window or a docked Tender are in the water like everything else. The depth test matters: custom depth/stencil is rendered without the main depth test, so an interior
  object hidden behind the opaque hull would otherwise leave a dry-looking, unfogged rectangle in the hull seen from outside
  (found in the integration pass, m2_exterior).

All numbers come from MPC_Descent, which ADescentEnvironment writes every frame. WaterAmount (0..1) blends the whole effect
(Descent.Look.Set WaterPP).
"""
import importlib.util
import os

import unreal

CODE = """
float3 sceneC = Scene.rgb;
if (DryCam > 0.5 && Stencil.x > 0.5 && CDepth.x <= Depth.x * 1.01 + 5.0)
{
    return sceneC;
}
float3 tw = SvPositionToTranslatedWorld(float4(Parameters.SvPosition.xy, 0.0001, 1.0));
float3 dir = normalize(tw - ResolvedView.TranslatedWorldCameraOrigin);
float cf = max(dot(dir, ResolvedView.ViewForward), 0.05);
float d = min(Depth.x / cf * 0.01, 1500.0);
float3 absr = exp(-d * Rho * float3(3.2, 1.05, 0.8));
float fog = 1.0 - exp(-pow(max(d * Rho, 0.0), 1.35));
float up = dir.z * 0.5 + 0.5;
float3 wc = lerp(Deep, Top, smoothstep(0.2, 1.0, up)) * (0.03 + 0.97 * SurfaceLight);
wc += float3(0.25, 0.62, 0.65) * pow(saturate(dot(dir, normalize(Sun))), 6.0) * SurfaceLight * 0.35;
wc += float3(0.35, 0.06, 0.22) * (1.0 - smoothstep(-1.0, -0.2, dir.z)) * FloorGlow * (1.0 + 2.0 * Pulse);
float3 outc = lerp(sceneC * absr, wc, saturate(fog));
return lerp(sceneC, outc, saturate(Amount));
"""


def run(ctx):
    h = _load(ctx)
    MEL = ctx.mel
    mpc = h.ensure_mpc(ctx)
    mat = ctx.get_or_create(h.MAT_DIR, "PP_Water", unreal.Material, unreal.MaterialFactoryNew())
    MEL.delete_all_material_expressions(mat)
    mat.set_editor_property("material_domain", unreal.MaterialDomain.MD_POST_PROCESS)
    mat.set_editor_property("blendable_location", unreal.BlendableLocation.BL_SCENE_COLOR_BEFORE_DOF)

    def scene_tex(tex_id, x, y):
        n = MEL.create_material_expression(mat, unreal.MaterialExpressionSceneTexture, x, y)
        n.set_editor_property("scene_texture_id", tex_id)
        return n

    scene = scene_tex(unreal.SceneTextureId.PPI_POST_PROCESS_INPUT0, -900, -300)
    depth = scene_tex(unreal.SceneTextureId.PPI_SCENE_DEPTH, -900, -150)
    stencil = scene_tex(unreal.SceneTextureId.PPI_CUSTOM_STENCIL, -900, 0)
    cdepth = scene_tex(unreal.SceneTextureId.PPI_CUSTOM_DEPTH, -900, -75)
    rho = h.mpc_param(ctx, mat, mpc, "FogDensity", -900, 150)
    light = h.mpc_param(ctx, mat, mpc, "SurfaceLight", -900, 250)
    glow = h.mpc_param(ctx, mat, mpc, "FloorGlow", -900, 350)
    pulse = h.mpc_param(ctx, mat, mpc, "Pulse", -900, 450)
    amount = h.mpc_param(ctx, mat, mpc, "WaterAmount", -900, 550)
    drycam = h.mpc_param(ctx, mat, mpc, "DryCamera", -900, 950)
    top = h.mpc_param(ctx, mat, mpc, "WaterTop", -900, 650)
    deep = h.mpc_param(ctx, mat, mpc, "WaterDeep", -900, 750)
    sun = h.mpc_param(ctx, mat, mpc, "SunDir", -900, 850)

    c = h.custom_node(ctx, mat, CODE, ["Scene", "Depth", "Stencil", "CDepth", "DryCam", "Rho", "SurfaceLight", "FloorGlow", "Pulse", "Amount", "Top", "Deep", "Sun"],
                      unreal.CustomMaterialOutputType.CMOT_FLOAT3, -450, 0, description="PP_Water")
    for src, out, name in ((scene, "Color", "Scene"), (depth, "Color", "Depth"), (stencil, "Color", "Stencil"), (cdepth, "Color", "CDepth"), (drycam, "", "DryCam"), (rho, "", "Rho"), (light, "", "SurfaceLight"),
                           (glow, "", "FloorGlow"), (pulse, "", "Pulse"), (amount, "", "Amount"), (top, "", "Top"), (deep, "", "Deep"), (sun, "", "Sun")):
        if not MEL.connect_material_expressions(src, out, c, name):
            ctx.fail("PP_Water: could not connect the input " + name)
    MEL.connect_material_property(c, "", unreal.MaterialProperty.MP_EMISSIVE_COLOR)
    MEL.recompile_material(mat)
    ctx.eal.save_loaded_asset(mat)
    ctx.log("PP_Water ready")


def _load(ctx):
    spec = importlib.util.spec_from_file_location("descent_helpers_look", os.path.join(ctx.project_dir, "Tools", "ue", "_helpers_look.py"))
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module