#!/usr/bin/python3

# Copyright (C) 2009-2018, Panagiotis Christopoulos Charitos and contributors.
# All rights reserved.
# Code licensed under the BSD License.
# http://www.anki3d.org/LICENSE

import os
import argparse

# Template
template_xml = """<?xml version="1.0" encoding="UTF-8" ?>
<!-- This file is auto generated by generate_program_permutations.py -->
<shaderProgram>
	<shaders>
		<shader>
			<type>vert</type>
			<inputs>
%vertInputs%
			</inputs>
			<source><![CDATA[#include "shaders/MsCommonVert.glsl"
void main() {
%vertSrc%
}
			]]></source>
		</shader>
		<shader>
			<type>frag</type>
			<inputs>
%fragInputs%
			</inputs>
			<source><![CDATA[#include "shaders/MsCommonFrag.glsl"
void main() {
%fragSrc%
}
			]]></source>
		</shader>
	</shaders>
</shaderProgram>
"""

# Input type
INPUT_NONE = 0
INPUT_TEXTURE = 1
INPUT_CONST = 2
INPUT_VALUE = 3 # non-const
INPUT_COUNT = 4

# Variables
DIFFUSE = 0
SPECULAR = 1
ROUGHNESS = 2
METALLIC = 3
NORMAL = 4
EMISSION = 5
SUBSURFACE = 6
HEIGHT = 7
VARIABLE_COUNT = 8

allowed_permutations = [
	[INPUT_TEXTURE, INPUT_CONST],
	[INPUT_TEXTURE, INPUT_CONST],
	[INPUT_TEXTURE, INPUT_CONST],
	[INPUT_TEXTURE, INPUT_CONST],
	[INPUT_TEXTURE, INPUT_NONE],
	[INPUT_TEXTURE, INPUT_CONST],
	[INPUT_CONST],
	[INPUT_NONE, INPUT_TEXTURE]]

def parse_commandline():
	""" Parse the command line arguments """

	parser = argparse.ArgumentParser(description = "Create shader program permutations",
			formatter_class = argparse.ArgumentDefaultsHelpFormatter)

	parser.add_argument("-o", "--output-dir", required = True, help = "specify the output directory")

	args = parser.parse_args()

	return args.output_dir

def spin_wheel(wheel):
	""" Spin the wheel """

	done = False
	start = VARIABLE_COUNT - 1

	while True:
		if start == -1:
			done = True
			break

		wheel[start] += 1
		if wheel[start] >= INPUT_COUNT:
			wheel[start] = INPUT_NONE
			start = start - 1
		else:
			break

	return done

def mutate(wheel):
	""" Create the mutation's source """
	
	parallax = wheel[HEIGHT] == INPUT_TEXTURE

	# Vert input & source
	if parallax:
		vert_inputs = """				<input><type>mat4</type><name>mvp</name></input>
				<input><type>mat3</type><name>normalMat</name><depth>0</depth></input>
				<input><type>mat4</type><name>modelViewMat</name><depth>0</depth></input>
"""

		vert_src = """#if COLOR
positionUvNormalTangent(mvp, normalMat);
parallax(vodelViewMat);
#else
ANKI_WRITE_POSITION(mvp * vec4(in_position, 1.0));
#endif
"""
	else:
		vert_inputs = """				<input><type>mat4</type><name>mvp</name></input>
				<input><type>mat3</type><name>normalMat</name><depth>0</depth></input>
"""

		vert_src = """#if COLOR
positionUvNormalTangent(mvp, normalMat);
#else
ANKI_WRITE_POSITION(mvp * vec4(in_position, 1.0));
#endif
"""

	# Frag inputs & source & filename
	frag_ins = ""
	frag_src = "#if COLOR\n"
	fname = "ms_"

	if parallax:
		frag_ins += "\t\t\t\t<input><type>float</type><name>heightMapScale</name><const>1</const><depth>0</depth></input>\n"

		frag_src += "vec2 uv = computeTextureCoordParallax(heightTex, in_uv, heightMapScale);\n"
	else:
		frag_src += "vec2 uv = in_uv;\n"

	if wheel[DIFFUSE] == INPUT_TEXTURE:
		frag_ins += "\t\t\t\t<input><type>sampler2D</type><name>diffTex</name><depth>0</depth></input>\n"

		frag_src += "vec3 diffColor = texture(diffTex, uv).rgb;\n"

		fname += "difft_"
	elif wheel[DIFFUSE] == INPUT_CONST:
		frag_ins += "\t\t\t\t<input><type>vec3</type><name>diffColor</name><const>1</const><depth>0</depth></input>\n"

		fname += "diffc_"
	else:
		assert 0

	if wheel[SPECULAR] == INPUT_TEXTURE:
		frag_ins += "\t\t\t\t<input><type>sampler2D</type><name>specTex</name><depth>0</depth></input>\n"

		frag_src += "vec3 specColor = texture(specTex, uv).rgb;\n"

		fname += "spect_"
	elif wheel[SPECULAR] == INPUT_CONST:
		frag_ins += "\t\t\t\t<input><type>vec3</type><name>specColor</name><const>1</const><depth>0</depth></input>\n"

		fname += "specc_"
	else:
		assert 0

	if wheel[ROUGHNESS] == INPUT_TEXTURE:
		frag_ins += "\t\t\t\t<input><type>sampler2D</type><name>roughnessTex</name><depth>0</depth></input>\n"

		frag_src += "float roughness = texture(roughnessTex, uv).r;\n"

		fname += "rought_"
	elif wheel[ROUGHNESS] == INPUT_CONST:
		frag_ins += "\t\t\t\t<input><type>float</type><name>roughness</name><const>1</const><depth>0</depth></input>\n"

		fname += "roughc_"
	else:
		assert 0
		
	if wheel[METALLIC] == INPUT_TEXTURE:
		frag_ins += "\t\t\t\t<input><type>sampler2D</type><name>metallicTex</name><depth>0</depth></input>\n"

		frag_src += "float metallic = texture(metallicTex, uv).r;\n"

		fname += "metalt_"
	elif wheel[METALLIC] == INPUT_CONST:
		frag_ins += "\t\t\t\t<input><type>float</type><name>metallic</name><const>1</const><depth>0</depth></input>\n"

		fname += "metalc_"
	else:
		assert 0

	if wheel[NORMAL] == INPUT_TEXTURE:
		frag_ins += "\t\t\t\t<input><type>sampler2D</type><name>normalTex</name><depth>0</depth></input>\n"

		frag_src += "vec3 normal = readNormalFromTexture(normalTex, uv);\n"

		fname += "normalt_"
	elif wheel[NORMAL] == INPUT_NONE:
		frag_src += "vec3 normal = in_normal;\n"

		fname += "normal0_"
	else:
		assert 0

	if wheel[EMISSION] == INPUT_TEXTURE:
		frag_ins += "\t\t\t\t<input><type>sampler2D</type><name>emissionTex</name><depth>0</depth></input>\n"

		frag_src += "vec3 emission = texture(emissionTex, uv).rgb;\n"

		fname += "emist_"
	elif wheel[EMISSION] == INPUT_CONST:
		frag_ins += "\t\t\t\t<input><type>vec3</type><name>emission</name><const>1</const><depth>0</depth></input>\n"

		fname += "emisc_"
	else:
		assert 0

	if wheel[SUBSURFACE] == INPUT_CONST:
		frag_ins += "\t\t\t\t<input><type>float</type><name>subsurface</name><const>1</const><depth>0</depth></input>\n"

		fname += "subsc_"
	else:
		assert 0

	if wheel[HEIGHT] == INPUT_NONE:
		# nothing
		fname += "par0"
	elif wheel[HEIGHT] == INPUT_TEXTURE:
		frag_ins += "\t\t\t\t<input><type>sampler2D</type><name>heightTex</name><depth>0</depth></input>\n"

		fname += "par1"
	else:
		assert 0

	frag_src += "writeRts(diffColor, normal, specColor, roughness, subsurface, emission, metallic);\n"
	frag_src += "#endif\n"

	xml = template_xml
	xml = xml.replace("%vertInputs%", vert_inputs)
	xml = xml.replace("%vertSrc%", vert_src)
	xml = xml.replace("%fragInputs%", frag_ins)
	xml = xml.replace("%fragSrc%", frag_src)

	return (xml, fname)

def main():
	""" Main function """

	out_dir = parse_commandline()

	wheel = [INPUT_NONE, INPUT_NONE, INPUT_NONE, INPUT_NONE, INPUT_NONE, INPUT_NONE, INPUT_NONE, INPUT_NONE]

	while(not spin_wheel(wheel)):
		if wheel[DIFFUSE] in allowed_permutations[DIFFUSE] \
				and wheel[SPECULAR] in allowed_permutations[SPECULAR] \
				and wheel[ROUGHNESS] in allowed_permutations[ROUGHNESS] \
				and wheel[METALLIC] in allowed_permutations[METALLIC] \
				and wheel[NORMAL] in allowed_permutations[NORMAL] \
				and wheel[EMISSION] in allowed_permutations[EMISSION] \
				and wheel[SUBSURFACE] in allowed_permutations[SUBSURFACE] \
				and wheel[HEIGHT] in allowed_permutations[HEIGHT]:
			(xml, fname) = mutate(wheel)

			file = open(fname + ".ankiprog", "w")
			file.write(xml)

if __name__ == "__main__":
	main()