using Unity.Burst;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using Unity.Jobs;
using UnityEngine;

// This example code demonstrates using Unity Burst directly on a static method without jobs.

// Unity NativeArrays can't be used directly in Unity Burst methods (only in Burst jobs),
// so we have to pass pointers to arrays instead.
// 
// This PointerArray wrapper can be created from a NativeArray via simple implicit assignment:
// 
// NativeArray<float> floatNativeArray = new NativeArray<float>(1000, Allocator.Persistent);
// PointerArray<float> floatPointerArray = floatNativeArray;
// 
// The PointerArray can then be used like an array itself, including getting the Length.
public unsafe struct PointerArray<T> where T : unmanaged {
	public readonly T* Array;
	public readonly int Length;

	public T this[uint index] {
		get { return Array[index]; }
		set { Array[index] = value; }
	}

	public T this[int index] {
		get { return Array[index]; }
		set { Array[index] = value; }
	}

	public PointerArray(T* Array, int Length) {
		this.Array = Array;
		this.Length = Length;
	}

	public static implicit operator PointerArray<T>(NativeArray<T> a) {
		return new PointerArray<T>((T*)a.GetUnsafePtr(), a.Length);
	}
}

[BurstCompile]
public class BurstMethodTester : MonoBehaviour {
	
	enum Mode { NonBurstMethod, NonBurstJob, BurstMethod, BurstJob }
	
	void Start() {
		Execute(Mode.NonBurstMethod);
		Execute(Mode.NonBurstJob);
		Execute(Mode.BurstMethod);
		Execute(Mode.BurstJob);
	}

	void Execute(Mode mode) {
		var input = new NativeArray<float>(10000000, Allocator.Persistent);
		var output = new NativeArray<float>(1, Allocator.Persistent);
		for (int i = 0; i < input.Length; i++)
			input[i] = 1.0f * i;

		var stopwatch = new System.Diagnostics.Stopwatch();
		stopwatch.Start();

		if (mode == Mode.NonBurstMethod) {
			PointerArray<float> outputArr = output;
			ExecuteNonBurstMethod(input, ref outputArr);
		}
		else if (mode == Mode.NonBurstJob) {
			var job = new MyNonBurstStruct { Input = input, Output = output };
			job.Schedule().Complete();
		}
		else if (mode == Mode.BurstMethod) {
			PointerArray<float> outputArr = output;
			ExecuteBurstMethod(input, ref outputArr);
		}
		else if (mode == Mode.BurstJob) {
			var job = new MyBurstStruct { Input = input, Output = output };
			job.Schedule().Complete();
		}

		stopwatch.Stop();
		Debug.Log("Took " + stopwatch.ElapsedMilliseconds + "ms with mode " + mode
			+ "\nThe result of the sum is: " + output[0]);

		input.Dispose();
		output.Dispose();
	}

	[BurstCompile]
	public static void ExecuteBurstMethod(in PointerArray<float> Input, ref PointerArray<float> Output) {
		float result = 0.0f;
		for (int i = 0; i < Input.Length; i++) {
			result += Input[i];
		}
		Output[0] = result;
	}

	public static void ExecuteNonBurstMethod(in PointerArray<float> Input, ref PointerArray<float> Output) {
		float result = 0.0f;
		for (int i = 0; i < Input.Length; i++) {
			result += Input[i];
		}
		Output[0] = result;
	}

	[BurstCompile(CompileSynchronously = true)]
	private struct MyBurstStruct : IJob {
		[ReadOnly]
		public NativeArray<float> Input;

		[WriteOnly]
		public NativeArray<float> Output;

		public void Execute() {
			float result = 0.0f;
			for (int i = 0; i < Input.Length; i++) {
				result += Input[i];
			}
			Output[0] = result;
		}
	}

	private struct MyNonBurstStruct : IJob {
		[ReadOnly]
		public NativeArray<float> Input;

		[WriteOnly]
		public NativeArray<float> Output;

		public void Execute() {
			float result = 0.0f;
			for (int i = 0; i < Input.Length; i++) {
				result += Input[i];
			}
			Output[0] = result;
		}
	}
}