from mathutils import Vector, Matrix, Color, Quaternion
import gpu
from bgl import *
from .math_utils import *
from .buffer_utils import buffer_to_str
from .gl_utils import gl_utils

__all__ = (
  "DataType",
  "ShaderType",
  "get_shader_src",
  "get_hair_refine_shader",
  "GLShader"
)

COMMON_VIEW_LIB_SRC = """
#define COMMON_VIEW_LIB

/* keep in sync with DRWManager.view_data */
layout(std140) uniform viewBlock
{
  /* Same order as DRWViewportMatrixType */
  //mat4 ViewProjectionMatrix;
  //mat4 ViewProjectionMatrixInverse;
  mat4 ViewMatrix;
  //mat4 ViewMatrixInverse;
  mat4 ProjectionMatrix;
  //mat4 ProjectionMatrixInverse;

  /* View frustum corners [NDC(-1.0, -1.0, -1.0) & NDC(1.0, 1.0, 1.0)].
   * Fourth components are near and far values. */
  //vec4 ViewVecs[2];
};

//#define ViewNear (ViewVecs[0].w)
//#define ViewFar (ViewVecs[1].w)

// # define cameraForward ViewMatrixInverse[2].xyz
// # define cameraPos ViewMatrixInverse[3].xyz
// # define cameraVec(P) ((ProjectionMatrix[3][3] == 0.0) ? normalize(cameraPos - P) : cameraForward)
// # define viewCameraVec(vP) ((ProjectionMatrix[3][3] == 0.0) ? normalize(-vP) : vec3(0.0, 0.0, 1.0))

#ifdef COMMON_GLOBALS_LIB
float mul_project_m4_v3_zfac(in vec3 co)
{
  return pixelFac * ((ViewProjectionMatrix[0][3] * co.x) + (ViewProjectionMatrix[1][3] * co.y) +
                     (ViewProjectionMatrix[2][3] * co.z) + ViewProjectionMatrix[3][3]);
}
#endif

/* Not the right place but need to be common to all overlay's.
 * TODO Split to an overlay lib. */
mat4 extract_matrix_packed_data(mat4 mat, out vec4 dataA, out vec4 dataB)
{
  const float div = 1.0 / 255.0;
  int a = int(mat[0][3]);
  int b = int(mat[1][3]);
  int c = int(mat[2][3]);
  int d = int(mat[3][3]);
  dataA = vec4(a & 0xFF, a >> 8, b & 0xFF, b >> 8) * div;
  dataB = vec4(c & 0xFF, c >> 8, d & 0xFF, d >> 8) * div;
  mat[0][3] = mat[1][3] = mat[2][3] = 0.0;
  mat[3][3] = 1.0;
  return mat;
}

/* Same here, Not the right place but need to be common to all overlay's.
 * TODO Split to an overlay lib. */
/* edge_start and edge_pos needs to be in the range [0..sizeViewport]. */
vec4 pack_line_data(vec2 frag_co, vec2 edge_start, vec2 edge_pos)
{
  vec2 edge = edge_start - edge_pos;
  float len = length(edge);
  if (len > 0.0) {
    edge /= len;
    vec2 perp = vec2(-edge.y, edge.x);
    float dist = dot(perp, frag_co - edge_start);
    /* Add 0.1 to diffenrentiate with cleared pixels. */
    return vec4(perp * 0.5 + 0.5, dist * 0.25 + 0.5 + 0.1, 1.0);
  }
  else {
    /* Default line if the origin is perfectly aligned with a pixel. */
    return vec4(1.0, 0.0, 0.5 + 0.1, 1.0);
  }
}

uniform mat4 ModelMatrix;
//uniform mat4 ModelMatrixInverse;
#define ViewMatrixInverse inverse(ViewMatrix)
#define ProjectionMatrixInverse inverse(ProjectionMatrix)
#define ViewProjectionMatrix (ProjectionMatrix * ViewMatrix)
#define ViewProjectionMatrixInverse inverse(ViewProjectionMatrix)
#define ModelMatrixInverse inverse(ModelMatrix)

/** Transform shortcuts. */
/* Rule of thumb: Try to reuse world positions and normals because converting through viewspace
 * will always be decomposed in at least 2 matrix operation. */

/**
 * Some clarification:
 * Usually Normal matrix is transpose(inverse(ViewMatrix * ModelMatrix))
 *
 * But since it is slow to multiply matrices we decompose it. Decomposing
 * inversion and transposition both invert the product order leaving us with
 * the same original order:
 * transpose(ViewMatrixInverse) * transpose(ModelMatrixInverse)
 *
 * Knowing that the view matrix is orthogonal, the transpose is also the inverse.
 * Note: This is only valid because we are only using the mat3 of the ViewMatrixInverse.
 * ViewMatrix * transpose(ModelMatrixInverse)
 */
#define NormalMatrix transpose(mat3(ModelMatrixInverse))
#define NormalMatrixInverse transpose(mat3(ModelMatrix))

#define normal_object_to_view(n) (mat3(ViewMatrix) * (NormalMatrix * n))
#define normal_object_to_world(n) (NormalMatrix * n)
#define normal_world_to_object(n) (NormalMatrixInverse * n)
#define normal_world_to_view(n) (mat3(ViewMatrix) * n)
#define normal_view_to_world(n) (mat3(ViewMatrixInverse) * n)

#define point_object_to_ndc(p) (ViewProjectionMatrix * vec4((ModelMatrix * vec4(p, 1.0)).xyz, 1.0))
#define point_object_to_view(p) ((ViewMatrix * vec4((ModelMatrix * vec4(p, 1.0)).xyz, 1.0)).xyz)
#define point_object_to_world(p) ((ModelMatrix * vec4(p, 1.0)).xyz)
#define point_view_to_ndc(p) (ProjectionMatrix * vec4(p, 1.0))
#define point_view_to_object(p) ((ModelMatrixInverse * (ViewMatrixInverse * vec4(p, 1.0))).xyz)
#define point_view_to_world(p) ((ViewMatrixInverse * vec4(p, 1.0)).xyz)
#define point_world_to_ndc(p) (ViewProjectionMatrix * vec4(p, 1.0))
#define point_world_to_object(p) ((ModelMatrixInverse * vec4(p, 1.0)).xyz)
#define point_world_to_view(p) ((ViewMatrix * vec4(p, 1.0)).xyz)


/* ---- Opengl Depth conversion ---- */

//float linear_depth(bool is_persp, float z, float zf, float zn)
//{
//  if (is_persp) {
//    return (zn * zf) / (z * (zn - zf) + zf);
//  }
//  else {
//    return (z * 2.0 - 1.0) * zf;
//  }
//}
//
//float buffer_depth(bool is_persp, float z, float zf, float zn)
//{
//  if (is_persp) {
//    return (zf * (zn - z)) / (z * (zn - zf));
//  }
//  else {
//    return (z / (zf * 2.0)) + 0.5;
//  }
//}
//
//float get_view_z_from_depth(float depth)
//{
//  if (ProjectionMatrix[3][3] == 0.0) {
//    float d = 2.0 * depth - 1.0;
//    return -ProjectionMatrix[3][2] / (d + ProjectionMatrix[2][2]);
//  }
//  else {
//    return ViewVecs[0].z + depth * ViewVecs[1].z;
//  }
//}
//
//float get_depth_from_view_z(float z)
//{
//  if (ProjectionMatrix[3][3] == 0.0) {
//    float d = (-ProjectionMatrix[3][2] / z) - ProjectionMatrix[2][2];
//    return d * 0.5 + 0.5;
//  }
//  else {
//    return (z - ViewVecs[0].z) / ViewVecs[1].z;
//  }
//}
//
//vec2 get_uvs_from_view(vec3 view)
//{
//  vec4 ndc = ProjectionMatrix * vec4(view, 1.0);
//  return (ndc.xy / ndc.w) * 0.5 + 0.5;
//}
//
//vec3 get_view_space_from_depth(vec2 uvcoords, float depth)
//{
//  if (ProjectionMatrix[3][3] == 0.0) {
//    return vec3(ViewVecs[0].xy + uvcoords * ViewVecs[1].xy, 1.0) * get_view_z_from_depth(depth);
//  }
//  else {
//    return ViewVecs[0].xyz + vec3(uvcoords, depth) * ViewVecs[1].xyz;
//  }
//}
//
//vec3 get_world_space_from_depth(vec2 uvcoords, float depth)
//{
//  return (ViewMatrixInverse * vec4(get_view_space_from_depth(uvcoords, depth), 1.0)).xyz;
//}
//
//vec3 get_view_vector_from_screen_uv(vec2 uv)
//{
//  if (ProjectionMatrix[3][3] == 0.0) {
//    return normalize(vec3(ViewVecs[0].xy + uv * ViewVecs[1].xy, 1.0));
//  }
//  else {
//    return vec3(0.0, 0.0, 1.0);
//  }
//}

"""


SHADER_INTERFACE_LIB_SRC = """
IN_OUT ShaderStageInterface
{
  vec3 normal_interp;
  vec3 color_interp;
  float alpha_interp;
  vec2 uv_interp;
  flat float packed_rough_metal;
};

"""


COMMON_LIB_SRC = """

#define EPSILON 0.00001

#define CAVITY_BUFFER_RANGE 4.0

#ifdef WORKBENCH_ENCODE_NORMALS

#  define WB_Normal vec2

/* From http://aras-p.info/texts/CompactNormalStorage.html
 * Using Method #4: Spheremap Transform */
vec3 workbench_normal_decode(vec4 enc)
{
  vec2 fenc = enc.xy * 4.0 - 2.0;
  float f = dot(fenc, fenc);
  float g = sqrt(1.0 - f / 4.0);
  vec3 n;
  n.xy = fenc * g;
  n.z = 1 - f / 2;
  return n;
}

/* From http://aras-p.info/texts/CompactNormalStorage.html
 * Using Method #4: Spheremap Transform */
WB_Normal workbench_normal_encode(bool front_face, vec3 n)
{
  n = normalize(front_face ? n : -n);
  float p = sqrt(n.z * 8.0 + 8.0);
  n.xy = clamp(n.xy / p + 0.5, 0.0, 1.0);
  return n.xy;
}

#else
#  define WB_Normal vec3
/* Well just do nothing... */
#  define workbench_normal_encode(f, a) (a)
#  define workbench_normal_decode(a) (a.xyz)
#endif /* WORKBENCH_ENCODE_NORMALS */

/* Encoding into the alpha of a RGBA16F texture. (10bit mantissa) */
#define TARGET_BITCOUNT 8u
#define METALLIC_BITS 3u /* Metallic channel is less important. */
#define ROUGHNESS_BITS (TARGET_BITCOUNT - METALLIC_BITS)

/* Encode 2 float into 1 with the desired precision. */
float workbench_float_pair_encode(float v1, float v2)
{
  // const uint v1_mask = ~(0xFFFFFFFFu << ROUGHNESS_BITS);
  // const uint v2_mask = ~(0xFFFFFFFFu << METALLIC_BITS);
  /* Same as above because some compiler are very dumb and think we use medium int. */
  const int v1_mask = 0x1F;
  const int v2_mask = 0x7;
  int iv1 = int(v1 * float(v1_mask));
  int iv2 = int(v2 * float(v2_mask)) << int(ROUGHNESS_BITS);
  return float(iv1 | iv2);
}

void workbench_float_pair_decode(float data, out float v1, out float v2)
{
  // const uint v1_mask = ~(0xFFFFFFFFu << ROUGHNESS_BITS);
  // const uint v2_mask = ~(0xFFFFFFFFu << METALLIC_BITS);
  /* Same as above because some compiler are very dumb and think we use medium int.  */
  const int v1_mask = 0x1F;
  const int v2_mask = 0x7;
  int idata = int(data);
  v1 = float(idata & v1_mask) * (1.0 / float(v1_mask));
  v2 = float(idata >> int(ROUGHNESS_BITS)) * (1.0 / float(v2_mask));
}

"""


MATERIAL_LIB_SRC = """

layout(std140) uniform material_block
{
  vec4 mat_data[100];
};

uniform int materialIndex;

void workbench_material_data_get(
    out vec3 color, out float alpha, out float roughness, out float metallic)
{
  vec4 data = mat_data[materialIndex];
  color = data.rgb;

  uint encoded_data = floatBitsToUint(data.w);
  alpha = float((encoded_data >> 16u) & 0xFFu) * (1.0 / 255.0);
  roughness = float((encoded_data >> 8u) & 0xFFu) * (1.0 / 255.0);
  metallic = float(encoded_data & 0xFFu) * (1.0 / 255.0);
}

"""


SHADING_DATA_LIB_SRC = """

struct LightData {
  vec4 direction;
  vec4 specular_color;
  vec4 diffuse_color_wrap; /* rgb: diffuse col a: wrapped lighting factor */
};

struct WorldData {
  vec4 viewport_size;
  /* - 16 bytes alignment-  */
  LightData lights[4];
  vec4 ambient_color;

  int matcap_orientation;
  bool use_specular;
  int _pad1;
  int _pad2;
};

#define viewport_size_inv viewport_size.zw

layout(std140) uniform world_block
{
  WorldData world_data;
};

"""


MATCAP_LIB_SRC = """

#pragma LIB_REQUIRE(shading_data_lib.glsl)

vec2 matcap_uv_compute(vec3 I, vec3 N, bool flipped)
{
  /* Quick creation of an orthonormal basis */
  float a = 1.0 / (1.0 + I.z);
  float b = -I.x * I.y * a;
  vec3 b1 = vec3(1.0 - I.x * I.x * a, b, -I.x);
  vec3 b2 = vec3(b, 1.0 - I.y * I.y * a, -I.y);
  vec2 matcap_uv = vec2(dot(b1, N), dot(b2, N));
  if (flipped) {
    matcap_uv.x = -matcap_uv.x;
  }
  return matcap_uv * 0.496 + 0.5;
}

uniform sampler2D matcapDiffuseImage;
uniform sampler2D matcapSpecularImage;

vec3 get_matcap_lighting(vec3 base_color, vec3 N, vec3 I)
{
  bool flipped = world_data.matcap_orientation != 0;
  vec2 uv = matcap_uv_compute(I, N, flipped);

  vec3 diffuse = textureLod(matcapDiffuseImage, uv, 0.0).rgb;
  vec3 specular = textureLod(matcapSpecularImage, uv, 0.0).rgb;

  return diffuse * base_color + specular * float(world_data.use_specular);
}

"""


WORLD_LIGHT_LIB_SRC = """

#pragma LIB_REQUIRE(shading_data_lib.glsl)

/* [Drobot2014a] Low Level Optimizations for GCN */
vec4 fast_rcp(vec4 v)
{
  return intBitsToFloat(0x7eef370b - floatBitsToInt(v));
}

vec3 brdf_approx(vec3 spec_color, float roughness, float NV)
{
  /* Very rough own approx. We don't need it to be correct, just fast.
   * Just simulate fresnel effect with roughness attenuation. */
  float fresnel = exp2(-8.35 * NV) * (1.0 - roughness);
  return mix(spec_color, vec3(1.0), fresnel);
}

void prep_specular(
    vec3 L, vec3 I, vec3 N, vec3 R, out float NL, out float wrapped_NL, out float spec_angle)
{
  wrapped_NL = dot(L, R);
  vec3 half_dir = normalize(L + I);
  spec_angle = clamp(dot(half_dir, N), 0.0, 1.0);
  NL = clamp(dot(L, N), 0.0, 1.0);
}

/* Normalized Blinn shading */
vec4 blinn_specular(vec4 shininess, vec4 spec_angle, vec4 NL)
{
  /* Pi is already divided in the light power.
   * normalization_factor = (shininess + 8.0) / (8.0 * M_PI) */
  vec4 normalization_factor = shininess * 0.125 + 1.0;
  vec4 spec_light = pow(spec_angle, shininess) * NL * normalization_factor;

  return spec_light;
}

/* NL need to be unclamped. w in [0..1] range. */
vec4 wrapped_lighting(vec4 NL, vec4 w)
{
  vec4 w_1 = w + 1.0;
  vec4 denom = fast_rcp(w_1 * w_1);
  return clamp((NL + w) * denom, 0.0, 1.0);
}

vec3 get_world_lighting(vec3 base_color, float roughness, float metallic, vec3 N, vec3 I)
{
  vec3 specular_color, diffuse_color;

  if (world_data.use_specular) {
    diffuse_color = mix(base_color, vec3(0.0), metallic);
    specular_color = mix(vec3(0.05), base_color, metallic);
  }
  else {
    diffuse_color = base_color;
    specular_color = vec3(0.0);
  }

  vec3 specular_light = world_data.ambient_color.rgb;
  vec3 diffuse_light = world_data.ambient_color.rgb;
  vec4 wrap = vec4(world_data.lights[0].diffuse_color_wrap.a,
                   world_data.lights[1].diffuse_color_wrap.a,
                   world_data.lights[2].diffuse_color_wrap.a,
                   world_data.lights[3].diffuse_color_wrap.a);

  if (world_data.use_specular) {
    /* Prepare Specular computation. Eval 4 lights at once. */
    vec3 R = -reflect(I, N);
    vec4 spec_angle, spec_NL, wrap_NL;
    prep_specular(world_data.lights[0].direction.xyz, I, N, R, spec_NL.x, wrap_NL.x, spec_angle.x);
    prep_specular(world_data.lights[1].direction.xyz, I, N, R, spec_NL.y, wrap_NL.y, spec_angle.y);
    prep_specular(world_data.lights[2].direction.xyz, I, N, R, spec_NL.z, wrap_NL.z, spec_angle.z);
    prep_specular(world_data.lights[3].direction.xyz, I, N, R, spec_NL.w, wrap_NL.w, spec_angle.w);

    vec4 gloss = vec4(1.0 - roughness);
    /* Reduce gloss for smooth light. (simulate bigger light) */
    gloss *= 1.0 - wrap;
    vec4 shininess = exp2(10.0 * gloss + 1.0);

    vec4 spec_light = blinn_specular(shininess, spec_angle, spec_NL);

    /* Simulate Env. light. */
    vec4 w = mix(wrap, vec4(1.0), roughness);
    vec4 spec_env = wrapped_lighting(wrap_NL, w);

    spec_light = mix(spec_light, spec_env, wrap * wrap);

    /* Multiply result by lights specular colors. */
    specular_light += spec_light.x * world_data.lights[0].specular_color.rgb;
    specular_light += spec_light.y * world_data.lights[1].specular_color.rgb;
    specular_light += spec_light.z * world_data.lights[2].specular_color.rgb;
    specular_light += spec_light.w * world_data.lights[3].specular_color.rgb;

    float NV = clamp(dot(N, I), 0.0, 1.0);
    specular_color = brdf_approx(specular_color, roughness, NV);
  }
  specular_light *= specular_color;

  /* Prepare diffuse computation. Eval 4 lights at once. */
  vec4 diff_NL;
  diff_NL.x = dot(world_data.lights[0].direction.xyz, N);
  diff_NL.y = dot(world_data.lights[1].direction.xyz, N);
  diff_NL.z = dot(world_data.lights[2].direction.xyz, N);
  diff_NL.w = dot(world_data.lights[3].direction.xyz, N);

  vec4 diff_light = wrapped_lighting(diff_NL, wrap);

  /* Multiply result by lights diffuse colors. */
  diffuse_light += diff_light.x * world_data.lights[0].diffuse_color_wrap.rgb;
  diffuse_light += diff_light.y * world_data.lights[1].diffuse_color_wrap.rgb;
  diffuse_light += diff_light.z * world_data.lights[2].diffuse_color_wrap.rgb;
  diffuse_light += diff_light.w * world_data.lights[3].diffuse_color_wrap.rgb;

  /* Energy conservation with colored specular look strange.
   * Limit this strangeness by using mono-chromatic specular intensity. */
  float spec_energy = dot(specular_color, vec3(0.33333));

  diffuse_light *= diffuse_color * (1.0 - spec_energy);

  return diffuse_light + specular_light;
}

// uniform bool forceShadowing = false;
// 
// float get_shadow(vec3 N)
// {
//   float light_factor = -dot(N, world_data.shadow_direction_vs.xyz);
//   float shadow_mix = smoothstep(world_data.shadow_shift, world_data.shadow_focus, light_factor);
//   shadow_mix *= forceShadowing ? 0.0 : world_data.shadow_mul;
//   return shadow_mix + world_data.shadow_add;
// }

"""

COMMON_HAIR_LIB_SRC = """

/**
 * Library to create hairs dynamically from control points.
 * This is less bandwidth intensive than fetching the vertex attributes
 * but does more ALU work per vertex. This also reduces the amount
 * of data the CPU has to precompute and transfer for each update.
 */

/**
 * hairStrandsRes: Number of points per hair strand.
 * 2 - no subdivision
 * 3+ - 1 or more interpolated points per hair.
 */
uniform int hairStrandsRes = 8;

/**
 * hairThicknessRes : Subdiv around the hair.
 * 1 - Wire Hair: Only one pixel thick, independent of view distance.
 * 2 - Polystrip Hair: Correct width, flat if camera is parallel.
 * 3+ - Cylinder Hair: Massive calculation but potentially perfect. Still need proper support.
 */
uniform int hairThicknessRes = 1;

/* Hair thickness shape. */
uniform float hairRadRoot = 0.01;
uniform float hairRadTip = 0.0;
uniform float hairRadShape = 0.5;
uniform float hairRadScale = 0.005;
uniform bool hairCloseTip = true;
uniform bool hairUseShapeCurve = false;
uniform bool hairUseShapeRandom = false;
uniform bool hairShowColor = false;
uniform int hairRandomSeed = 0;
uniform int hairTotCurves = 0;
uniform float hairRadStart = 0.0;
uniform float hairRandomThres = 0.0;

uniform mat4 hairDupliMatrix;

/* -- Per control points -- */
uniform samplerBuffer hairPointBuffer; /* RGBA32F */
#define point_position xyz
#define point_time w /* Position along the hair length */

/* -- Per strands data -- */
uniform usamplerBuffer hairStrandBuffer;    /* R32UI */
uniform usamplerBuffer hairStrandSegBuffer; /* R16UI */

/* Not used, use one buffer per uv layer */
// uniform samplerBuffer hairUVBuffer; /* RG32F */
// uniform samplerBuffer hairColBuffer; /* RGBA16 linear color */

/* -- Subdivision stage -- */
/**
 * We use a transform feedback to preprocess the strands and add more subdivision to it.
 * For the moment these are simple smooth interpolation but one could hope to see the full
 * children particle modifiers being evaluated at this stage.
 *
 * If no more subdivision is needed, we can skip this step.
 */

#ifdef HAIR_PHASE_SUBDIV
int hair_get_base_id(float local_time, int strand_segments, out float interp_time)
{
  float time_per_strand_seg = 1.0 / float(strand_segments);

  float ratio = local_time / time_per_strand_seg;
  interp_time = fract(ratio);

  return int(ratio);
}

void hair_get_interp_attrs(
    out vec4 data0, out vec4 data1, out vec4 data2, out vec4 data3, out float interp_time)
{
  float local_time = float(gl_VertexID % hairStrandsRes) / float(hairStrandsRes - 1);

  int hair_id = gl_VertexID / hairStrandsRes;
  int strand_offset = int(texelFetch(hairStrandBuffer, hair_id).x);
  int strand_segments = int(texelFetch(hairStrandSegBuffer, hair_id).x);

  int id = hair_get_base_id(local_time, strand_segments, interp_time);

  int ofs_id = id + strand_offset;

  data0 = texelFetch(hairPointBuffer, ofs_id - 1);
  data1 = texelFetch(hairPointBuffer, ofs_id);
  data2 = texelFetch(hairPointBuffer, ofs_id + 1);
  data3 = texelFetch(hairPointBuffer, ofs_id + 2);

  if (id <= 0) {
    /* root points. Need to reconstruct previous data. */
    data0 = data1 * 2.0 - data2;
  }
  if (id + 1 >= strand_segments) {
    /* tip points. Need to reconstruct next data. */
    data3 = data2 * 2.0 - data1;
  }
}
#endif

/* -- Drawing stage -- */
/**
 * For final drawing, the vertex index and the number of vertex per segment
 */

#if !defined(HAIR_PHASE_SUBDIV) && defined(GPU_VERTEX_SHADER)
uniform usamplerBuffer hairStrandColorBuffer; /* R8UI */
uniform usamplerBuffer hairFrandBuffer; /* RBG32UI */

layout(std140) uniform shapeCurveBlock
{
  /* Used as float[260], because the array members in the uniformblock must
   * be 16-byte aligned, we use vec4. The members layout is:
   * values...(257)|min|max|range| */
  vec4 shapeCurveTable[65];
};

layout(std140) uniform colorGroupBlock
{
  vec4 colorGroups[256];
};

#define CurveTableMin shapeCurveTable[64][1]
#define CurveTableMax shapeCurveTable[64][2]
#define CurveTableRange shapeCurveTable[64][3]
#define PSYS_FRAND_COUNT 1024

int hair_get_strand_id(void)
{
  return gl_VertexID / (hairStrandsRes * hairThicknessRes);
}

int hair_get_base_id(void)
{
  return gl_VertexID / hairThicknessRes;
}

void child_frand_vec(int seed, int rand, out vec3 vec)
{
  int offset = int(texelFetch(hairFrandBuffer, rand % PSYS_FRAND_COUNT).g);
  int multi = int(texelFetch(hairFrandBuffer, 1).b);
  vec.x = uintBitsToFloat(texelFetch(hairFrandBuffer, (offset + seed * multi) % PSYS_FRAND_COUNT).r);
  vec.y = uintBitsToFloat(texelFetch(hairFrandBuffer, (offset + (seed + 1) * multi) % PSYS_FRAND_COUNT).r);
  vec.z = uintBitsToFloat(texelFetch(hairFrandBuffer, (offset + (seed + 2) * multi) % PSYS_FRAND_COUNT).r);
}

float seed_frand(int seed)
{
  ivec2 d = ivec2(texelFetch(hairFrandBuffer, 1).gb);
  return uintBitsToFloat(texelFetch(hairFrandBuffer, (d.x + seed * d.y) % PSYS_FRAND_COUNT).r);
}

float RadiusGetCustomdata(int path_num)
{
  float data = seed_frand(path_num + hairRandomSeed);
  if (hairRandomThres < 0 && path_num < int(-hairRandomThres * hairTotCurves)) {
    data /= 3;
  }
  return clamp(data, hairRadStart, 1.0);
}

/* Copied from cycles. */
float hair_shaperadius(float shape, float root, float tip, float time)
{
  float radius = 1.0 - time;

  if (shape < 0.0) {
    radius = pow(radius, 1.0 + shape);
  }
  else {
    radius = pow(radius, 1.0 / (1.0 - shape));
  }

  if (hairCloseTip && (time > 0.99)) {
    return 0.0;
  }

  return (radius * (root - tip)) + tip;
}

float hair_shape_curve_radius(float time)
{
  /* Clamp time!!! */
  time = clamp(time, CurveTableMin, CurveTableMax);
  /* index in table */
  float fi = (time - CurveTableMin) * CurveTableRange;
  int i = int(fi);

  if (hairCloseTip && (time > 0.99)) {
    return 0.0;
  }

  int ti = int(i / 4), tsi = i % 4;
  int tni = int((i + 1) / 4), tnsi = (i + 1) % 4;
  fi = fi - float(i);
  return ((1.0f - fi) * shapeCurveTable[ti][tsi] + fi * shapeCurveTable[tni][tnsi]) * hairRadScale;
}

#  ifdef OS_MAC
in float dummy;
#  endif

void hair_get_pos_tan_binor_time(bool is_persp,
                                 mat4 invmodel_mat,
                                 vec3 camera_pos,
                                 vec3 camera_z,
                                 out vec3 wpos,
                                 out vec3 wtan,
                                 out vec3 wbinor,
                                 out float time,
                                 out float thickness,
                                 out float thick_time)
{
  int id = hair_get_base_id();
  float rad_fac = 1.0;
  vec4 data = texelFetch(hairPointBuffer, id);
  wpos = data.point_position;
  time = data.point_time;

#  ifdef OS_MAC
  /* Generate a dummy read to avoid the driver bug with shaders having no
   * vertex reads on macOS (T60171) */
  wpos.y += dummy * 0.0;
#  endif

  if (time == 0.0) {
    /* Hair root */
    wtan = texelFetch(hairPointBuffer, id + 1).point_position - wpos;
  }
  else {
    wtan = wpos - texelFetch(hairPointBuffer, id - 1).point_position;
  }

  mat4 obmat = hairDupliMatrix;

  wpos = (obmat * vec4(wpos, 1.0)).xyz;
  wtan = -normalize(mat3(obmat) * wtan);

  vec3 camera_vec = (is_persp) ? camera_pos - wpos : camera_z;
  wbinor = normalize(cross(camera_vec, wtan));

  if (hairUseShapeRandom) {
    vec3 vec;
    int strand_id = hair_get_strand_id();
    child_frand_vec(strand_id + 27, hairRandomSeed, vec);
    if (hairRandomThres <= 0 || abs(-1.5 + vec.x + vec.y + vec.z) >= 1.5 * hairRandomThres) {
      rad_fac = RadiusGetCustomdata(strand_id);
    }
  }

  if (hairUseShapeCurve) {
    thickness = hair_shape_curve_radius(time) * rad_fac;
  }
  else {
    thickness = hair_shaperadius(hairRadShape, hairRadRoot, hairRadTip, time) * rad_fac;
  }

  if (hairThicknessRes > 1) {
    thick_time = float(gl_VertexID % hairThicknessRes) / float(hairThicknessRes - 1);
    thick_time = thickness * (thick_time * 2.0 - 1.0);

    /* Take object scale into account.
     * NOTE: This only works fine with uniform scaling. */
    float scale = 1.0 / length(mat3(invmodel_mat) * wbinor);

    wpos += wbinor * thick_time * scale;
  }
}

vec3 hair_get_shown_color(vec3 color)
{
  if (hairShowColor) {
    int strand_id = hair_get_strand_id();
    int index = int(texelFetch(hairStrandColorBuffer, strand_id).x);
    if (index < 255) {
      return colorGroups[index].rgb;
    }
  }
  return color;
}

vec2 hair_get_customdata_vec2(const samplerBuffer cd_buf)
{
  int id = hair_get_strand_id();
  return texelFetch(cd_buf, id).rg;
}

vec3 hair_get_customdata_vec3(const samplerBuffer cd_buf)
{
  int id = hair_get_strand_id();
  return texelFetch(cd_buf, id).rgb;
}

vec4 hair_get_customdata_vec4(const samplerBuffer cd_buf)
{
  int id = hair_get_strand_id();
  return texelFetch(cd_buf, id).rgba;
}

vec3 hair_get_strand_pos(void)
{
  int id = hair_get_strand_id() * hairStrandsRes;
  return texelFetch(hairPointBuffer, id).point_position;
}

vec2 hair_get_barycentric(void)
{
  /* To match cycles without breaking into individual segment we encode if we need to invert
   * the first component into the second component. We invert if the barycentricTexCo.y
   * is NOT 0.0 or 1.0. */
  int id = hair_get_base_id();
  return vec2(float((id % 2) == 1), float(((id % 4) % 3) > 0));
}

#endif

/* To be fed the result of hair_get_barycentric from vertex shader. */
vec2 hair_resolve_barycentric(vec2 vert_barycentric)
{
  if (fract(vert_barycentric.y) != 0.0) {
    return vec2(vert_barycentric.x, 0.0);
  }
  else {
    return vec2(1.0 - vert_barycentric.x, 0.0);
  }
}

"""


EMPTY_FRAG_SHADER_SRC = """

void main()
{
  // no color output, only depth (line below is implicit)
  // gl_FragDepth = gl_FragCoord.z;
}

"""


LIB_SRCS = {
  "common_view_lib.glsl": COMMON_VIEW_LIB_SRC,
  "common_hair_lib.glsl": COMMON_HAIR_LIB_SRC,
  "common_lib.glsl": COMMON_LIB_SRC,
  "shader_interface_lib.glsl": SHADER_INTERFACE_LIB_SRC,
  "material_lib.glsl": MATERIAL_LIB_SRC,
  "shading_data_lib.glsl": SHADING_DATA_LIB_SRC,
  "matcap_lib.glsl": MATCAP_LIB_SRC,
  "world_light_lib.glsl": WORLD_LIGHT_LIB_SRC,
}


MESH_VERT_SHADER_SRC = """
#pragma LIB_REQUIRE(common_view_lib.glsl)
#pragma LIB_REQUIRE(shader_interface_lib.glsl)
#pragma LIB_REQUIRE(common_lib.glsl)
#pragma LIB_REQUIRE(material_lib.glsl)

in vec3 pos;
in vec3 nor;
in vec4 ac; /* active color */
in vec2 au; /* active texture layer */

void main()
{
  vec3 world_pos = point_object_to_world(pos);
  gl_Position = point_world_to_ndc(world_pos);

  uv_interp = au;

  normal_interp = normalize(normal_object_to_view(nor));

  float metallic, roughness;

  workbench_material_data_get(color_interp, alpha_interp, roughness, metallic);

  packed_rough_metal = workbench_float_pair_encode(roughness, metallic);
}

"""

MESH_FRAG_SHADER_SRC = """

#pragma LIB_REQUIRE(common_view_lib.glsl)
#pragma LIB_REQUIRE(shader_interface_lib.glsl)
#pragma LIB_REQUIRE(common_lib.glsl)
#pragma LIB_REQUIRE(world_light_lib.glsl)

//layout(location = 0) out vec4 materialData;
//layout(location = 1) out WB_Normal normalData;
out vec4 fragColor;

uniform bool useMatcap = false;

void main()
{
  /* Encode and decode can make normal smooth... */
  WB_Normal normalData = workbench_normal_encode(gl_FrontFacing, normal_interp);
  vec3 N = workbench_normal_decode(vec4(normalData, 0.0, 0.0));

  //materialData = vec4(color_interp, packed_rough_metal);

  //if (useMatcap) {
  //  /* For matcaps, save front facing in alpha channel. */
  //  materialData.a = float(gl_FrontFacing);
  //}
  //vec3 base_color = mat_data.rgb;

  float roughness, metallic;
  workbench_float_pair_decode(packed_rough_metal, roughness, metallic);

#ifdef V3D_LIGHTING_MATCAP
  /* When using matcaps, mat_data.a is the backface sign. */
  N = (mat_data.a > 0.0) ? N : -N;

  fragColor.rgb = get_matcap_lighting(color_interp, N, I);
#endif

#ifdef V3D_LIGHTING_STUDIO
  fragColor.rgb = get_world_lighting(color_interp, roughness, metallic, N, vec3(0, 0, 1));
#endif

#ifdef V3D_LIGHTING_FLAT
  fragColor.rgb = color_interp;
#endif

  // fragColor.rgb *= get_shadow(N);

  fragColor.a = 1.0;
}

"""

HAIR_VERT_SHADER_SRC = """
#pragma LIB_REQUIRE(common_hair_lib.glsl)
#pragma LIB_REQUIRE(common_view_lib.glsl)
#pragma LIB_REQUIRE(shader_interface_lib.glsl)
#pragma LIB_REQUIRE(common_lib.glsl)
#pragma LIB_REQUIRE(material_lib.glsl)

//uniform samplerBuffer ac; /* active color layer */
//uniform samplerBuffer au; /* active texture layer */

/* From http://libnoise.sourceforge.net/noisegen/index.html */
float integer_noise(int n)
{
  n = (n >> 13) ^ n;
  int nn = (n * (n * n * 60493 + 19990303) + 1376312589) & 0x7fffffff;
  return (float(nn) / 1073741824.0);
}

vec3 workbench_hair_random_normal(vec3 tan, vec3 binor, float rand)
{
  /* To "simulate" anisotropic shading, randomize hair normal per strand. */
  vec3 nor = cross(tan, binor);
  nor = normalize(mix(nor, -tan, rand * 0.1));
  float cos_theta = (rand * 2.0 - 1.0) * 0.2;
  float sin_theta = sqrt(max(0.0, 1.0 - cos_theta * cos_theta));
  nor = nor * sin_theta + binor * cos_theta;
  return nor;
}

void workbench_hair_random_material(float rand,
                                    inout vec3 color,
                                    inout float roughness,
                                    inout float metallic)
{
  /* Center noise around 0. */
  rand -= 0.5;
  rand *= 0.1;
  /* Add some variation to the hairs to avoid uniform look. */
  metallic = clamp(metallic + rand, 0.0, 1.0);
  roughness = clamp(roughness + rand, 0.0, 1.0);
  /* Modulate by color intensity to reduce very high contrast when color is dark. */
  color = clamp(color + rand * (color + 0.05), 0.0, 1.0);
}

void main()
{
  bool is_persp = (ProjectionMatrix[3][3] == 0.0);
  float time, thick_time, thickness;
  vec3 world_pos, tan, binor;
  hair_get_pos_tan_binor_time(is_persp,
                              ModelMatrixInverse,
                              ViewMatrixInverse[3].xyz,
                              ViewMatrixInverse[2].xyz,
                              world_pos,
                              tan,
                              binor,
                              time,
                              thickness,
                              thick_time);

  gl_Position = point_world_to_ndc(world_pos);

  float hair_rand = integer_noise(hair_get_strand_id());
  vec3 nor = workbench_hair_random_normal(tan, binor, hair_rand);

  //uv_interp = hair_get_customdata_vec2(au);

  normal_interp = normalize(normal_world_to_view(nor));

  float metallic, roughness;
  workbench_material_data_get(color_interp, alpha_interp, roughness, metallic);
  color_interp = hair_get_shown_color(color_interp);

  /* Hairs have lots of layer and can rapidly become the most prominent surface.
   * So we lower their alpha artificially. */
  alpha_interp *= 0.3;

  workbench_hair_random_material(hair_rand, color_interp, roughness, metallic);

  packed_rough_metal = workbench_float_pair_encode(roughness, metallic);
}

"""


HAIR_REFINE_VERT_SRC = """

/* To be compiled with common_hair_lib.glsl */

out vec4 finalColor;

vec4 get_weights_cardinal(float t)
{
  float t2 = t * t;
  float t3 = t2 * t;
#if defined(CARDINAL)
  float fc = 0.71;
#else /* defined(CATMULL_ROM) */
  float fc = 0.5;
#endif

  vec4 weights;
  /* GLSL Optimized version of key_curve_position_weights() */
  float fct = t * fc;
  float fct2 = t2 * fc;
  float fct3 = t3 * fc;
  weights.x = (fct2 * 2.0 - fct3) - fct;
  weights.y = (t3 * 2.0 - fct3) + (-t2 * 3.0 + fct2) + 1.0;
  weights.z = (-t3 * 2.0 + fct3) + (t2 * 3.0 - (2.0 * fct2)) + fct;
  weights.w = fct3 - fct2;
  return weights;
}

/* TODO(fclem): This one is buggy, find why. (it's not the optimization!!) */
vec4 get_weights_bspline(float t)
{
  float t2 = t * t;
  float t3 = t2 * t;

  vec4 weights;
  /* GLSL Optimized version of key_curve_position_weights() */
  weights.xz = vec2(-0.16666666, -0.5) * t3 + (0.5 * t2 + 0.5 * vec2(-t, t) + 0.16666666);
  weights.y = (0.5 * t3 - t2 + 0.66666666);
  weights.w = (0.16666666 * t3);
  return weights;
}

vec4 interp_data(vec4 v0, vec4 v1, vec4 v2, vec4 v3, vec4 w)
{
  return v0 * w.x + v1 * w.y + v2 * w.z + v3 * w.w;
}

#ifdef TF_WORKAROUND
uniform int targetWidth;
uniform int targetHeight;
uniform int idOffset;
#endif

void main(void)
{
  float interp_time;
  vec4 data0, data1, data2, data3;
  hair_get_interp_attrs(data0, data1, data2, data3, interp_time);

  vec4 weights = get_weights_cardinal(interp_time);
  finalColor = interp_data(data0, data1, data2, data3, weights);

#ifdef TF_WORKAROUND
  int id = gl_VertexID - idOffset;
  gl_Position.x = ((float(id % targetWidth) + 0.5) / float(targetWidth)) * 2.0 - 1.0;
  gl_Position.y = ((float(id / targetWidth) + 0.5) / float(targetHeight)) * 2.0 - 1.0;
  gl_Position.z = 0.0;
  gl_Position.w = 1.0;

  gl_PointSize = 1.0;
#endif
}

"""


HAIR_FRAG_SHADER_SRC = MESH_FRAG_SHADER_SRC

COMP_VERT_SHADER_SRC = """

out vec4 uvcoordsvar;

void main()
{
  int v = gl_VertexID % 3;
  float x = -1.0 + float((v & 1) << 2);
  float y = -1.0 + float((v & 2) << 1);
  gl_Position = vec4(x, y, 1.0, 1.0);
  uvcoordsvar = vec4((gl_Position.xy + 1.0) * 0.5, 0.0, 0.0);
}

"""

COMP_FRAG_SHADER_SRC = """

#pragma LIB_REQUIRE(common_view_lib.glsl)
#pragma LIB_REQUIRE(common_lib.glsl)
#pragma LIB_REQUIRE(world_light_lib.glsl)

uniform sampler2D materialBuffer;
uniform sampler2D normalBuffer;

in vec4 uvcoordsvar;

out vec4 fragColor;

void main()
{
  /* Normal and Incident vector are in viewspace. Lighting is evaluated in viewspace. */
  vec3 I = get_view_vector_from_screen_uv(uvcoordsvar.st);
  vec3 N = workbench_normal_decode(texture(normalBuffer, uvcoordsvar.st));
  vec4 mat_data = texture(materialBuffer, uvcoordsvar.st);

  vec3 base_color = mat_data.rgb;

  float roughness, metallic;
  workbench_float_pair_decode(mat_data.a, roughness, metallic);

#ifdef V3D_LIGHTING_MATCAP
  /* When using matcaps, mat_data.a is the backface sign. */
  N = (mat_data.a > 0.0) ? N : -N;

  fragColor.rgb = get_matcap_lighting(base_color, N, I);
#endif

#ifdef V3D_LIGHTING_STUDIO
  fragColor.rgb = get_world_lighting(base_color, roughness, metallic, N, I);
#endif

#ifdef V3D_LIGHTING_FLAT
  fragColor.rgb = base_color;
#endif

  // fragColor.rgb *= get_shadow(N);

  fragColor.a = 1.0;
}

"""


EDIT_HAIR_VERT_SHADER_SRC = """
#pragma LIB_REQUIRE(common_view_lib.glsl)

uniform int hairStrandsRes = 8;
uniform int colorTexWidth = 1024;
uniform sampler2D strandsColorTex;

in vec3 pos;

out vec4 finalColor;

void main()
{
  gl_Position = point_world_to_ndc(pos);

  int hair_id = gl_VertexID / hairStrandsRes;
  int width = hair_id % colorTexWidth;
  int height = hair_id / colorTexWidth;
  finalColor = texelFetch(strandsColorTex, ivec2(width, height), 0);
}

"""


EDIT_HAIR_FRAG_SHADER_SRC = """

in vec4 finalColor;
out vec4 fragColor;

void main()
{
  if ((finalColor.r + finalColor.g + finalColor.b) > 0) {
    fragColor = vec4(finalColor.rgb, 1.0);
    /* Add a little depth to override color. */
    gl_FragDepth = gl_FragCoord.z - 0.000001f;
  } else {
    discard;
  }
}

"""


SHADER_SRC = [[MESH_VERT_SHADER_SRC, MESH_FRAG_SHADER_SRC],
              [HAIR_VERT_SHADER_SRC, HAIR_FRAG_SHADER_SRC],
              [EDIT_HAIR_VERT_SHADER_SRC, EDIT_HAIR_FRAG_SHADER_SRC],
              [COMP_VERT_SHADER_SRC, COMP_FRAG_SHADER_SRC]]


class DataType:
    MESH = 0
    HAIR = 1
    EDIT_HAIR = 2
    COMP = 3


class ShaderType:
    VERT = 0
    FRAG = 1


def _get_src_deps_libnames(src: str):
    lines = src.splitlines()
    libnames = []
    for l in lines:
        if not l:
            continue
        li = l.find("LIB_REQUIRE(")
        if li != -1:
            ln = l[li + 12:-1]
            if ln in LIB_SRCS:
                llns = _get_src_deps_libnames(LIB_SRCS[ln])
                for lln in llns:
                    if libnames.count(lln) == 0:
                        libnames.append(lln)
            if libnames.count(ln) == 0:
                libnames.append(ln)
        else:
            break
    return libnames


def _get_lib_src_prepended_src(src):
    lns = _get_src_deps_libnames(src)
    if lns:
        for ln in reversed(lns):
            if ln not in LIB_SRCS:
                continue
            src = LIB_SRCS[ln] + src
    return src


def get_shader_src(data_type, shader_type):
    return _get_lib_src_prepended_src(SHADER_SRC[data_type][shader_type])


def get_hair_refine_shader():
    return GLShader(COMMON_HAIR_LIB_SRC + HAIR_REFINE_VERT_SRC,
                    EMPTY_FRAG_SHADER_SRC,
                    defines="#define HAIR_PHASE_SUBDIV\n",
                    tf_type=GL_POINTS,
                    tf_names=("finalColor",))


def get_edit_hair_color_shader():
    pass


class GLShader():
    def __init__(self,
                 vertcode,
                 fragcode,
                 geomcode=None,
                 defines="",
                 tf_type=None,
                 tf_names=None):
        assert vertcode and fragcode

        self._program = glCreateProgram()
        self._status = Buffer(GL_INT, 1)
        self._log = Buffer(GL_BYTE, 5000)
        self._attr_locs = None
        self._uniform_locs = {}
        self._uniform_block_locs = {}
        self._uniform_sampler_locs = {}

        if vertcode:
            pre_defs = "#define GPU_VERTEX_SHADER\n#define IN_OUT out\n" + defines
            self._vert_shader = self._create_shader_from_glsl(
                GL_VERTEX_SHADER, pre_defs + vertcode)

        if fragcode:
            pre_defs = "#define GPU_FRAGMENT_SHADER\n#define IN_OUT in\n" + defines
            self._frag_shader = self._create_shader_from_glsl(
                GL_FRAGMENT_SHADER, pre_defs + fragcode)

        if geomcode:
            pre_defs = "#define GPU_GEOMETRY_SHADER\n" + defines
            self._geom_shader = self._create_shader_from_glsl(
                GL_GEOMETRY_SHADER, pre_defs + geomcode)
        else:
            self._geom_shader = None

        self._tf_type = tf_type
        if tf_names:
            assert not (tf_type is None)
            gl_utils.transform_feedback_names_set(self._program, tf_names)

        glLinkProgram(self._program)
        glGetProgramiv(self._program, GL_LINK_STATUS, self._status)
        if not self._status[0]:
            glGetProgramInfoLog(self._program, 5000, self._status, self._log)
            raise RuntimeError(buffer_to_str(self._log, self._status[0]))

    def _create_shader_from_glsl(self, shader_type, src):
        shader = glCreateShader(shader_type)
        if shader == 0:
            raise RuntimeError("Could not create shader object.")

        glShaderSource(shader, "#version 330\n" + src)
        glCompileShader(shader)

        glGetShaderiv(shader, GL_COMPILE_STATUS, self._status)
        if not self._status[0]:
            glGetShaderInfoLog(shader, 5000, self._status, self._log)
            glDeleteShader(shader)
            raise RuntimeError(buffer_to_str(self._log, self._status[0]))

        glAttachShader(self._program, shader)
        return shader

    def bind(self):
        assert self._program != 0
        glUseProgram(self._program)
        return self # for ContextManager

    def unbind(self):
        glUseProgram(0)

    def transform_feedback_enable(self, vbo_id):
        assert vbo_id != 0
        if self._tf_type is None:
            return

        glBindBufferBase(GL_TRANSFORM_FEEDBACK_BUFFER, 0, vbo_id)
        gl_utils.begin_transform_feedback(self._tf_type)

    def transform_feedback_disable(self):
        gl_utils.end_transform_feedback()

    def get_attrib_locations(self):
        if self._attr_locs is None:
            self._attr_locs = {}
            glGetProgramiv(self._program, GL_ACTIVE_ATTRIBUTES, self._status)
            attr_len = self._status[0]
            sizebf = Buffer(GL_INT, 1)
            name_lenbf = Buffer(GL_INT, 1)
            for i in range(attr_len):
                name_lenbf[0] = 0
                glGetActiveAttrib(self._program, i, 1000, name_lenbf, sizebf, self._status, self._log)
                name = buffer_to_str(self._log, name_lenbf[0])
                loc = glGetAttribLocation(self._program, name)
                if loc != -1:
                    self._attr_locs[name] = loc

        return self._attr_locs

    def get_uniform_location(self, name):
        if name in self._uniform_locs:
            return self._uniform_locs[name]
        loc = glGetUniformLocation(self._program, name)
        if loc != -1:
            self._uniform_locs[name] = loc
        return loc

    def _get_valid_binding_point(self):
        i = 10
        exist_points = {v[1] for v in self._uniform_sampler_locs.values()}
        while i in exist_points:
            i += 1
        return i

    def get_uniform_sampler_binding(self, name):
        if name in self._uniform_sampler_locs:
            return self._uniform_sampler_locs[name]
        loc = self.get_uniform_location(name)
        if loc < 16:
            binding = loc
        else:
            binding = self._get_valid_binding_point()
        self._uniform_sampler_locs[name] = (loc, binding)
        return (loc, binding)

    def get_uniform_block_index(self, name):
        if name in self._uniform_block_locs:
            return self._uniform_block_locs[name]
        idx = glGetUniformBlockIndex(self._program, name)
        if idx != -1:
            glUniformBlockBinding(self._program, idx, idx)
            self._uniform_block_locs[name] = idx
        return idx

    def uniform_float(self, name, value):
        loc = self.get_uniform_location(name)
        if loc == -1:
            return

        values = value
        t = type(value)
        if t is float or t is int:
            values = [float(value)]
        elif t is Matrix:
            values = mat_to_list(value)
        elif t is Vector:
            values = value.to_tuple()
        elif t is Color:
            values = [value.r, value.g, value.b]
        elif t is Quaternion:
            values = [value.w, value.x, value.y, value.z]

        self._uniform_float_vector(loc, values)

    def _uniform_float_vector(self, loc, values):
        l = len(values)
        if l == 1:
            glUniform1f(loc, values[0])
        elif l == 2:
            glUniform2f(loc, *values)
        elif l == 3:
            glUniform3f(loc, *values)
        elif l == 4:
            glUniform4f(loc, *values)
        else:
            buf = Buffer(GL_FLOAT, l, values)
            if l == 9:
                glUniformMatrix3fv(loc, 1, False, buf)
            elif l == 16:
                glUniformMatrix4fv(loc, 1, False, buf)

    def uniform_block(self, name, ubo):
        assert type(ubo) is gpu.types.GPUUniformBuf
        idx = self.get_uniform_block_index(name)
        if idx != -1:
            gl_utils.bind_uniform_buffer(ubo, idx)

    def uniform_texture(self, name, tex):
        loc, binding = self.get_uniform_sampler_binding(name)
        if loc != -1:
            tex.bind(binding)
            glUniform1i(loc, binding)

    def uniform_int(self, name, val):
        loc = self.get_uniform_location(name)
        if loc != -1:
            glUniform1i(loc, val)
    
    def uniform_bool(self, name, val):
        self.uniform_int(name, int(val))

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        if traceback:
            print(traceback)

    def __del__(self):
        glDeleteShader(self._vert_shader)
        if self._geom_shader:
            glDeleteShader(self._geom_shader)
        glDeleteShader(self._frag_shader)
        glDeleteProgram(self._program)


if __name__ == "__main__":
    print(get_shader_src(DataType.HAIR, ShaderType.VERT))
