import sys

import numpy as np
import pickle as pkl
from scipy import sparse

from time import time

from data.supersmpl_io import writeOBJ, readOBJ


def makeGraph(F):
	G = {}
	for f in F:
		N = len(f)
		for i,v in enumerate(f):
			if not v in G: G[v] = {}
			G[v][f[(i+1) % N]] = 1
	return G

def extrapolateMissings(M, G):
	for i in range(6890, M.shape[0]):
		if np.isnan(M[i]).all():
			M[i] = neighAvg(M, G, i)
		elif np.isnan(M[i]).any():
			print("HOW?")
			print(M[i])
	if np.isnan(M).any(): extrapolateMissings(M, G)

def neighAvg(M, G, i):
	neigh = set(G[i].keys())
	avg = 0
	count = 0
	for n in neigh:
		if not np.isnan(M[n]).any():
			avg += M[n]
			count += 1
	if count: return avg / count
	return np.nan


if __name__ == '__main__':

	g = 'f'
	# Read Super topology
	V, F, Vt, Ft = readOBJ('super_f_avg.obj')
	N = V.shape[0]
	G = makeGraph(F)

	# Read PKL
	print("Loading .pkl file")
	with open('custom_data/model_' + g + '.pkl', 'rb') as f:
		M = pkl.load(f, encoding='latin1')
		
	# Fields
	# print("Checking fields and shapes")
	# for k in M:
	# 	print(k, M[k].shape)
	# Faces
	print("New faces...")
	M['f'] = F
	# J_regressor (24, 6890)
	# Sparse matrix, just change shape to add zeros
	print("New regressor")
	A = np.zeros((24, N), np.float32)
	A[:,:6890] = M['J_regressor'].todense()
	M['J_regressor'] = sparse.csr_matrix(A)

	# Weights (6890, 24)
	print("New weights")
	start = time()
	A = np.empty((N, 24), np.float32)
	A[:] = np.nan
	A[:6890] = M['weights'][:6890]
	extrapolateMissings(A, G)
	M['weights'] = A
	print(time() - start)

	# Posedirs (6890, 3, 207)
	print("New posedirs")
	start = time()
	A = np.empty((N, 3, 207), np.float32)
	A[:] = np.nan
	A[:6890] = M['posedirs'][:6890]
	extrapolateMissings(A, G)
	M['posedirs'] = A
	print(time() - start)

	# v_tempalte (6890, 3)
	print("New template")
	start = time()
	A = np.empty((N, 3), np.float32)
	A[:] = np.nan
	A[:6890] = M['v_template'][:6890]
	extrapolateMissings(A, G)
	M['v_template'] = A
	print(time() - start)

	# Shapedirs (6890, 3, 10)
	print("New shapedirs")
	start = time()
	A = np.empty((N, 3, 10), np.float32)
	A[:] = np.nan
	A[:6890] = M['shapedirs'][:6890]
	extrapolateMissings(A, G)
	M['shapedirs'] = A
	print(time() - start)

	# Dump
	print("Saving new model")
	with open('super_model_' + g + '.pkl', 'wb') as f:
		pkl.dump(M, f)