//
// PRISM Wood material shader, based on the RTSL implementation by Karl Schmidt
// which in turn was based on work by Zhao Dong and Milos Hasan
// OSL port by Zap Andersson
//
// Modified: 2021-09-23
// Copyright 2021 Autodesk Inc, All rights reserved. This file is licensed under Apache 2.0 license
//    https://github.com/ADN-DevTech/3dsMax-OSL-Shaders/blob/license/LICENSE.txt
//

///////////////////////////////////////////////////////////////////////////
// Constants
///////////////////////////////////////////////////////////////////////////

#include <vector4.h>

#define float4 vector4

#define PI 3.14159265
#define TWO_PI (2.0 * PI)
#define HALF_PI (0.5 * PI)
#define bool int
#define true 1
#define false 0

#define PERFECT_SPECULAR_ROUGNESS 0.001

#define WOOD_PORE_BOTH  0
#define WOOD_PORE_EARLY 1
#define WOOD_PORE_LATE  2

///////////////////////////////////////////////////////////////////////////
// Utility functions for RapidRT Prism shaders
///////////////////////////////////////////////////////////////////////////

float frac(float f)
{
	return f - floor(f);
}

float saturate(float x) 
{
	return clamp(x,0.0,1.0);
}


float partialSmoothstep(float min, float max, float x){
  float tmp = saturate((x - min) / (max - min));
  if (tmp < 0.5) {
    return tmp;
  }
  else {
    return tmp*tmp*(3.0 - 2.0*tmp);
  }
}


///////////////////////////////////////////////////////////////
// Shader /////////////////////////////////////////////////////
shader prism_wood 
[[ string label = "Advanced Wood",
   string version = "1.0.0",
   string help  = "<h3>Advanced Wood</h3>"
                  "This OSL shader is similar to the classic 'Advanced Wood', except the UVW "
				  "can be modified with external shaders, and a random seed can be driven "
				  "to have per-object variations.<br>" 
				  "It is also fully supported in a 'High Quality' viewport and will<br>"
				  "render <i>exactly the same</i> on any OSL capable renderer."
   ]]
(
  vector UVW = transform("object", P),
  int    Seed = 0,
  //////////
  // Wood //
  //////////

  float scale = 1.0        [[ string label = "Scale" ]],
  int   unit_dependent = 1 [[ string label = "Unit Dependent", string widget = "checkBox", int connectable = 0]],
  int   axis           = 2 [[ string label = "Axis", string widget = "mapper", string options = "X:0|Y:1|Z:2", int connectable = 0]],

  // Wood Seasonal Growth Params
  float     ring_thickness      =  0.75  [[ string label = "Ring Thickness",       float min = 0.01, float max = 1.0 ]],
  float     late_wood_ratio      = 0.238 [[ string label = "Late Wood Ratio",      float min = 0.0,  float max = 1.0 ]],
  float     early_wood_sharpness = 0.395 [[ string label = "Early Wood Sharpness", int connectable = 0, float min = 0.0,  float max = 1.0 ]],
  float     late_wood_sharpness  = 0.109 [[ string label = "Late Wood Sharpness",  int connectable = 0, float min = 0.0,  float max = 1.0 ]],

  // Wood Grain Color Params
  color     early_wood_color           = color(0.2, 0.08, 0.024) [[ string label = "Early Wood Color" ]],
  bool      use_manual_late_wood_color = false [[ string label = "Use Manual Late Wood Color", string widget = "checkBox", int connectable = 0  ]],
  color     manual_late_wood_color     = color(0.217, 0.071, 0.01) [[ string label = "Late Wood Color" ]],
  float     late_wood_color_power      = 1.07 [[ string label = "Late Wood Color Power", float min = 0.01,  float max = 5.0 ]],

  // Wood Pores Params
  bool      use_pores        = true  [[ string label = "Use Pores", string widget = "checkBox", int connectable = 0  ]],
  int       pore_type        = 0     [[ string label = "Pore Type", int connectable = 0, string widget="mapper", string options = "Both:0|Early Only:1|Late Only:2" ]],
  float     pore_cell_dim    = 0.12  [[ string label = "Pore Cell Dim", int connectable = 0, float min = 0.01, float max = 3.0 ]],
  float     pore_radius      = 0.04  [[ string label = "Pore Radius", int connectable = 0, float min = 0.01, float max = 2.0 ]],
  float     pore_color_power = 2.13  [[ string label = "Pore Color Power", float min = 0.01, float max = 5.0 ]],
  float     pore_depth       = 0.01  [[ string label = "Pore Depth", int connectable = 0, float min = 0.0,  float max = 1.0 ]],

  // Wood Rays Params
  bool      use_rays             = true  [[ string label = "Use Rays", string widget = "checkBox", int connectable = 0  ]],
  float     ray_color_power      = 1.1   [[ string label = "Ray Color Power", float min = 0.01, float max = 5.0 ]],
  float     ray_seg_length_z     = 5.0   [[ string label = "Ray Seg Length Z", int connectable = 0, float min = 0.01, float max = 10.0 ]],
  int       ray_num_slices       = 160   [[ string label = "Ray Num Slices", int connectable = 0, int   min = 1,    int   max = 400  ]],
  float     ray_ellipse_depth    = 2.0   [[ string label = "Ray Ellipse Depth", int connectable = 0, float min = 0.01, float max = 10.0 ]],
  float     ray_ellipse_scale_x  = 0.2   [[ string label = "Ray Ellipse Scale Z", int connectable = 0, float min = 0.01, float max = 1.0 ]],
  float     ray_ellipse_z2x      = 10.0  [[ string label = "Ray Ellipse Z2X", int connectable = 0, float min = 0.0,  float max = 40.0 ]],
  float     ray_roughness        = 0.1   [[ string label = "Ray Ellipse Roughness", int connectable = 0, float min = 0.0,  float max = 1.0 ]],
  float     ray_bump_depth       = 0.02  [[ string label = "Ray Ellipse Bump Depth", int connectable = 0, float min = 0.0,  float max = 1.0 ]],


  // Wood Shading Params
  float     roughness            = 0.25   [[ string label = "Roughness", float min = 0.0,  float max = 1.0 ]],
  float     diffuse_lobe_weight  = 0.85   [[ string label = "Diffuse Lobe Weight" ]],
  bool      use_late_wood_bump   = true   [[ string widget = "checkBox", int connectable = 0  ]],
  float     late_wood_bump_depth = 0.002  [[ int connectable = 0, float min = 0.0,  float max = 1.0 ]],
  bool      use_groove_roughness = true   [[ string widget = "checkBox", int connectable = 0  ]],
  float     groove_roughness     = 0.19   [[ string label = "Groove Roughness" ]],


  // Wood Distortion Noise Params. The actual noise profiles are hidden from the user and only come from the presets,
  // just like the C++ Advanced Wood
  bool      use_fiber_perlin         = true [[ string widget = "checkBox", int connectable = 0 ]],
  int       fiber_perlin_band_count  = 3    [[ int connectable = 0, string widget="null" ]],
  float4     fiber_perlin_frequencies = {0.025, 0.05, 0.286, 0.0 } [[ int connectable = 0, string widget="null" ]],
  float4     fiber_perlin_weights     = {3.0, 1.0, 0.2, 0.0 }    [[ int connectable = 0, string widget="null" ]],
  
  float     fiber_perlin_scale_z     = 0.3  [[ int connectable = 0, float min = 0.05, float max = 1.0 ]],
  bool      use_fiber_cosine         = true [[ string widget = "checkBox", int connectable = 0  ]],
  int       fiber_cosine_band_count  = 2    [[ int connectable = 0, string widget="null" ]],
  float4     fiber_cosine_frequencies = {0.067, 0.25, 0.0, 0.0} [[ int connectable = 0, string widget="null" ]],
  float4     fiber_cosine_weights     = {2.5, 0.5, 0.0, 0.0} [[ int connectable = 0, string widget="null" ]],

  // Wood Color Noise Params
  bool      use_growth_perlin         = true[[ string widget = "checkBox", int connectable = 0  ]],
  int       growth_perlin_band_count  = 3 [[ int connectable = 0, string widget="null" ]],
  float4     growth_perlin_frequencies = {1.0,  0.2, 0.077, 0.0} [[ int connectable = 0, string widget="null" ]],
  float4     growth_perlin_weights     = {1.0, 2.0, 1.0, 0.0} [[ int connectable = 0, string widget="null" ]],

  bool      use_diffuse_perlin      = true [[ string widget = "checkBox", int connectable = 0  ]],
  int       diffuse_perlin_band_count  = 3 [[ int connectable = 0, string widget="null" ]],
  float4     diffuse_perlin_frequencies = {50.0, 10.0, 0.33, 0.0} [[ int connectable = 0, string widget="null" ]],
  float4     diffuse_perlin_weights     = {0.25, 0.2, 0.05, 0.0} [[ int connectable = 0, string widget="null" ]],
  float     diffuse_perlin_scale_z     = 0.15 [[ int connectable = 0, float min = 0.01,  float max = 1.0 ]],

  bool      use_early_wood_color_perlin      = true [[ string widget = "checkBox", int connectable = 0  ]],
  int       earlycolor_perlin_band_count  = 2 [[ int connectable = 0, string widget="null" ]],
  float4     earlycolor_perlin_frequencies = {0.123, 0.333, 0.0, 0.0} [[ int connectable = 0, string widget="null" ]],
  float4     earlycolor_perlin_weights     = {0.3, 0.5, 0.0, 0.0} [[ int connectable = 0, string widget="null" ]],

  bool      use_late_wood_color_perlin      = true [[ string widget = "checkBox", int connectable = 0  ]],
  int       latecolor_perlin_band_count    = 1 [[ int connectable = 0, string widget="null" ]],
  float4     latecolor_perlin_frequencies  = {0.222, 0.0, 0.0, 0.0} [[ int connectable = 0, string widget="null" ]],
  float4     latecolor_perlin_weights      = {0.75, 0.0, 0.0, 0.0} [[ int connectable = 0, string widget="null" ]],

  // This parameter should not be set by the user and is just used as information to the shader about scene
  // units, since it's defined in worldunits and should equal 1 cm.
  float cm_scale = 0.01    [[ int worldunits = 1, int connectable=0, string widget="null" ]],

   
  output color Col = 1 [[ string label = "Diffuse Color" ]],
  output float Roughness = 0.2,
  output float Bump = 0.0
)
{ 
	// Needs to be limited for numerical stability
	int adjustedSeed = Seed & 0xfff;

	// NOTE: This is not the pixel-identical noise to the C++ implementation, it WILL be slightly different
	float woodnoise1(float x)
	{
		return noise("perlin", x)/2.0;
	}
	
	float woodnoise1(vector x)
	{
		return noise("perlin", x)/2.0;
	}
	
	
	// Evaluate a 1D -> 1D wood noise profile
	// The profile is treated as set of pairs (weight, wavelength), each representing a single band of
	// Perlin noise. All the bands are evaluated as "weight * noise(frequency * x)" and summed.
	float woodNoise1(float4 profile_frequencies, float4 profile_weights, int band_count, float x)
	{
	  float offset = 0.0;
	  if (band_count > 0) {
	    float frequency = profile_frequencies.x;
	    float weight = profile_weights.x;
	    offset += weight * woodnoise1(frequency * x);
	  }
	  if (band_count > 1) {
	    float frequency = profile_frequencies.y;
	    float weight = profile_weights.y;
	    offset += weight * woodnoise1(frequency * x);
	  }
	  if (band_count > 2) {
	    float frequency = profile_frequencies.z;
	    float weight = profile_weights.z;
	    offset += weight * woodnoise1(frequency * x);
	  }
	  if (band_count > 3) {
	    float frequency = profile_frequencies.w;
	    float weight = profile_weights.w;
	    offset += weight * woodnoise1(frequency * x);
	  }
	  return offset;
	}
	
	
	// Evaluate a 3D -> 1D wood noise profile.
	// The profile is treated as set of pairs (weight, wavelength), each representing a single band of
	// Perlin noise. All the bands are evaluated as "weight * noise(frequency * x)" and summed.
	float woodNoise1(float4 profile_frequencies, float4 profile_weights, int band_count, vector p)
	{
	  float offset = 0.0;
	  if (band_count > 0) {
	    float frequency = profile_frequencies.x;
	    float weight = profile_weights.x;
	    offset += weight * woodnoise1(frequency * p);
	  }
	  if (band_count > 1) {
	    float frequency = profile_frequencies.y;
	    float weight = profile_weights.y;
	    offset += weight * woodnoise1(frequency * p);
	  }
	  if (band_count > 2) {
	    float frequency = profile_frequencies.z;
	    float weight = profile_weights.z;
	    offset += weight * woodnoise1(frequency * p);
	  }
	  if (band_count > 3) {
	    float frequency = profile_frequencies.w;
	    float weight = profile_weights.w;
	    offset += weight * woodnoise1(frequency * p);
	  }
	  return offset;
	}
	
	
	// Evaluate a cosine noise profile
	vector cosineNoiseRadialDir(float4 profile_frequencies, float4 profile_weights, int band_count,
	                            vector p)
	{
	  float radius = length(vector(p[0],p[1],0.0));
	  if ( radius < 0.00001 ) {
	    return p;
	  }
	
	  float cos_theta = p[0] / radius;
	  float sin_theta = p[1] / radius;
	
	  float radius_shift = 0.0;
	  float angdist = p[2] / TWO_PI;
	  if (band_count > 0) {
	    float frequency = profile_frequencies.x;
	    float weight = profile_weights.x;
	    radius_shift += weight * cos(frequency * angdist);
	  }
	  if (band_count > 1) {
	    float frequency = profile_frequencies.y;
	    float weight = profile_weights.y;
	    radius_shift += weight * cos(frequency * angdist);
	  }
	  if (band_count > 2) {
	    float frequency = profile_frequencies.z;
	    float weight = profile_weights.z;
	    radius_shift += weight * cos(frequency * angdist);
	  }
	  if (band_count > 3) {
	    float frequency = profile_frequencies.w;
	    float weight = profile_weights.w;
	    radius_shift += weight * cos(frequency * angdist);
	  }
	
	  // Radius at which the distortion reaches full power.
	  float min_radius = 1.5;
	  float shift_weight = partialSmoothstep(0.0, min_radius, radius) * radius_shift;
	
	    vector p2 = p;
	  p2[0] += cos_theta * shift_weight;
	  p2[1] += sin_theta * shift_weight;
	
	  return p2;
	}
	
	
	// Wyvill kernel
	float wyvill(float r)
	{
	  float tmp = 1.0 - r*r;
	  return r >= 1.0 ? 0.0 : tmp * tmp * tmp;
	}
	
	
	// Wyvill kernel with pre-squared input.
	float wyvillsq(float rsq)
	{
	  float tmp = 1.0 - rsq;
	  return rsq >= 1.0 ? 0.0 : tmp * tmp * tmp;
	}
	
  // The weight ratio of earlywood to latewood. 1.0 indicates all earlywood (lighter). 0.0
  // indicates latewood (darker). Computes the radial position as a side-effect.
  float computeEarlywoodRatio(vector position, output float radius)
  {
    // Seed is applied to radius
    radius = length(vector(position[0],position[1],0.0))  + abs(adjustedSeed) * 20;

    // Growth rate variation
    if (use_growth_perlin) {    
      radius += woodNoise1(growth_perlin_frequencies, growth_perlin_weights,
          growth_perlin_band_count, radius);
      if (radius < 0.0) {
        radius = 0.0;
      }
    }

    // ring_thickness will never be 0, this is worked around in PrismWood rtsi class.
    float fraction = frac(radius/ring_thickness);

    // As fraction goes from 0 to 1 the wood coloration sectors are
    // Earlywood -> Fall -> Latewood -> Rise
    // We find the transition points and handle each interval separately.
    float fall_to_late_wood =  1.0 - late_wood_ratio;
    float early_wood_to_fall = fall_to_late_wood * early_wood_sharpness;
    float late_wood_to_rise = fall_to_late_wood + late_wood_ratio * late_wood_sharpness;

    // Earlywood
    if (fraction <= early_wood_to_fall) {
      return 1.0;
    }

    // Fall
    if (fraction <= fall_to_late_wood) {
      float fall_width = fall_to_late_wood - early_wood_to_fall;
      float t = 1.0 - (fraction - early_wood_to_fall) / fall_width;
      return t;
    }

    // Latewood
    if (fraction <= late_wood_to_rise) {
      return 0.0;
    }

    // Rise
    float rise_width = 1.0 - late_wood_to_rise;
    float t = (fraction - late_wood_to_rise) / rise_width;
    return t;
  }


  // The total wyvill impulse from the radial fibers (rays) at position P.
  float weight3DRayImpulses(vector pin)
  {
    // The segment is the longitudinal subdivision. As rays can have a significant longitudinal
    // extent we sum the effect of the ray in both the segment P is in (id_0) and the nearest other
    // segment (id_1)

	vector p = pin;
	p[2] /= ray_ellipse_depth;

    int segment_id_0 = int(floor(p[2] / ray_seg_length_z));
    float factor = p[2] / ray_seg_length_z - float(segment_id_0);
    int segment_id_1 = factor > 0.5 ? segment_id_0 + 1 : segment_id_0 - 1;
    int segments[2] = { segment_id_0, segment_id_1 };

    // The slice is the circumferential subidivison. Rays have very little circumferential extent so
    // we simplify and only use the current slice. This will cause some clipping for rays near the
    // edge of a slice.
    float theta = atan2(p[1], p[0]);
    int slice_id = int(floor((theta + PI) / TWO_PI * float(ray_num_slices)));
    if (slice_id == ray_num_slices) {
      slice_id--;
    }

    // Sum the impulses for the two segments affecting P.
    float weight = 0.0;
    for (int i = 0; i < 2; i++) {
      int current_segment = segments[i];
      vector hash = noise("hash", slice_id + adjustedSeed * 39, current_segment);

      // Each ray starts at a position randomly offset from the center of the tree, at most 5 cm.
      float base_radial_offset = 5.0;
      if (length(vector(p[0],p[1],0.0)) >= base_radial_offset * hash[0]) {
	      // The direction of the ray
	      float ray_theta = (float(slice_id) + hash[0]) / float(ray_num_slices) * TWO_PI - PI;
	      float ray_pos_z = (float(current_segment) + hash[1]) * ray_seg_length_z;
	      vector ray_direction = vector(cos(ray_theta), sin(ray_theta), 0.0);
	
	      // Scale down the Z component so the ray cross-section is circular.
	      vector p_normalized = p;
	      p_normalized[2] = (p_normalized[2] - ray_pos_z) / ray_ellipse_z2x;
	
	      // Compute the point to line distance and use it for the wyvill impulse.
	      float dist = length(cross(ray_direction, p_normalized)) / length(ray_direction);
	      weight += wyvill(dist / ray_ellipse_scale_x);
	  }
    }

    return weight;
  }


  // The total wyvil impulse of the longitudinal pores at position P.
  float weight2DNeighborImpulses(vector p, float wood_weight)
  {
    if (wood_weight <= 0.0) {
      return 0.0;
    }
    float pore_radius = pore_radius * wood_weight;

    // Determine the boundaries of the pore cells that may affect us, given the pore radius and
    // cell dimensions.
    float x_left  = floor((p[0] - pore_radius) / pore_cell_dim);
    float x_right = floor((p[0] + pore_radius) / pore_cell_dim);
    float y_left  = floor((p[1] - pore_radius) / pore_cell_dim);
    float y_right = floor((p[1] + pore_radius) / pore_cell_dim);

    // Sum the wyvill impulse from all potentially contributing cells.
    float weight = 0.0;
    float inverse_radius_sq = 1.0 / (pore_radius * pore_radius);
    for (int j = int(y_left); j <= int(y_right); j++) {
      for (int i = int(x_left); i <= int(x_right); i++) {
        // Determine where the pore is in the cell
        vector offset =  noise("cell", i + adjustedSeed * 37, j);
        offset[2] = 0.0;
        vector pore_position = (offset + vector(float(i), float(j), 0.0)) * pore_cell_dim;

        // Compute the distance to the pore and use it for the wyvill impulse.
        vector diff = pore_position - vector(p[0],p[1],0.0);
        float diff_sq = dot(diff, diff);
        weight += wyvillsq(diff_sq * inverse_radius_sq);
      }
    }

    return weight;
  }


  // Darken a color based on whether or not that color is on a pore.
  color darkenColorWithPores(vector position, color col, float wood_weight,
  	float pore_color_power,
  	float pore_radius,
  	float pore_cell_dim )
  {
    float pores_weight = weight2DNeighborImpulses(position, wood_weight);
    float exponent = (pore_color_power - 1.0) * pores_weight + 1.0;
    return pow(col, exponent);
  }


  // Darken a color based on whether or not that color is on a ray.
  color darkenColorWithRays(vector position, color col, output float weight)
  {
    weight = weight3DRayImpulses(position);
    float exponent = (ray_color_power - 1.0) * weight + 1.0;
    return pow(col, exponent);
  }


  // Compute the diffuse color.
  color computeDiffuseColor(vector position, float early_wood_ratio, float radius, output float rayweight)
  {
    // Earlywood color
    color early_color = early_wood_color;
    if (use_early_wood_color_perlin) {
      float exponent = 1.0 + woodNoise1(earlycolor_perlin_frequencies, earlycolor_perlin_weights, earlycolor_perlin_band_count, radius);
      early_color = pow(early_color, exponent);
    }

    // Latewood color
    color late_wood_color = pow(early_color, late_wood_color_power);
    if (use_manual_late_wood_color) {
      late_wood_color = manual_late_wood_color;
    }
    if (use_late_wood_color_perlin) {
      float exponent = 1.0 + woodNoise1(latecolor_perlin_frequencies, latecolor_perlin_weights, latecolor_perlin_band_count, radius);
      late_wood_color = pow(late_wood_color, exponent);
    }

    // Overall diffuse albedo
    color diffuse_albedo =
        early_wood_ratio * early_color + (1.0 - early_wood_ratio) * late_wood_color;
    if (use_diffuse_perlin) {
      vector diffuse_position = position;
      diffuse_position[2] = diffuse_position[2] * diffuse_perlin_scale_z;
      float exponent = 1.0 + woodNoise1(diffuse_perlin_frequencies, diffuse_perlin_weights, diffuse_perlin_band_count, diffuse_position);
      diffuse_albedo = pow(diffuse_albedo, exponent);
    }

    // Darkening from pores
    if (use_pores) {
      if (pore_type == WOOD_PORE_EARLY) {
        diffuse_albedo = darkenColorWithPores(position, diffuse_albedo, early_wood_ratio,
			pore_color_power,
			pore_radius,
			pore_cell_dim);
      }
      else if (pore_type == WOOD_PORE_LATE) {
        diffuse_albedo = darkenColorWithPores(position, diffuse_albedo, 1.0 - early_wood_ratio,
			pore_color_power,
			pore_radius,
			pore_cell_dim);
      }
      else {
        diffuse_albedo = darkenColorWithPores(position, diffuse_albedo, 1.0,
			pore_color_power,
			pore_radius,
			pore_cell_dim);
      }
    }

    // Darkening from rays
    if (use_rays) {
      diffuse_albedo = darkenColorWithRays(position, diffuse_albedo, rayweight);
    }
    return diffuse_albedo;
  }


  // Apply radial cosine distortion to P.
  vector distort3DCosineRadialDir(vector p)
  {
    return cosineNoiseRadialDir(fiber_cosine_frequencies, fiber_cosine_weights,
        fiber_cosine_band_count, p);
  }


  // Apply Perlin distortion to P.
  vector distort3DPerlin(vector pp)
  {
  	vector p = pp;
    vector p_aniso = p;
    p_aniso[2] *= fiber_perlin_scale_z;

    float radial_shift = woodNoise1(fiber_perlin_frequencies, fiber_perlin_weights,
        fiber_perlin_band_count, p_aniso);
    p[0] += radial_shift;
    p[1] += radial_shift;

    return p;
  }

    // Apply distortion to P.
  vector distort(vector pi)
  {
  	vector p = pi;
    // Cosine distortion
    if (use_fiber_cosine) {
      p = distort3DCosineRadialDir(p);
    }

    // Perlin distortion
    if (use_fiber_perlin) {
      p = distort3DPerlin(p);
    }

    return p;
  }

  // Compute the height differential from pores
  float heightFrom2DNeighborPores(vector position, float wood_weight)
  {
    // Assume pores are always negative depth
    return -weight2DNeighborImpulses(position, wood_weight) * pore_depth;
  }


  // Compute the total bump displacement
  float heightVarForBump(vector position)
  {
    float height = 0.0;

    // compute the weight from the early wood ratio
    float radius;
    float early_wood_ratio = computeEarlywoodRatio(position, radius);

    // consider pores case
    if (use_pores) {
      if (pore_type == WOOD_PORE_EARLY) {
        height += heightFrom2DNeighborPores(position, early_wood_ratio);
      }
      else if (pore_type == WOOD_PORE_LATE) {
        height += heightFrom2DNeighborPores(position, 1.0 - early_wood_ratio);
      }
      else {
        height += heightFrom2DNeighborPores(position, 1.0);
      }
    }

    if (use_late_wood_bump) {
      height += (1.0 - early_wood_ratio) * late_wood_bump_depth;
    }

    return height;
  }

  // Compute treespace coordinates, pick the right axis
  vector tree_position = UVW;
  if (axis == 0)  // Adjust for X axis
  {
	tree_position[0] = UVW[1];
	tree_position[1] = UVW[2];
	tree_position[2] = UVW[0];
  }
  else if (axis == 1) // Adjust for Y axis
  {
	tree_position[0] = UVW[2];
	tree_position[1] = UVW[0];
	tree_position[2] = UVW[1];
  }

  // Apply the scale
  if (scale != 0.0)
  {
    float scaling = scale;
	if (unit_dependent)
	{
		// Scale by the unit-dependent parameter. This should not be set by the user
		// and always be kept at whatever value matches 1 centimeter in the current
		// scene units
		scaling *= cm_scale;
	}
	tree_position /= scaling;
  }
 
    // Apply random seed simply as a Z shift, simplifies all the noise
    // calls and is just as "random". Seed still applies to the cellnoise 
    // functions that are not position driven.
    tree_position[2] += adjustedSeed * 50;
    
    vector tree_distorted_position = distort(tree_position);

	Bump = heightVarForBump(tree_distorted_position);
    
    float surface_roughness = roughness;

    // Wood lobe
    float radius;

    float early_wood_ratio = computeEarlywoodRatio(tree_distorted_position, radius);

    // Compute the surface roughness
    if (use_groove_roughness) {
      surface_roughness =
          early_wood_ratio * groove_roughness + (1.0 - early_wood_ratio) * surface_roughness;
    }
    float lobe_roughness = surface_roughness;

    // Wood diffuse lobe.
	float rayweight = 0.0;
    vector wood_f0 = vector(0.04);
    vector diffuse_color =
        computeDiffuseColor(tree_distorted_position, early_wood_ratio, radius, rayweight) * diffuse_lobe_weight;

	// Output
	Col       = diffuse_color;
	Roughness = clamp(lobe_roughness + ray_roughness * rayweight, 0.0, 1.0);
	Bump      -= rayweight * ray_bump_depth;
}