#!/usr/bin/env python
"""Main script for importing submeshes from one model to another by throwing 
away the bones and weights of the original and copying them from nearby 
vertices in the new.

Recommended Procedure
---------------------

The basic procedures goes like this. Firstly, move, stretch and deform 
the new submeshes so they fit together with the parts which will remain 
of the original model, properly, and save them. Try to get them to fit 
closely over the original otherwise they might get the wrong bones, the only 
way the scripts will know that the index finger wants index finger bones 
rather than thumb bones is them being in the same place.

Where the imports will be connecting to the originals, you will probably 
need to get vertices of each into precisely the same location, although 
sometimes you can have them overlap and this looks perfectly fine. 
Using Blender's snap facility is probably easiest, then the process 
of running this script will automatically give the vertices matching 
bones and weights, but the script can do this too if they are just near 
each other.

Remove parts of the original model which are getting replaced, and save 
those submeshes too, but without overwriting the originals.

Work out which original submeshes you are importing into, this will normally 
be anything with overlapping material. Use the non-modifed versions.

Look at the bones in the vgmap files of the original submeshes you are 
importing into, and work out which bones seem appropriate. Remove any bones 
which the things you're importing shouldn't be using, and save them in a 
badbones file (which actually contains the good bones, sorry,) ending in .50 
The 'bones' here are the textual descriptions found between the quotes.

Now run this script on the meshes being imported, with the original submeshes 
you are importing into as recipients, and the badbones file.

If all goes well, this will give you outputs you can use merge_into.simple
to merge into submeshes of the original. Use the modified versions of the 
originals now. It may be that you are keeping nothing from an original submesh 
but the script wants to import things into it, particularly if you are 
replacing a lot. In this case, the output of this script can just be a 
straight replacement for it.

Now we need to sort out textures. The merge_into.simple script will create 
results which project the texture mapping of the original proportionally into 
the top left quarter of a notional new texture file, and the texture mapping 
of the new mesh into the top left corner, so to make the textures for the 
combined file, make dds files twice as big on each side, with the original 
textures in the top left, and the textures being added on the top right. 
(You can then scale it down if you're worried about file size.)

You can add new textures to the model here by editing mesh_metadata.json 
and the dds.json, but the most efficient thing to do is replace the textures 
for the original with the new thing, and then use move_tex on all other 
submeshes which use that texture to put their texture mapping into the 
top left quarter.

What this script specifically does
----------------------------------

Run this script with some meshes with bones and weights for the wrong model 
as the first arguments, followed by RECIPIENTS, and then some meshes with 
bones and weights for the correct model, and this script will give bones 
and weights to the first set of meshes based on what nearby things in the 
second set have, and produce new submeshes which can be merged in to the 
submeshes in the second set.

It's not always necessary, but it's generally a good idea to tell the script 
what bones it can use with one or more BADBONES file. These should follow a 
BADBONES argument at the end of the argument list, and should be files 
containing lists of bones by their textual name, with filenames ending in an 
integer less than 100. These can work in two ways - if the integer is 0, 
e.g. bb.0 then any vertex in the meshes we're getting the weights from will 
be totally ignored, this is useful, for instance, for preventing skirts which 
are being added from taking anything from arms or hands. Otherwise, the script 
will only adopt bones from this file, and will ignore any vertex in the meshes 
we're getting the weights from which is less than that percentage weighted 
with approved bones. Using a single file with percentage 50 is what I 
generally do.

If you haven't already done this, the script can also move vertices a little 
bit to snap to vertices in the copy-bones-from meshes. After the RECIPIENTS 
(and before any BADBONES) add some more submeshes (which can be the same as 
previous ones) after a SNAP_TO argument, and then even more submeshes after 
a NEAR argument, and a threshold, which is supplied as a numbwe which can be 
an argument at any point. The script will then move any vertices from the 
meshes we're importing to the nearest vertex it can find in a SNAP_TO mesh 
if it is closer than the threshold to a vertex in a NEAR mesh.

Generally, the SNAP_TO meshes will be the bone donor meshes. Depending on 
the situation, the NEAR meshes might be the importing meshes if every vertex 
should be snapped, or the donor meshes if you're confident that you've got 
every vertex which should be snapped closer than any non-equivalent vertex 
pairs, but most commonly, you can make new meshes from the import meshes by 
just deleting any vertices you don't want to snap or dragging them way away
from where everything is.
"""

import sys
import os.path

import vb
import snapmesh
import facemapper

def parse_args():
	modifying = []
	recipients = []
	snapping_to = []
	nearby = []
	badbones = []
	thresh = None
	l = modifying
	for arg in sys.argv[1:]:
		if thresh is None:
			try:
				thresh = float(arg)
				continue
			except ValueError:
				pass
		if l == modifying and arg == 'RECIPIENTS':
			l = recipients
			continue
		if l == recipients and arg == 'SNAP_TO':
			l = snapping_to
			continue
		if l == snapping_to and arg == 'NEAR':
			l = nearby
			continue
		if arg == 'BADBONES':
			badbones.append(list())
			l = badbones[-1]
			continue
		if badbones and not os.path.exists(vb.component_path(arg, 'vb')):
			with open(arg, 'r') as fi:
				l.append((int(arg.rsplit('.')[-1]), [x.strip() for x in fi]))
		else:
			l.append(vb.Mesh(arg))
	assert((thresh is None) == (not bool(snapping_to)))
	return modifying, recipients, snapping_to, nearby, thresh, badbones

class Restriction:
	def __init__(self, l, thresh):
		self.need_match = False
		self.meshes = []
		self.poison = []
		self.only = None
		self.threshholds = []
		self.thresh = thresh
		for thing in l:
			if type(thing) == tuple:
				if thing[0] == 0:
					self.poison.extend(thing[1])
				else:
					if self.only is None:
						self.only = set(thing[1])
					else:
						self.only.intersection_update(set(thing[1]))
					self.threshholds.append(thing)
			else:
				self.meshes.append(thing)
				self.need_match = True

	def apply(self, v):
		if not self.need_match:
			return True
		for m in self.meshes:
			if m.distance_to(v) < self.thresh:
				return True
		return False

	def allowed(self, v):
		totals = [0] * len(self.threshholds)
		assessing = v.bones_according_to_initial()
		for bone, weight in v.bones_according_to_initial():
			if bone in self.poison and weight >= self.thresh:
				return False
			for i, t in enumerate(self.threshholds):
				if bone not in t[1]:
					totals[i] += weight
		for i, t in enumerate(self.threshholds):
			# t[0] is a percentage
			if 100 * totals[i] > t[0]:
				return False
		return True

	def badbone(self, b):
		return b in self.poison or (self.only is not None and b not in self.only)

def bm(vertexmap, restrictions):
	if len(vertexmap) == 1:
		ret = vertexmap[0][0].bones_according_to_initial()
	else:
		ret = {}
		for v, weight in vertexmap:
			thesebones = v.bones_according_to_initial()
			for index, bw in thesebones:
				if index in ret:
					ret[index] += (weight * bw)
				else:
					ret[index] = weight * bw
		ret = list(ret.items())
	ret.sort(key=lambda x: -x[1])
	ret = [(x, y) for x, y in ret if True not in [r.badbone(x) for r in restrictions]]
	assert(ret)
	return ret

def avail_bones(base_avail, restrictions):
	ret = [x for x in base_avail if False not in [r.allowed(x) for r in restrictions]]
	if not ret:
		raise Exception('All vertices ruled out by badbones restrictions')
	return ret

def main():
	modifying, recipients, snapping_to, nearby, thresh, badbones = parse_args()
	fixsets = []
	realbones = []
	for i in range(len(modifying)):
		fixsets.append(list())
		realbones.append(list())
	bone_refs = vb.combined_bone_refs(recipients) # v, bonemap
	base_avail = [x[0] for x in bone_refs]
	if thresh is not None:
		assert(nearby or badbones)
		if nearby:
			origlengths = [x.lv for x in modifying]
			snapping_modified, modified = snapmesh.snap_meshes(modifying, nearby, thresh, snapping_to, sort_bones=False)
			for m, i, j, v in modified.keys():
				fixsets[i].append(j)
				# XXX just obsolete, since we have sort_bones=False?
				#if j >= len(realbones[i]):
				#	for _____ in range(len(realbones[i]) + 1 - j):
				#		realbones[i].append(None)
				#realbones[i][j] = [(m.bones[x], y) for x, y in v.bones()]
	if thresh is None:
		thresh = 0.000001
	restrictions = [Restriction(bb, thresh) for bb in badbones]
	for i, m in enumerate(modifying):
		for v in m.vertices:
			r = [x for x in restrictions if x.apply(v)]
			realbones[i].append(bm(vb.find_bone_donors(avail_bones(base_avail, r), v), r)) # real bone map
	mapper = facemapper.FaceMapper(modifying, realbones, fixsets, 
			[(i, list(x.bones.values())) for i, x in enumerate(recipients)])
	assignment = mapper.decide()
	bone_mappings = [m.bones for m in modifying]
	for i, r in enumerate(recipients):
		if i not in assignment.values():
			continue
		print ('Doing', i)
		new_mesh = facemapper.make_new_mesh(r, 
				[f for f in mapper.faces if assignment[f] == i], 
				mapper.vertices, bone_mappings)
		name = r.orig_base_path + '_incoming.vb'
		print ('Writing', name)
		new_mesh.write(name)
	for i, s in enumerate(snapping_to):
		if snapping_modified[i]:
			m = s.orig_base_path + '_automodded.vb'
			print (f'Modified {sys.argv[4+i]} too, writing to {m}')
			s.write(m)

if __name__ == '__main__':
	main()

