SimpleJavaEngine/src/main/java/speiger/src/coreengine/rendering/textures/custom/TextureAtlas.java

package speiger.src.coreengine.rendering.textures.custom;

import java.util.Iterator;
import java.util.List;
import java.util.Set;
import java.util.function.BiConsumer;

import speiger.src.collections.objects.lists.ObjectArrayList;
import speiger.src.collections.objects.maps.impl.hash.Object2ObjectLinkedOpenHashMap;
import speiger.src.collections.objects.maps.interfaces.Object2ObjectMap;
import speiger.src.collections.objects.misc.pairs.ObjectObjectPair;
import speiger.src.collections.objects.sets.ObjectOpenHashSet;
import speiger.src.collections.objects.utils.ObjectIterators;
import speiger.src.coreengine.assets.AssetLocation;
import speiger.src.coreengine.math.vector.ints.Vec2i;

/**
 * Inspired by: <a href=https://github.com/lukaszdk/texture-atlas-generator/blob/master/AtlasGenerator.java>AtlasGenerator</a>
 */
public class TextureAtlas
{
	int width;
	int height;
	Object2ObjectMap<AssetLocation, AtlasEntry> textures;
	
	protected TextureAtlas(Vec2i bounds, Object2ObjectMap<AssetLocation, AtlasEntry> textures)
	{
		width = bounds.getX();
		height = bounds.getY();
		this.textures = textures;
	}
	
	public int getWidth()
	{
		return width;
	}
	
	public int getHeight()
	{
		return height;
	}
	
	public AtlasEntry getTexture(AssetLocation texture)
	{
		return textures.getObject(texture);
	}
	
	public Iterator<AtlasEntry> getContents()
	{
		return ObjectIterators.unmodifiable(textures.values().iterator());
	}
	
	public static Builder create() 
	{
		return new Builder();
	}
	
	public static class Builder
	{
		Set<AssetLocation> names = new ObjectOpenHashSet<>();
		List<Record> records = new ObjectArrayList<>();
		int pixelsUsed = 0;
		
		public boolean add(AssetLocation location, int width, int height)
		{
			return add(location, width, height, 0);
		}
		
		public boolean add(AssetLocation location, int width, int height, int padding)
		{
			if(location == null || width <= 0 || height <= 0 || padding < 0 || !names.add(location)) return false;
			records.add(new Record(location, width, height, padding));
			pixelsUsed += (width + padding) * (height + padding);
			return true;
		}
		
		private ObjectObjectPair<Vec2i, Object2ObjectMap<AssetLocation, AtlasEntry>> stitch()
		{
			int textureWidth = 2;
			int textureHeight = 2;
			boolean height = false;
			for(;textureHeight * textureWidth <= pixelsUsed;height = !height)
			{
				if(height) textureHeight *= 2;
				else textureWidth *= 2;
			}
			records.sort(null);
			int attempts = 0;
			while(attempts < 50)
			{
				Slot slot = new Slot(0, 0, textureWidth, textureHeight);
				boolean failed = false;
				for(int i = 0,m=records.size();i<m;i++)
				{
					if(slot.addRecord(records.get(i))) continue;
					failed = true;
					break;
				}
				if(failed)
				{
					if(height) textureHeight *= 2;
					else textureWidth *= 2;
					attempts++;
					height = !height;
					continue;
				}
				Object2ObjectMap<AssetLocation, AtlasEntry> entries = new Object2ObjectLinkedOpenHashMap<>();
				slot.build(entries::put);
				return ObjectObjectPair.of(Vec2i.newVec(textureWidth, textureHeight), entries);
			}
			throw new IllegalStateException("Couldn't fit Texture Atlas after growing it 5 Times");
		}
		
		public void buildHollow(BiConsumer<Vec2i, Iterable<AtlasEntry>> builder)
		{
			ObjectObjectPair<Vec2i, Object2ObjectMap<AssetLocation, AtlasEntry>> pairs = stitch();
			builder.accept(pairs.getKey(), pairs.getValue().values());
		}
		
		public TextureAtlas build(BiConsumer<Vec2i, Iterable<AtlasEntry>> builder)
		{
			ObjectObjectPair<Vec2i, Object2ObjectMap<AssetLocation, AtlasEntry>> pairs = stitch();
			builder.accept(pairs.getKey(), pairs.getValue().values());
			return new TextureAtlas(pairs.getKey(), pairs.getValue());
		}
		
		public TextureAtlas build()
		{
			ObjectObjectPair<Vec2i, Object2ObjectMap<AssetLocation, AtlasEntry>> pairs = stitch();
			return new TextureAtlas(pairs.getKey(), pairs.getValue());
		}
	}
	
	public static class AtlasEntry
	{
		AssetLocation location;
		int x;
		int y;
		int width;
		int height;
		
		private AtlasEntry(Slot slot)
		{
			Record record = slot.record;
			location = record == null ? null : record.getLocation();
			x = slot.x;
			y = slot.y;
			width = record == null ? slot.width : record.width;
			height = record == null ? slot.height : record.height;
		}
		
		public AssetLocation getLocation()
		{
			return location;
		}
		
		public int getX()
		{
			return x;
		}
		
		public int getY()
		{
			return y;
		}
		
		public int getWidth()
		{
			return width;
		}
		
		public int getHeight()
		{
			return height;
		}
	}
	
	private static class Slot
	{
		int x;
		int y;
		int width;
		int height;
		Record record;
		Slot[] children = null;
		
		public Slot(int x, int y, int width, int height)
		{
			this.x = x;
			this.y = y;
			this.width = width;
			this.height = height;
		}
		
		public boolean isLeaf()
		{
			return children == null;
		}
		
		public void build(BiConsumer<AssetLocation, AtlasEntry> builder)
		{
			if(!isLeaf()) 
			{
				for(int i = 0,m=children.length;i<m;i++)
				{
					children[i].build(builder);
				}
				return;
			}
			if(record == null) return;
			builder.accept(record.getLocation(), new AtlasEntry(this));
		}
		
		public boolean addRecord(Record record)
		{
			if(isLeaf())
			{
				int rw = record.getWidth();
				int rh = record.getHeight();
				if(this.record != null || rw > width || rh > height) return false;
				if(rw == width && rh == height)
				{
					this.record = record;
					return true;
				}
				int p = record.getPadding();
				int dw = width - rw;
				int dh = height - rh;
				children = new Slot[dw > 0 && dh > 0 ? 3 : 2];
				children[0] = new Slot(x, y, rw, rh);
				if(dw > 0 && dh > 0)
				{
					if(dw > dh)
					{
						children[1] = new Slot(x + rw + p, y, dw - p, rh);
						children[2] = new Slot(x, y + rh + p, width, dh - p);
					}
					else
					{
						children[1] = new Slot(x, y + rh + p, rw, dh - p);
						children[2] = new Slot(x + rw + p, y, dw - p, height);
					}
				}
				else if(dw == 0) children[1] = new Slot(x, y + rh + p, rw, dh - p);
				else if(dh == 0) children[1] = new Slot(x + rw + p, y, dw - p, rh);
				return children[0].addRecord(record);
			}
			for(int i = 0,m=children.length;i<m;i++)
			{
				if(children[i].addRecord(record)) return true;
			}
			return false;
		}
	}
	
	private static class Record implements Comparable<Record>
	{
		AssetLocation location;
		int width;
		int height;
		int padding;
		
		public Record(AssetLocation location, int width, int height, int padding)
		{
			this.location = location;
			this.width = width;
			this.height = height;
			this.padding = padding;
		}
		
		public int getHeight()
		{
			return height;
		}
		
		public int getWidth()
		{
			return width;
		}
		
		public int getPadding()
		{
			return padding;
		}
		
		public AssetLocation getLocation()
		{
			return location;
		}
		
		private int getPixelCount()
		{
			return (width + padding) * (height + padding);
		}
		
		@Override
		public int compareTo(Record o)
		{
			int result = Integer.compare(o.getPixelCount(), getPixelCount());
			return result == 0 ? getLocation().compareTo(o.getLocation()) : result;
		}
	}
}