using Unity.Collections;
using Unity.Jobs;
using UnityEngine;
using System.Diagnostics;
using System.Collections.Generic;
using Unity.Burst;

public struct JobNoBurst : IJobParallelFor {
    [ReadOnly]
    public NativeArray<float> a;
    [ReadOnly]
    public NativeArray<float> b;
    public NativeArray<float> result;
    public void Execute(int i) {
        result[i] = a[i] * b[i];
    }
}

[BurstCompile]
public struct JobWithBurst : IJobParallelFor {
    [ReadOnly]
    public NativeArray<float> a;
    [ReadOnly]
    public NativeArray<float> b;
    public NativeArray<float> result;
    public void Execute(int i) {
        result[i] = a[i] * b[i];
    }
}

public class JobSystemTest : MonoBehaviour {
    public int NumberOfTestsToRun = 10;
    public int NumberOfArraysElements = 1000000;
    public int BatchSize = 32;
    public float ValueOfArrayA = 0.5f;
    public float ValueOfArrayB = 1.5f;
    private List<long> NoBurstResults = new List<long>();
    private List<long> WithBurstResults = new List<long>();

    private void Awake() {
        QualitySettings.vSyncCount = 0;
    }

    private void Update() {
        if (Input.GetKeyDown(KeyCode.Space)) {
            ParallelArrayOperation();
        }
    }

    private void ParallelArrayOperation() {
        UnityEngine.Debug.Log($"Summing two arrays of {NumberOfArraysElements} elements for {NumberOfTestsToRun} times");
        NoBurstResults.Clear();
        WithBurstResults.Clear();
        for (int i = 0; i < NumberOfTestsToRun; i++) {
            ExecuteJob<JobNoBurst>();
        }
        for (int i = 0; i < NumberOfTestsToRun; i++) {
            ExecuteJob<JobWithBurst>();
        }
        long averageNoBurst = 0;
        long averageWithBurst = 0;
        for (int i = 0; i < NumberOfTestsToRun; i++) {
            averageNoBurst += NoBurstResults[i];
            averageWithBurst += WithBurstResults[i];
        }
        double averageMsNoBurst = (double)(averageNoBurst * 1000L) / (NumberOfTestsToRun * Stopwatch.Frequency);
        double averageMsWithBurst = (double)(averageWithBurst * 1000L) / (NumberOfTestsToRun * Stopwatch.Frequency);

        UnityEngine.Debug.Log($"Average execution time - BURST OFF: {averageMsNoBurst}");
        UnityEngine.Debug.Log($"Average execution time - BURST ON: {averageMsWithBurst}");

        float percentageDiff = (float)(((averageMsNoBurst / averageMsWithBurst) - 1) * 100);
        UnityEngine.Debug.Log($"Burst compiling was {percentageDiff}% faster");
    }

    private void ExecuteJob<T>() where T : IJobParallelFor {
        var nativeA = new NativeArray<float>(NumberOfArraysElements, Allocator.TempJob);
        var nativeB = new NativeArray<float>(NumberOfArraysElements, Allocator.TempJob);
        var nativeResult = new NativeArray<float>(NumberOfArraysElements, Allocator.TempJob);
        for (int i = 0; i < nativeA.Length; i++) {
            nativeA[i] = ValueOfArrayA;
            nativeB[i] = ValueOfArrayB;
        }
        if (typeof(T) == typeof(JobNoBurst)) {
            var job = new JobNoBurst();
            job.a = nativeA;
            job.b = nativeB;
            job.result = nativeResult;
            Stopwatch sw = Stopwatch.StartNew();
            var handle = job.Schedule(nativeResult.Length, BatchSize);
            handle.Complete();
            sw.Stop();
            job.a.Dispose();
            job.b.Dispose();
            job.result.Dispose();
            NoBurstResults.Add(sw.ElapsedTicks);
        }
        else {
            var job = new JobWithBurst();
            job.a = nativeA;
            job.b = nativeB;
            job.result = nativeResult;
            Stopwatch sw = Stopwatch.StartNew();
            var handle = job.Schedule(nativeResult.Length, BatchSize);
            handle.Complete();
            sw.Stop();
            job.a.Dispose();
            job.b.Dispose();
            job.result.Dispose();
            WithBurstResults.Add(sw.ElapsedTicks);
        }
    }
}