# Copyright (C) 2026, Princeton University.
# This source code is licensed under the BSD 3-Clause license found in the LICENSE file in the root directory of this source tree.

# Authors:
# - Yiming Zuo: original Infinigen v1 nodegroup (https://github.com/princeton-vl/infinigen/blob/05a09759fe9478595a3323ec2d6e26ce3513223f/infinigen/assets/materials/ceramic/concrete.py)
# - Alexander Raistrick: transpile to procfunc/v2
# Acknowledgement: This file draws inspiration https://www.youtube.com/watch?v=XDqRa0ExDqs by Ryan King Art

import numpy as np
import procfunc as pf
from procfunc.nodes import types as t

__all__ = [
    "concrete",
    "concrete_rand",
]


@pf.nodes.node_function
def _crack(
    vector: pf.ProcNode[pf.Vector],
    seed: t.SocketOrVal[float] = 0.0,
    amount: t.SocketOrVal[float] = 0.0,
    scale: t.SocketOrVal[float] = 0.0,
    snake_crack: t.SocketOrVal[float] = 0.0,
) -> pf.ProcNode[pf.Color]:
    # Code generated by procfunc v0.2.0

    fac_7 = pf.nodes.texture.noise(
        vector=vector,
        w=seed,
        scale=scale,
        detail=14.0,
        roughness=0.8706,
        noise_dimensions="4D",
        normalize=False,
    )
    noise_vector = pf.nodes.texture.noise(
        vector=vector,
        w=seed,
        scale=scale * 0.5,
        detail=15.0,
        noise_dimensions="4D",
    )
    fac_value_2 = pf.nodes.texture.voronoi_distance(
        vector=noise_vector.fac.astype(dtype=pf.Vector),
        scale=1.2,
    )
    fac_6 = pf.nodes.math.map_range(
        value=fac_value_2,
        from_max=0.02,
        to_min=2.0,
        to_max=0.0,
    )
    fac_5 = pf.nodes.color.mix_rgb(
        factor=snake_crack,
        a=fac_7.fac.astype(dtype=pf.Color),
        b=fac_6.astype(dtype=pf.Color),
        blend_type="ADD",
    )
    fac_value_1 = pf.nodes.texture.noise(
        vector=vector,
        w=seed,
        scale=scale * 0.6,
        detail=14.0,
        noise_dimensions="4D",
        normalize=False,
    )
    fac_4 = pf.nodes.math.map_range(value=amount, to_min=1.0, to_max=-0.5)
    fac_3 = pf.nodes.math.map_range(
        value=fac_value_1.fac,
        from_min=fac_4,
        from_max=fac_4 + 0.1,
    )
    fac_2 = pf.nodes.color.mix_rgb(
        factor=1.0,
        a=fac_5,
        b=fac_3.astype(dtype=pf.Color),
        blend_type="DARKEN",
    )
    fac_value = pf.nodes.texture.noise(
        vector=vector,
        w=seed,
        scale=scale * 0.3,
        detail=14.0,
        noise_dimensions="4D",
        normalize=False,
    )
    fac_1 = pf.nodes.math.map_range(
        value=fac_value.fac,
        from_min=fac_4,
        from_max=fac_4 + 0.1,
    )
    fac = pf.nodes.color.mix_rgb(
        factor=1.0,
        a=fac_2,
        b=fac_1.astype(dtype=pf.Color),
        blend_type="DARKEN",
    )
    result = pf.nodes.color.color_ramp(
        fac=fac.astype(dtype=float),
        interpolation="LINEAR",
        points=[(0.0, (0.0, 0.0, 0.0, 1.0)), (1.0, (1.0, 1.0, 1.0, 1.0))],
    )
    return result.color


@pf.nodes.node_function
def concrete(
    vector: pf.ProcNode[pf.Vector],
    base_color: t.SocketOrVal[pf.Color] = pf.Color((0.636, 0.633, 0.604)),
    scale: t.SocketOrVal[float] = 1.0,
    seed: t.SocketOrVal[float] = -560.1302,
    roughness: t.SocketOrVal[float] = 0.8433,
    crack_amount: t.SocketOrVal[float] = 0.6801,
    crack_scale: t.SocketOrVal[float] = 2.7978,
    snake_crack: t.SocketOrVal[float] = 0.4541,
    color_variation: t.SocketOrVal[float] = 0.2,
    crack_darkness: t.SocketOrVal[float] = 0.2,
    surface_bump_strength: t.SocketOrVal[float] = 0.3,
    displacement_depth: t.SocketOrVal[float] = 0.005,
    micro_grain: t.SocketOrVal[float] = 0.0,
    micro_grain_scale: t.SocketOrVal[float] = 200.0,
) -> pf.Material:
    # Code generated by procfunc v0.2.0

    crack_result_scale = scale * 10.0
    crack_result = _crack(
        vector=vector,
        seed=seed,
        amount=crack_amount,
        scale=crack_result_scale * crack_scale,
        snake_crack=snake_crack,
    )

    base_color_factor = pf.nodes.math.map_range(crack_result.astype(dtype=float))
    base_color_value = pf.nodes.texture.noise(
        vector=vector,
        w=seed,
        scale=scale * 2.0,
        detail=15.0,
        noise_dimensions="4D",
    )

    surface_factor_1 = pf.nodes.texture.noise(
        vector=vector,
        w=seed,
        scale=scale * 5.0,
        detail=14.0,
        roughness=0.3333,
        lacunarity=3.0,
        noise_dimensions="4D",
        normalize=False,
    )
    surface_factor = pf.nodes.color.mix_rgb(
        factor=0.5,
        a=surface_factor_1.fac.astype(dtype=pf.Color),
        b=pf.Color((0.5, 0.5, 0.5)),
    )

    color_color_1 = pf.nodes.color.hue_saturation(fac=1.0, value=0.6, color=base_color)

    displacement_factor_2 = pf.nodes.texture.noise(
        vector=vector,
        w=seed,
        scale=scale * 20.0,
        detail=15.0,
        distortion=0.2,
        noise_dimensions="4D",
    )
    displacement_factor_1 = pf.nodes.texture.noise(
        vector=vector,
        w=seed,
        scale=scale * 20.0,
        detail=14.0,
        roughness=0.8706,
        noise_dimensions="4D",
        normalize=False,
    )
    displacement_factor = pf.nodes.color.mix_rgb(
        factor=0.5,
        a=displacement_factor_2.fac.astype(dtype=pf.Color),
        b=displacement_factor_1.fac.astype(dtype=pf.Color),
    )
    displacement_6 = pf.nodes.color.hue_saturation(fac=1.0, value=1.4, color=base_color)
    displacement_5 = pf.nodes.math.mix(
        factor=displacement_factor.astype(dtype=float),
        a=base_color,
        b=displacement_6,
    )

    base_color_color = pf.nodes.color.mix_rgb(
        factor=surface_factor.astype(dtype=float),
        a=color_color_1,
        b=displacement_5,
    )
    base_color_2 = pf.nodes.color.hue_saturation(
        value=base_color_value.fac,
        fac=color_variation,
        color=base_color_color,
    )
    base_color_1 = pf.nodes.color.hue_saturation(
        fac=1.0, value=crack_darkness, color=base_color
    )

    surface_base_color = pf.nodes.color.mix_rgb(
        factor=base_color_factor, a=base_color_2, b=base_color_1
    )

    # Micro grain texture (fine sand/aggregate with sharp cell boundaries)
    grain_voronoi = pf.nodes.texture.voronoi(
        vector=vector,
        w=seed,
        scale=micro_grain_scale,
        voronoi_dimensions="4D",
    )
    grain_color = pf.nodes.color.mix_rgb(
        factor=grain_voronoi.color,
        a=base_color,
        b=pf.Color((0.5, 0.5, 0.5)),
        blend_type="SOFT_LIGHT",
    )
    surface_base_color = pf.nodes.color.mix_rgb(
        factor=micro_grain,
        a=surface_base_color,
        b=grain_color,
    )

    surface = pf.nodes.shader.principled_bsdf(
        base_color=surface_base_color,
        roughness=roughness,
        subsurface_scale=1.0,
    )

    crack_result_1 = _crack(
        vector=vector,
        seed=seed,
        amount=crack_amount * 0.6,
        scale=(crack_result_scale * crack_scale) * 5.0,
        snake_crack=0.3,
    )

    displacement_4 = crack_result_1.astype(dtype=float) * (roughness * 1.0)
    displacement_3 = (roughness * 1.0) * crack_result.astype(dtype=float)
    displacement_1 = surface_bump_strength * roughness
    displacement_height = (
        displacement_4 + displacement_3
    ) + displacement_1 * displacement_5.astype(dtype=float)

    displacement = pf.nodes.shader.displacement(
        height=displacement_height,
        midlevel=0.0,
        scale=displacement_depth * -1.0,
        normal=(0.0, 0.0, 0.0),
    )

    return pf.Material(
        surface=surface,
        displacement=displacement,
    )


def concrete_rand(
    rng: pf.RNG,
    vector: pf.ProcNode[pf.Vector],
    base_color: t.SocketOrVal[pf.Color] | None = None,
    roughness: t.SocketOrVal[float] | None = None,
    crack_amount: t.SocketOrVal[float] | None = None,
    crack_scale: t.SocketOrVal[float] | None = None,
    snake_crack: t.SocketOrVal[float] | None = None,
    color_variation: t.SocketOrVal[float] | None = None,
    crack_darkness: t.SocketOrVal[float] | None = None,
    surface_bump_strength: t.SocketOrVal[float] | None = None,
    displacement_depth: t.SocketOrVal[float] | None = None,
    micro_grain: t.SocketOrVal[float] | None = None,
    micro_grain_scale: t.SocketOrVal[float] | None = None,
) -> pf.Material:
    if base_color is None:
        hue_sat = pf.random.uniform(
            rng,
            np.array([0.0, 0.0]),
            np.array([1.0, 0.1]),
        )
        hue, sat = hue_sat[0], hue_sat[1]
        val = pf.random.clip_gaussian(rng, 0.5, 0.2, 0.2, 0.95)
        base_color = pf.color.hsv_color(hue=hue, saturation=sat, value=val)

    if roughness is None:
        roughness = pf.random.uniform(rng, 0.5, 1.0)
    if crack_amount is None:
        crack_amount = pf.random.uniform(rng, 0.2, 0.8)
    if crack_scale is None:
        crack_scale = pf.random.uniform(rng, 1.0, 3.0)
    if snake_crack is None:
        snake_crack = pf.random.uniform(rng, 0.0, 1.0)
    if color_variation is None:
        color_variation = pf.random.uniform(rng, 0.05, 0.4)
    if crack_darkness is None:
        crack_darkness = pf.random.uniform(rng, 0.1, 0.4)
    if surface_bump_strength is None:
        surface_bump_strength = pf.random.uniform(rng, 0.1, 0.5)
    if displacement_depth is None:
        displacement_depth = pf.random.uniform(rng, 0.002, 0.01)

    if micro_grain is None:
        micro_grain = pf.control.choice(
            rng, [(pf.random.uniform(rng, 0.2, 0.6), 1.0), (0.0, 1.0)]
        )
    if micro_grain_scale is None:
        micro_grain_scale = pf.random.uniform(rng, 100.0, 400.0)

    return concrete(
        vector=vector,
        base_color=base_color,
        scale=1.0,
        roughness=roughness,
        crack_amount=crack_amount,
        crack_scale=crack_scale,
        snake_crack=snake_crack,
        color_variation=color_variation,
        crack_darkness=crack_darkness,
        surface_bump_strength=surface_bump_strength,
        displacement_depth=displacement_depth,
        micro_grain=micro_grain,
        micro_grain_scale=micro_grain_scale,
    )