require 'torch';
require 'sys';
require 'image';
local dir = require 'pl.dir';
local ffi = require 'ffi';

function getImagePaths(folder)
	-- obtain list of image files
	local classes = {}
	local classPaths = {}
	local dirs = dir.getdirectories(folder);
	for k,dirpath in ipairs(dirs) do
		local class = paths.basename(dirpath)
		table.insert(classes, class)
		table.insert(classPaths, dirpath)
	end

	-- define command-line tools, try your best to maintain OSX compatibility
	local wc = 'wc'
	local cut = 'cut'
	local find = 'find'
	if ffi.os == 'OSX' then
		wc = 'gwc'
		cut = 'gcut'
		find = 'gfind'
	end

	-- options for the GNU find command
	local extensionList	 = {'jpg', 'JPG', 'png', 'PNG', 'jpeg', 'JPEG', 'ppm', 'PPM', 'bmp', 'BMP'}
	local findOptions = ' -iname "*.' .. extensionList[1] .. '"'
	for i=2,#extensionList do
		findOptions = findOptions .. ' -o -iname "*.' .. extensionList[i] .. '"'
	end

	-- find the image path names
	local imagePaths = torch.CharTensor()	-- path to each image in dataset
	local imageClass = torch.LongTensor()	-- class index of each image (class index in self.classes)
	local classList = {}			-- index of imageList to each image of a particular class

	-- create file listing the paths to every image
	local classFindFiles = {}
	for i=1,#classes do
		classFindFiles[i] = os.tmpname()
	end
	local combinedFindList = os.tmpname()

	local tmpfile = os.tmpname()
	local tmphandle = assert(io.open(tmpfile, 'w'))
	for i,class in ipairs(classes) do
		local command = find .. ' "' .. classPaths[i] .. '" ' .. findOptions .. ' >>"' .. classFindFiles[i] .. '" \n'
		tmphandle:write(command)
	end
	io.close(tmphandle)
	os.execute('bash ' .. tmpfile)
	os.execute('rm -f ' .. tmpfile)

	local tmpfile = os.tmpname()
	local tmphandle = assert(io.open(tmpfile, 'w'))
	-- concat all finds to a single large file in the order of self.classes
	for i=1,#classes do
		local command = 'cat "' .. classFindFiles[i] .. '" >>' .. combinedFindList .. ' \n'
		tmphandle:write(command)
	end
	io.close(tmphandle)
	os.execute('bash ' .. tmpfile)
	os.execute('rm -f ' .. tmpfile)

	local maxPathLength = tonumber(sys.fexecute(wc .. " -L '" .. combinedFindList .. "' |" .. cut .. " -f1 -d' '")) + 1
	local length = tonumber(sys.fexecute(wc .. " -l '" .. combinedFindList .. "' |" .. cut .. " -f1 -d' '"))

	imagePaths:resize(length, maxPathLength):fill(0)
	local s_data = imagePaths:data()
	for line in io.lines(combinedFindList) do
		ffi.copy(s_data, line)
		s_data = s_data + maxPathLength
	end
	numSamples = imagePaths:size(1)
	print(numSamples ..  ' samples found.')

	imageClass:resize(numSamples)
	local runningIndex = 0
	for i=1,#classes do
		local length = tonumber(sys.fexecute(wc .. " -l '" .. classFindFiles[i] .. "' |" .. cut .. " -f1 -d' '"))
		classList[i] = torch.linspace(runningIndex + 1, runningIndex + length, length):long()
		imageClass[{{runningIndex + 1, runningIndex + length}}]:fill(i)
		runningIndex = runningIndex + length
	end

	local tmpfilelistall = ''
	for i=1,#(classFindFiles) do
		tmpfilelistall = tmpfilelistall .. ' "' .. classFindFiles[i] .. '"'
		if i % 1000 == 0 then
			os.execute('rm -f ' .. tmpfilelistall)
			tmpfilelistall = ''
		end
	end
	os.execute('rm -f '  .. tmpfilelistall)
	os.execute('rm -f "' .. combinedFindList .. '"')

	return classes, classList, imagePaths
end

function getSample(classes, sampleList, imagePaths)
	dataTable = {}
	scalarTable = {}
	N = 0
	for i=1,#classes do
		for j=1,sampleList[i]:nElement() do
			local imgpath = ffi.string(torch.data(imagePaths[sampleList[i][j]]))
			out = image.load(imgpath, 3, 'float')
			table.insert(dataTable, out)
			table.insert(scalarTable, i)
			N = N + 1
		end
	end
	data = torch.Tensor(N, 3, 101, 101)
	scalarLabels = torch.LongTensor(N):fill(-1111)
	for i=1,#dataTable do
		data[i]:copy(dataTable[i])
		scalarLabels[i] = scalarTable[i]
	end
	dataset = {}
	dataset.data = data
	dataset.label = scalarLabels

	setmetatable(dataset,
		{__index = function(t, i) 
						return {t.data[i], t.label[i]}
					end}
	);

	function dataset:size() 
		return self.data:size(1) 
	end

	-- data normalization
	mean = {}
	stdv  = {}
	for i=1,3 do
		mean[i] = dataset.data[{ {}, {i}, {}, {}  }]:mean()
		dataset.data[{ {}, {i}, {}, {}  }]:add(-mean[i])
		
		stdv[i] = dataset.data[{ {}, {i}, {}, {}  }]:std()
		if stdv[i] ~= 0 then
			dataset.data[{ {}, {i}, {}, {}  }]:div(stdv[i])
		end
	end

	return dataset
end

function getRandomSample(classes, batchSize, classList, imagePaths)
	dataTable = {}
	scalarTable = {}
	N = 0
	for i=1,#classes do
		for j=1,batchSize do
			local index = math.max(1, math.ceil(torch.uniform() * classList[i]:nElement()))
			local imgpath = ffi.string(torch.data(imagePaths[classList[i][index]]))
			out = image.load(imgpath, 3, 'float')
			table.insert(dataTable, out)
			table.insert(scalarTable, i)
			N = N + 1
		end
	end
	data = torch.Tensor(N, 3, 101, 101)
	scalarLabels = torch.LongTensor(N):fill(-1111)
	for i=1,#dataTable do
		data[i]:copy(dataTable[i])
		scalarLabels[i] = scalarTable[i]
	end
	dataset = {}
	dataset.data = data
	dataset.label = scalarLabels

	setmetatable(dataset,
		{__index = function(t, i) 
						return {t.data[i], t.label[i]}
					end}
	);

	function dataset:size() 
		return self.data:size(1) 
	end

	-- data normalization
	mean = {}
	stdv  = {}
	for i=1,3 do
		mean[i] = dataset.data[{ {}, {i}, {}, {}  }]:mean()
		dataset.data[{ {}, {i}, {}, {}  }]:add(-mean[i])
		
		stdv[i] = dataset.data[{ {}, {i}, {}, {}  }]:std()
		if stdv[i] ~= 0 then
			dataset.data[{ {}, {i}, {}, {}  }]:div(stdv[i])
		end
	end

	return dataset
end

function getBatchSizes(classes, classList, batchSize)
	local numSamples = 0
	for i=1,#classes do
		numSamples = numSamples + classList[i]:nElement()
	end
	local numBatches = math.ceil(numSamples/batchSize)

	local batchSizes = {}
	for i=1,#classes do
		local roundFlag = 0
		local batchSum = 0
		batchSizes[i] = {}
		for j=1,numBatches-1 do
			if roundFlag == 0 then
				batchSizes[i][j] = math.floor(classList[i]:nElement()/numBatches)
			else
				batchSizes[i][j] = math.ceil(classList[i]:nElement()/numBatches)
			end
			
			batchSum = batchSum + batchSizes[i][j]
			
			if j*classList[i]:nElement()/numBatches > batchSum then
				roundFlag = 1
			else
				roundFlag = 0
			end
		end
		batchSizes[i][numBatches] = classList[i]:nElement() - batchSum
	end
	
	return batchSizes, numBatches, numSamples
end

function shuffleImages(classList, classes)
	local temp = {}
	for i=1,#classes do
		local perm = torch.randperm(classList[i]:size(1))
		temp[i] = torch.LongTensor(classList[i]:size(1))
		for j=1,classList[i]:size(1) do
			temp[i][j] = classList[i][perm[j]]
		end
	end
	return temp
end