//#include "UnityShaderVariables.cginc"
//#include "Noise.cginc"
#pragma only_renderers d3d12

//#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Macros.hlsl"
#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Common.hlsl"
//#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Packing.hlsl"
//#include "Packages/com.unity.render-pipelines.core/ShaderLibrary/Sampling/Sampling.hlsl"

#pragma max_recursion_depth 1

#define M_PI 3.1415926535897932384626433832795
#define M_PHI 1.618033988749895

// Input
RaytracingAccelerationStructure g_SceneAccelStruct;

//Global inputs 
uint	PointLightCount;
uint	ConeLightCount;
uint	DirLightCount;
uint	AreaLightCount;
uint	AreaLightSamples;
uint	EnvLightSamples;
int	PerDispatchRayCount;
int	StartRayIdx;
float HalfVoxelSize;

struct PointLightData{
float3	PointLightsWS;
float4	PointLightsColor;
};

struct ConeLightData{
float3	ConeLightsWS	;
float4	ConeLightsColor	;
float3	ConeLightsDir	;
float2	ConeLightsPram	; //Outter , inner
};

struct DirLightData{
float3	DirLightsDir;
float4	DirLightsColor;
};

struct AreaLightData{
float4x4	AreaLightsMatrix;
float4x4	AreaLightsMatrixInv;
float3	AreaLightsWS;
float4	AreaLightsColor;
float3	AreaLightsSize;
};

StructuredBuffer<PointLightData> PLD;
StructuredBuffer<ConeLightData> CLD;
StructuredBuffer<DirLightData> DLD;
StructuredBuffer<AreaLightData> ALD;

float3	WPosition;
float3	Size = float3(1,1,1);

// Output
RWTexture3D<float4> g_Output;

//Enviorment
TextureCube  _SkyTexture;
SamplerState  sampler_SkyTexture;

struct RayPayload
{
    float4 color;
    float3 dir;
};

struct EnvRayPayload
{
    float4 color;
    float3 dir;
};

[shader("miss")]
void MainMissShader(inout RayPayload payload : SV_RayPayload)
{
    payload.color = float4(1, 1, 1, 0);
}

[shader("miss")]
void zEnvMissShader(inout EnvRayPayload payload : SV_RayPayload)
{
    //Enviormental    
    payload.color = _SkyTexture.SampleLevel(sampler_SkyTexture, payload.dir, 0);
}

float InverseSquare(float distance){
	return	1 / (4 * M_PI * distance * distance);
}

//float _Seed;

float rand(float2 Pixel)
{
    float _Seed = Pixel.x + Pixel.y ;

	float result = frac(sin(_Seed / 100.0f * dot(Pixel, float2(12.9898f, 78.233f) ) ) * 43758.5453f);
//	_Seed = _Seed + 1.0f;
	return result;
}

float rand(float3 Pixel)
{
    float _Seed = Pixel.x + Pixel.y + Pixel.z ;

	float result = frac(sin(_Seed / 100.0f * dot(Pixel, float3(12.9898f, 49.1165f, 29.1165f))) * 43758.5453f);
//	_Seed += 1.0f;
	return result;
}

//RTgems 3.11 FIBONACCI SPHERE
//current index, number of samples
 float3 sphericalFibonacci(float i, float n) 
 {
     const float PHI = sqrt(5.) * 0.5f + 0.5f;
     float fraction = (i * (PHI - 1)) - floor(i * (PHI - 1));
     float phi = 2.f * M_PI * fraction;
     float cosTheta = 1.f - (2.f * i + 1.f) * (1.f / n);
     float sinTheta = sqrt(saturate(1.f - cosTheta*cosTheta));
    
     return float3(cos(phi) * sinTheta , sin(phi) * sinTheta , cosTheta);
}

///
///Point light
///
float4 PointLightCast(float3 VoxelWS, uint Num){

	float LightRadius = distance(VoxelWS, PLD[Num].PointLightsWS ); 
    float3 rayDirection =  -normalize(VoxelWS - PLD[Num].PointLightsWS)  ;

    RayDesc ray;
    ray.Origin    = VoxelWS;
    ray.Direction = rayDirection;
    ray.TMin      = 0.0f;
    ray.TMax      = LightRadius;

    RayPayload payload;
    payload.color = float4(0, 0, 0, 1);

    uint missShaderIndex = 0;
    TraceRay(g_SceneAccelStruct, 0, 0xFF, 0, 1, missShaderIndex, ray, payload); //Add an anyhit shader to support transparencies 

	return (1-payload.color.a) * InverseSquare(max(LightRadius, HalfVoxelSize)) * PLD[Num].PointLightsColor;
}

///
/// Cone light
///

float4 ConeLightCast(float3 VoxelWS, uint Num){


    float3 rayDirection =  -normalize(VoxelWS - CLD[Num].ConeLightsWS)  ;

	//Currently taking a point light and adding attenuation
	float attenuation = (dot(CLD[Num].ConeLightsDir, -rayDirection));

    if (attenuation <= 0) return float4(0, 0, 0, 0); //early out

    float LightRadius = (distance(VoxelWS, CLD[Num].ConeLightsWS )); 

	/////
	float flOuterConeCos = CLD[Num].ConeLightsPram.x;
	float flTemp = dot(CLD[Num].ConeLightsDir, -rayDirection) - flOuterConeCos;
	float vSpotAtten = saturate(flTemp * CLD[Num].ConeLightsPram.y);
	///

    RayDesc ray;
    ray.Origin    = VoxelWS;
    ray.Direction = rayDirection;
    ray.TMin      = 0.0f;
    ray.TMax      = LightRadius;

    RayPayload payload;
    payload.color = float4(0, 0, 0, 1);

    uint missShaderIndex = 0;
    TraceRay(g_SceneAccelStruct, 0, 0xFF, 0, 1, missShaderIndex, ray, payload); //Add an anyhit shader to support transparencies 

	return  InverseSquare(max(LightRadius, HalfVoxelSize) ) * vSpotAtten * (1-payload.color.a) * CLD[Num].ConeLightsColor;
}

//
//Directional Light
//

float4 DirLightCast(float3 VoxelWS, uint Num){

    float3 rayDirection = -DLD[Num].DirLightsDir;

    RayDesc ray;
    ray.Origin    = VoxelWS;
    ray.Direction = rayDirection;
    ray.TMin      = 0.0f;
    ray.TMax      =  1.#INF;

    RayPayload payload;
    payload.color = float4(0, 0, 0, 1);

    uint missShaderIndex = 0;
    TraceRay(g_SceneAccelStruct, 0, 0xFF, 0, 1, missShaderIndex, ray, payload); //Add an anyhit shader to support transparencies 

	return  (1-payload.color.a) * DLD[Num].DirLightsColor;
}

//
//Area Light
//

float4 AreaLightCast(float3 VoxelWS, uint Num, int startIdx, int endIdx){
  //  
    float3 lsPos = float3(ALD[Num].AreaLightsMatrix[0][3], ALD[Num].AreaLightsMatrix[1][3],ALD[Num].AreaLightsMatrix[2][3]);
    float3 VoxelLS =  mul( float4(VoxelWS.xyz,1)-lsPos, ALD[Num].AreaLightsMatrix);
    float4 areaLightAccumulation = float4(0,0,0,0);
    if (VoxelLS.z <= 0) return areaLightAccumulation; //Early out
    uint3 id =  DispatchRaysIndex().xyz; //redundent, oh well

  for (int j = startIdx; j < endIdx; j++) 
  {
    //int loop64 = fmod(j + AreaLightSamples + id.x+id.y+id.z,64);
    float3 LocalPos = mul(ALD[Num].AreaLightsWS.xyz -lsPos , ALD[Num].AreaLightsMatrix).xyz ;
    float3 LightPosSample = LocalPos + float3( (rand(id.xyz + j)-0.5) *  ALD[Num].AreaLightsSize.x , (rand(id.xyz+j+20)-0.5) * ALD[Num].AreaLightsSize.y,0);
    //float3 LightPosSample = LocalPos + float3( BlueNoiseInDisk[loop64].x *  AreaLightsSize[Num].x * 0.5  , BlueNoiseInDisk[loop64].y * AreaLightsSize[Num].y *0.5, 0   );

    float LightRadius = distance(VoxelLS, LightPosSample ); 

	float3 rayDirection = -normalize(VoxelLS - LightPosSample);

	float attenuation = saturate(dot(float3(0,0,1), -rayDirection));

    RayDesc ray;
    ray.Origin    = VoxelWS;
    ray.Direction = mul(rayDirection, ALD[Num].AreaLightsMatrixInv);
    ray.TMin      = 0.0f;
    ray.TMax      = LightRadius;

    RayPayload payload;
    payload.color = float4(0, 0, 0, 1);

    uint missShaderIndex = 0;
    TraceRay(g_SceneAccelStruct, 0, 0xFF, 0, 1, missShaderIndex, ray, payload); //Add an anyhit shader to support transparencies 

    areaLightAccumulation += saturate(InverseSquare(LightRadius) * (1-payload.color.a) * ALD[Num].AreaLightsColor * attenuation) / AreaLightSamples;
   }
	return  areaLightAccumulation;
}

//
//Disk light
//

float4 DiscLightCast(float3 VoxelWS, uint Num, int startIdx, int endIdx){
  //  
    float3 lsPos = float3(ALD[Num].AreaLightsMatrix[0][3], ALD[Num].AreaLightsMatrix[1][3],ALD[Num].AreaLightsMatrix[2][3]);
    float3 VoxelLS =  mul( float4(VoxelWS.xyz,1)-lsPos, ALD[Num].AreaLightsMatrix);
    float3 IntialDirection =  -normalize(VoxelWS - ALD[Num].AreaLightsWS)  ;
    uint3 id =  DispatchRaysIndex().xyz;
    float4 areaLightAccumulation = float4(0,0,0,0);
    if (VoxelLS.z <= 0) return areaLightAccumulation;

  for (int j = startIdx; j < endIdx; j++) 
  {
    float3 LocalPos = mul(ALD[Num].AreaLightsWS.xyz -lsPos , ALD[Num].AreaLightsMatrix).xyz ;
    //https://stackoverflow.com/questions/5837572/generate-a-random-point-within-a-circle-uniformly
    float t = 2 * M_PI * rand(id.xyz + j + 30);
    float u = rand(id.xyz + j) + rand(id.xyz + j + 20);
    float r;
    if (u > 1) r = (2 - u);
    else  r = u;
    //	[r * cos(t), r * sin(t)]

    float3 LightPosSample = LocalPos +  
        float3( ( r * cos(t)) * ALD[Num].AreaLightsSize.x,
                ( r * sin(t)) * ALD[Num].AreaLightsSize.x,
                0);


    float LightRadius = distance(VoxelLS, LightPosSample ); 

	float3 rayDirection = -normalize(VoxelLS - LightPosSample);

	float attenuation = saturate(dot(float3(0,0,1), -rayDirection));

    RayDesc ray;
    ray.Origin    = VoxelWS;
    ray.Direction = mul(rayDirection, ALD[Num].AreaLightsMatrixInv);
    ray.TMin      = 0.0f;
    ray.TMax      = LightRadius;

    RayPayload payload;
    payload.color = float4(0, 0, 0, 1);

    uint missShaderIndex = 0;
    TraceRay(g_SceneAccelStruct, 0, 0xFF, 0, 1, missShaderIndex, ray, payload); //Add an anyhit shader to support transparencies 

    areaLightAccumulation += saturate(InverseSquare(LightRadius) * (1-payload.color.a) * ALD[Num].AreaLightsColor * attenuation) / AreaLightSamples;
   }
	return  areaLightAccumulation;
}

//Cast Ray from vox to env
float4 EnvCast(float3 VoxelWS, int startIdx, int endIdx){

    float4 Accumulation = float4(0,0,0,0);

       for (int j = startIdx; j < endIdx; j++) 
        {
            uint3 id =  DispatchRaysIndex().xyz;
            float3 rayDirection = sphericalFibonacci(j, EnvLightSamples) + ( normalize(float3(rand(id.xyz+j),rand(id.xyz+j+33) , rand(id.xyz+j+120)) -.5) * 2 * (4*M_PI / EnvLightSamples) )  ;

            RayDesc ray;
            ray.Origin    = VoxelWS;
            ray.Direction = rayDirection;
            ray.TMin      = 0.0f;
            ray.TMax      = 100000;

            RayPayload payload;
            payload.color = float4(0, 0, 0, 0);

            uint missShaderIndex = 1;
            TraceRay(g_SceneAccelStruct, 0, 0xFF, 0, 1, missShaderIndex, ray, payload);
            
            Accumulation += payload.color / EnvLightSamples ;
        }
    return Accumulation;
}

// [shader("closesthit")]
// void MyClosestHit(inout RayPayload data,
// BuiltInTriangleIntersectionAttributes attribs) {
// data.color = float4( 1, 0, 0, 1 );
// }

//
// Main shader
//

[shader("raygeneration")]
void MainRayGenShader()
{
    float3 launchIndex = DispatchRaysIndex().xyz ; //id
    float3 launchDim = DispatchRaysDimensions().xyz; //whd
    float3 VoxelWorldPosition = WPosition + Size * ( (launchIndex + 0.5 ) / launchDim );

	float4 LightAccumulation = g_Output[launchIndex];

	int start = StartRayIdx;
	int end = StartRayIdx + PerDispatchRayCount;
	int maxIter = min(PointLightCount, end);
	for (int i = start; i < maxIter; i++) LightAccumulation += PointLightCast(VoxelWorldPosition, i);
	
	start = max(0, start - (int)PointLightCount);
	end -= PointLightCount;
	maxIter = min(ConeLightCount, end);
	for (int i = start; i < maxIter; i++)  LightAccumulation += ConeLightCast(VoxelWorldPosition, i); 

	start = max(0, start - (int)ConeLightCount);
	end -= ConeLightCount;
	maxIter = min(DirLightCount, end);
	for (int i = start; i < maxIter; i++)   LightAccumulation += DirLightCast(VoxelWorldPosition, i);

	start = max(0, start - (int)DirLightCount);
	end -= (int)DirLightCount;
	maxIter = min(EnvLightSamples, end);
    LightAccumulation += EnvCast(VoxelWorldPosition, start, maxIter);

	start = max(0, start - (int)EnvLightSamples);
	end -= (int)EnvLightSamples;
	maxIter = min(AreaLightSamples, end);
	int light = 0;
	while (end > 0)
	{

		if (ALD[light].AreaLightsSize.z == 0)  LightAccumulation += AreaLightCast(VoxelWorldPosition, light, start, maxIter); 
		else  LightAccumulation += DiscLightCast(VoxelWorldPosition, light, start, maxIter);

		start = max(0, start - (int)AreaLightSamples);
		end -= (int)AreaLightSamples;
		maxIter = min(AreaLightSamples, end);
		light++;
    }


  //  LightAccumulation.a = 1; //filling alpha
    g_Output[launchIndex] = LightAccumulation;
}