import numpy as np
import trimesh

thumb = [[1267,1268,1270,1271,1272,1273,1274,1275,1276,1277,1278,1279,1280,1281,1289,1290,1291,1292,1295,1296,1298,1299,1300,1304,1305,1306,1311,1312,1313,1314,1315,1316,1317,1318,1319,1320,1323,1327,1328,1339,1340,1350,1357,1368,1369,1370,1371,1373,1374,1375,1376,1379,1380,1381,1382,1269,1282,1287,1288,1297,1301,1302,1303,1307,1308,1324,1325,1326,1329,1330,1377,1321,1322,1363,1367,1372,1354,1358,1349,1351,1364,1359,1361,1378,1355,1356,1360,1353,1365,1362,1352,1366,1265,1293,1309,1337,1344,1283,1285,1348,1284,1294,1266,1347,1345,1346,1342,1341,1338,1343,1310,1286
],[1234,1235,1236,1237,1238,1247,1248,1249,1245,1246,1250,92,233,331,1232,1241,1242,1244,1251,1253,1255,1257,1259,1261,1263,1331,1333,1336,1233,1240,235,361,234,330,335,360,1231,1243,1252,1254,1256,1258,1260,1262,1264,1332,1334,1335,1239,236,93,1479],
[4,5,223,224,1480,1481,1535,1536,1537,327,25,54,55,56,2,3,22,83,23,24,1533,1534,7,1,16,57,63,72,81,82,1462,1465,1466,1468,1475,1477,1482,1500,1531,326,1470,1471,1473,1532,259,1478,6,62,0,17,36,37,73,80,214,218,334,1464,1467,1469,1472,1474,1476,1483,1499,1501,1503,1530,1463]]

index = [[398,411,412,413,433,434,437,438,445,447,448,451,452,453,454,477,483,484,485,486,487,488,489,490,491,495,496,499,500,501,502,503,504,509,510,401,402,403,405,406,410,507,508,511,512,478,482,404,409,432,444,492,493,494,505,399,400,414,423,429,430,497,498,481,427,439,443,506,424,431,397,421,408,446,415,435,450,473,428,475,464,422,480,440,416,407,449,436,463,474,476,479],
[458,392,393,394,385,386,389,390,391,457,459,388,381,382,383,384,460,420,387,456,277,315,350,461,465,470,472,305,301,396,195,299,455,467,321,442,419,418,426,220,441,417,196,230,276,300,304,314,320,351,462,466,469,471,395,425,298,468,229,219],
[134,157,185,254,346,516,197,198,352,531,155,100,101,146,289,353,184,203,204,253,529,530,156,518,515,524,347,522,513,514,525,526,517,135,264,265,521,527,128,519,520,306,147,288,307,286,287,523,154,129,53,200,275,528,365,132,532,29,35,117,103,256,131,303,149,375,369,255,308,364,52,116,130,133,199,274,302,28,34,102,142,148,374,368]]

middle = [[629,630,632,633,634,635,636,637,638,639,640,641,642,651,652,657,658,659,660,661,662,663,664,673,674,675,676,677,679,680,683,684,685,686,687,688,711,712,716,717,718,719,720,721,722,723,724,725,726,727,728,729,730,735,736,737,739,740,741,742,743,744,667,668,733,734,738,715,731,732,631,643,649,654,665,669,678,682,708,709,698,714,644,650,653,666,670,681,697,707,710,713],
[571,572,573,574,587,589,590,599,600,618,619,620,621,622,689,690,693,694,597,598,604,563,614,588,691,581,582,617,626,692,613,623,624,603,616,564,593,594,625,615,533,535,537,543,567,579,607,611,628,648,672,695,702,706,591,609,699,704,646,656,534,536,538,544,568,580,608,610,612,627,645,647,655,671,696,701,703,705,592,700],
[540,551,552,577,578,605,606,747,748,752,755,756,762,761,539,575,576,749,750,751,542,545,553,556,558,559,561,565,583,595,745,753,759,763,602,757,569,586,548,549,764,541,546,547,554,555,557,560,562,566,584,585,596,746,754,758,760,601,570,550]]

ring = [[865,866,867,868,869,870,871,872,875,876,877,878,879,880,885,886,887,888,893,894,896,897,898,903,909,910,911,912,917,918,920,922,923,924,934,947,948,949,950,951,952,953,954,955,956,957,958,959,960,961,962,963,964,969,971,972,973,974,975,976,977,978,979,980,933,943,944,945,946,873,874,889,890,904,913,914,967,968,970,895,905,915,916,919,921,965,966,902,899,900,901,906,863,881,883,891,936,939,941,932,907,938,940,942,864,882,884,892,931,935,937,908],
[800,808,817,818,823,824,825,829,830,833,834,835,836,839,840,807,809,810,799,826,775,777,781,803,827,843,845,847,849,851,853,855,862,925,927,930,773,815,857,859,816,774,776,778,782,804,828,844,846,848,850,852,854,856,926,928,929,860,861,858],
[779,780,790,811,812,813,814,837,838,981,982,983,984,985,986,987,989,990,991,992,993,994,995,996,999,1388,789,997,998,841,842,988,127,343,370,377,379,783,785,787,793,805,819,822,831,1000,1387,1395,797,801,791,795,798,126,159,342,371,376,380,784,786,788,792,794,802,806,820,821,832,1396,796]]

little = [[1101,1102,1103,1104,1105,1106,1107,1108,1109,1110,1111,1112,1121,1122,1125,1126,1129,1130,1131,1132,1133,1134,1135,1136,1138,1141,1142,1143,1144,1145,1146,1147,1149,1150,1152,1155,1156,1165,1166,1175,1176,1179,1181,1185,1186,1187,1188,1189,1190,1191,1192,1193,1194,1195,1196,1199,1200,1201,1202,1203,1204,1205,1206,1207,1208,1209,1210,1211,1212,1148,1151,1153,1154,1197,1198,1097,1098,1099,1100,1118,1119,1127,1128,1137,1177,1178,1180,1182,1183,1184,1117,1120,1113,1115,1139,1095,1123,1163,1168,1170,1171,1174,1116,1114,1140,1096,1124,1164,1167,1169,1172,1173],
[1031,1032,1039,1040,1041,1042,1049,1055,1056,1057,1058,1061,1062,1065,1066,1071,1072,1050,1067,1068,1087,1001,1003,1005,1011,1035,1047,1059,1075,1077,1079,1089,1091,1094,1157,1159,1162,1083,1081,1085,1088,1002,1004,1006,1012,1036,1048,1060,1076,1078,1080,1082,1090,1092,1093,1158,1160,1161,1086,1084],
[1007,1008,1043,1044,1045,1046,1073,1074,1215,1216,1223,1224,1227,1019,1020,1228,1217,1218,1219,1220,1010,1013,1016,1017,1022,1024,1025,1037,1054,1063,1070,1213,1222,1225,1229,1433,1029,1033,1027,1051,1009,1014,1015,1018,1021,1023,1026,1028,1030,1038,1053,1064,1069,1214,1221,1226,1230,1434,1034,1052]]

hand = [thumb, index, middle, ring, little]

contact_thumb = [1353,1354,1357,1358,1366,1375,1376,1377,1378,1368,1363,1364,1365,1367,1379,1380,1361,1362,1373,1374,1369,1359,1372,1306,1360,1305,1371,1351,1356,1382,1381,1350,1301,1339,1370,1321,1349,1302,1322,1300,1323,1340,1352,1355,1311,1312,1299]
contact_index = [437,497,498,499,500,503,504,434,438,488,501,505,433,481,432,464,502,506,435,487,482,430,492,480,431,491,463,479,490,443,429,436,439,440,489,478,427,428,477,444,493,494,475,476,485,486]
contact_middle = [738,731,733,734,737,732,698,664,667,668,665,739,663,697,662,735,722,666,715,713,740,660,669,696,661,714,701,659,736,721,716,725,726,699,700,657,673,695,670,671,672,702]
contact_ring = [967,968,969,970,971,958,973,951,972,974,904,952,957,903,899,934,960,900,975,959,950,961,948,898,901,962,933,976,902,947,897,949,945,955,956,895,896,946,906,905,953]
contact_little = [1183,1184,1204,1199,1192,1180,1200,1203,1202,1179,1201,1191,1182,1136,1131,1132,1190,1188,1165,1189,1177,1128,1185,1206,1127,1135,1186,1130,1205,1187,1166,1178,1181,1133,1119,1193,1129,1137,1207,1176,1125,1134,1194,1208,1175,1198,1152,1117,1120,1126,1138,1151,1195]
contact_palm = [1416,68,1417,1413,1497,69,1415,227,1409,1496,1425,228,1418,1410,1414,1441,1444,1445,1495,1407,1408,1419,1439,1440,1442,1443,1446,1494]

contact = [contact_thumb, contact_index, contact_middle, contact_ring, contact_little, contact_palm]

seal = [1538,1539,1540,1541,1542,1543,1544,1545,1546,1547,1548,1549,1550,1551,1552,1553]

N_faces_sealed = 1554
N_faces = 1538

def creat_face2label_array(output_dir = './',sealed = True):
	face2label_array = []
	if sealed:
		face2label_array = np.zeros((N_faces_sealed, 4), dtype = np.int8)
	else:
		face2label_array = np.zeros((N_faces, 3), dtype = np.int8)

	for i in range(len(hand)):
		for j in range(len(hand[i])):
			for f in hand[i][j]:
				face2label_array[f,0] = i+1
				face2label_array[f,1] = i*3 + j + 1

	for i in range(len(contact)):
		for f in contact[i]:
			face2label_array[f, 2] = i+1

	if sealed:
		for f in seal:
			face2label_array[f, 3] = 1

	if sealed:
		np.save(output_dir+'/face2label_sealed.npy', face2label_array) 
	else:
		np.save(output_dir+'/face2label.npy', face2label_array)

def writeFaceColorOff(output, vertex, face, color):
    with open(output, 'w') as f:
        f.write("OFF\n")
        f.write("%d %d 0\n" %(vertex.shape[0], face.shape[0]))
        for row in range(0, vertex.shape[0]):
            f.write("%f %f %f\n" %(vertex[row, 0], vertex[row, 1], vertex[row, 2]))
        for row in range(0, face.shape[0]):
        	f.write("3 %d %d %d %d %d %d\n" %(face[row, 0], face[row, 1], face[row, 2], color[row, 0], color[row, 1], color[row, 2]))


def draw_model_seg(output_dir = './', hand_mesh = './model_normalized_sealed.obj'):
	mesh = trimesh.load(hand_mesh, process=False)
	color_h = np.zeros((N_faces_sealed, 3), dtype = np.int16)
	color_l = np.zeros((N_faces_sealed, 3), dtype = np.int16)
	color_c = np.zeros((N_faces_sealed, 3), dtype = np.int16)
	color_s = np.zeros((N_faces_sealed, 3), dtype = np.int16)

	color_dict = np.random.randint(256, size=(16, 3))

	face2label_array = np.load('./face2label_sealed.npy')
	print(face2label_array)
	for f in range(0, N_faces_sealed):
		color_l[f] = color_dict[face2label_array[f,0],:]
		color_h[f] = color_dict[face2label_array[f,1],:]
		color_c[f] = color_dict[face2label_array[f,2],:]
		if face2label_array.shape[1] == 4:
			color_s[f] = color_dict[face2label_array[f,3],:]

	writeFaceColorOff(output_dir+'/seg_coarse.off', mesh.vertices, mesh.faces, color_l)
	writeFaceColorOff(output_dir+'/seg_fine.off', mesh.vertices, mesh.faces, color_h)
	writeFaceColorOff(output_dir+'/contact.off', mesh.vertices, mesh.faces, color_c)
	if face2label_array.shape[1] == 4:
		if not mesh.faces.shape[0] == color_s.shape[0]:
			print("Face labels are for sealed hand, but the model is not sealed!")
			print("Current model is %s" % hand_mesh)
			print("Please change model to a sealed model!")
			exit()
		writeFaceColorOff(output_dir+'/seal.off', mesh.vertices, mesh.faces, color_s)

creat_face2label_array()
draw_model_seg()
print("Successfully Done!")