#include "bounded_collective.hpp"
#include "row.h"
#include "sort_algorithms.hpp"
#include <iostream>

inline uint8_t& PIX(uint8_t* video, int f , int i, int j, int k,int H,int W,int C) {
	return video[(size_t)f* H + i) *W+j)*C+k];
}

BoundedCollective::BoundedCollective(int num_threads,int max_size , uint8_t* video , int width , int height, int channels):
	max_queue_size(max_size) , stop_flag(false) , rows_completed(0), target_rows(0), video_data(video) , W(width) , H(height) , C(channels)
	{
		for(int i=0; i<num_threads; i++) {
			workers.emplace_back(&BoundedCollective::worker_loop, this , i);
	}
}

BoundedCollective::~BoundedCollective() {
	{
		std::lock_guard<std::mutex> lock(q_mtx);
		stop_flag=true;
	}
	cv_not_empty.notify_all();
	for(auto& t : workers) {	
		if(t.joinable()) t.join();
	}
}

void BoundedCollective::worker_loop(int id){
	while(true) {
		Task t;
			{
				std::unique_lock<std::mutex> lock(q_mtx);
				cv_not_empty.wait(lock,[this]() {return !task.queue.empty() || stop_flag;});
				
				if(stop_flag && task_queue.empty()) return;
				
				t=task_queue.front();
				task_queue.pop();
				cv_not_full.notify_one();
			}
	
	int ff=t.frame_id;
	int i=t.row_id;
	
	uint8_t* blueRow=new uint8_t[W];
	uint8_t* greenRow = new uint8_t[W];
	uint8_t* redRow = new uint8_t[W];
	
	for(int j=0; j<W; j++){
		blueRow[j] = PIX(video_data, ff , i , j , 0 , H , W , C);
		greenRow[j]=PIX(video_data,ff,i,j,1,H,W,C);		
		redRow[j]=PIX(video_data,ff,i,j,2,H,W,C);
	}

	row<uint8_t> b(blueRow, W);
	row<uint8_t> g(greenRow, W);
	row<uint8_t> r(redRow, W);

	quick_sort(b);
	quick_sort(g);
	quick_sort(r);

	for (int j=0; j<W; j++) {
		PIX(video_data , ff , i , j , 0, H , W , C) = b[j];
		PIX(video_data , ff , i , j , 1, H , W , C) = g[j];
		PIX(video_data , ff , i , j , 2, H , W , C) = r[j];
		

	}

	delete[] blueRow;
	delete[] greenRow;
	delete[] redRow;
	{
		std::lock_guard<std::mutex> lock(barrier_mtx);
		rows_completed++;

		if(rows_completed == target_rows) {
			cv_frame_done.notify_one();
			}
		}
	}
}