import numpy as np
import torch
from torch.autograd import Variable


trees = [
[1, 7],
[2, 6],
[3, 6],
[4, 6],
[5, 6],
[6, 3],
[7, 4],
[8, 2],
[9, 3],
[10, 2],
[11, 3],
[12, 2],
[13, 3],
[14, 3],
[15, 3],
[16, 6],
[17, 6],
[18, 9],
[19, 4],
[20, 6],
[21, 2],
[22, 3],
[23, 6],
[24, 6],
[25, 3],
[26, 6],
[27, 2],
[28, 3],
[29, 6],
[30, 9],
[31, 2],
[32, 7],
[33, 7],
[34, 3],
[35, 6],
[36, 2],
[37, 7],
[38, 2],
[39, 2],
[40, 6],
[41, 6],
[42, 3],
[43, 3],
[44, 6],
[45, 2],
[46, 3],
[47, 6],
[48, 7],
[49, 6],
[50, 7],
[51, 6],
[52, 7],
[53, 1],
[54, 1],
[55, 2],
[56, 3],
[57, 3],
[58, 7],
[59, 2],
[60, 5],
[61, 6],
[62, 7],
[63, 6],
[64, 8],
[65, 1],
[66, 3],
[67, 6],
[68, 7],
[69, 1],
[70, 1],
[71, 8],
[72, 3],
[73, 6],
[74, 1],
[75, 1],
[76, 7],
[77, 2],
[78, 5],
[79, 6],
[80, 2],
[81, 2],
[82, 9],
[83, 9],
[84, 9],
[85, 5],
[86, 1],
[87, 1],
[88, 8],
[89, 7],
[90, 1],
[91, 1],
[92, 9],
[93, 3],
[94, 7],
[95, 7],
[96, 6],
[97, 6],
[98, 4],
[99, 3],
[100, 2],
[101, 3],
[102, 2],
[103, 2],
[104, 3],
[105, 6],
[106, 1],
[107, 2],
[108, 5],
[109, 7],
[110, 7],
[111, 1],
[112, 3],
[113, 1],
[114, 1],
[115, 6],
[116, 9],
[117, 6],
[118, 7],
[119, 8],
[120, 7],
[121, 7],
[122, 1],
[123, 2],
[124, 1],
[125, 1],
[126, 5],
[127, 5],
[128, 3],
[129, 6],
[130, 4],
[131, 7],
[132, 7],
[133, 7],
[134, 6],
[135, 6],
[136, 6],
[137, 6],
[138, 6],
[139, 4],
[140, 6],
[141, 3],
[142, 7],
[143, 7],
[144, 3],
[145, 7],
[146, 7],
[147, 7],
[148, 7],
[149, 7],
[150, 3],
[151, 3],
[152, 3],
[153, 3],
[154, 7],
[155, 7],
[156, 6],
[157, 2],
[158, 2],
[159, 7],
[160, 3],
[161, 2],
[162, 6],
[163, 3],
[164, 6],
[165, 6],
[166, 8],
[167, 6],
[168, 4],
[169, 8],
[170, 4],
[171, 3],
[172, 3],
[173, 6],
[174, 5],
[175, 2],
[176, 6],
[177, 6],
[178, 4],
[179, 2],
[180, 3],
[181, 6],
[182, 6],
[183, 4],
[184, 6],
[185, 6],
[186, 7],
[187, 6],
[188, 6],
[189, 7],
[190, 4],
[191, 4],
[192, 4],
[193, 4],
[194, 6],
[195, 7],
[196, 2]
]


trees_order_to_family = [
[1, 53, 54, 65, 69, 70, 74, 75, 86, 87, 90, 91, 106, 111, 113, 114, 122, 124, 125], 
[2, 8, 10, 12, 21, 27, 31, 36, 38, 39, 45, 55, 59, 77, 80, 81, 100, 102, 103, 107, 123, 157, 158, 161, 175, 179, 196], 
[3, 6, 9, 11, 13, 14, 15, 22, 25, 28, 34, 42, 43, 46, 56, 57, 66, 72, 93, 99, 101, 104, 112, 128, 141, 144, 150, 151, 152, 153, 160, 163, 171, 172, 180], 
[4, 7, 19, 98, 130, 139, 168, 170, 178, 183, 190, 191, 192, 193], 
[5, 60, 78, 85, 108, 126, 127, 174], 
[6, 2, 3, 4, 5, 16, 17, 20, 23, 24, 26, 29, 35, 40, 41, 44, 47, 49, 51, 61, 63, 67, 73, 79, 96, 97, 105, 115, 117, 129, 134, 135, 136, 137, 138, 140, 156, 162, 164, 165, 167, 173, 176, 177, 181, 182, 184, 185, 187, 188, 194], 
[7, 1, 32, 33, 37, 48, 50, 52, 58, 62, 68, 76, 89, 94, 95, 109, 110, 118, 120, 121, 131, 132, 133, 142, 143, 145, 146, 147, 148, 149, 154, 155, 159, 186, 189, 195], 
[8, 64, 71, 88, 119, 166, 169], 
[9, 18, 30, 82, 83, 84, 92, 116]
]




def get_order_family_target(targets):


    order_target_list = []
    # family_target_list = []


    for i in range(targets.size(0)):

        order_target_list.append(trees[targets[i]][1]-1)
        # family_target_list.append(trees[targets[i]][2]-1)



    order_target_list = Variable(torch.from_numpy(np.array(order_target_list)).cuda())   
    # family_target_list = Variable(torch.from_numpy(np.array(family_target_list)).cuda())

    # return order_target_list, family_target_list
    return order_target_list


import numpy as np
import torch
from torch.autograd import Variable


trees = [
[1, 7],
[2, 6],
[3, 6],
[4, 6],
[5, 6],
[6, 3],
[7, 4],
[8, 2],
[9, 3],
[10, 2],
[11, 3],
[12, 2],
[13, 3],
[14, 3],
[15, 3],
[16, 6],
[17, 6],
[18, 9],
[19, 4],
[20, 6],
[21, 2],
[22, 3],
[23, 6],
[24, 6],
[25, 3],
[26, 6],
[27, 2],
[28, 3],
[29, 6],
[30, 9],
[31, 2],
[32, 7],
[33, 7],
[34, 3],
[35, 6],
[36, 2],
[37, 7],
[38, 2],
[39, 2],
[40, 6],
[41, 6],
[42, 3],
[43, 3],
[44, 6],
[45, 2],
[46, 3],
[47, 6],
[48, 7],
[49, 6],
[50, 7],
[51, 6],
[52, 7],
[53, 1],
[54, 1],
[55, 2],
[56, 3],
[57, 3],
[58, 7],
[59, 2],
[60, 5],
[61, 6],
[62, 7],
[63, 6],
[64, 8],
[65, 1],
[66, 3],
[67, 6],
[68, 7],
[69, 1],
[70, 1],
[71, 8],
[72, 3],
[73, 6],
[74, 1],
[75, 1],
[76, 7],
[77, 2],
[78, 5],
[79, 6],
[80, 2],
[81, 2],
[82, 9],
[83, 9],
[84, 9],
[85, 5],
[86, 1],
[87, 1],
[88, 8],
[89, 7],
[90, 1],
[91, 1],
[92, 9],
[93, 3],
[94, 7],
[95, 7],
[96, 6],
[97, 6],
[98, 4],
[99, 3],
[100, 2],
[101, 3],
[102, 2],
[103, 2],
[104, 3],
[105, 6],
[106, 1],
[107, 2],
[108, 5],
[109, 7],
[110, 7],
[111, 1],
[112, 3],
[113, 1],
[114, 1],
[115, 6],
[116, 9],
[117, 6],
[118, 7],
[119, 8],
[120, 7],
[121, 7],
[122, 1],
[123, 2],
[124, 1],
[125, 1],
[126, 5],
[127, 5],
[128, 3],
[129, 6],
[130, 4],
[131, 7],
[132, 7],
[133, 7],
[134, 6],
[135, 6],
[136, 6],
[137, 6],
[138, 6],
[139, 4],
[140, 6],
[141, 3],
[142, 7],
[143, 7],
[144, 3],
[145, 7],
[146, 7],
[147, 7],
[148, 7],
[149, 7],
[150, 3],
[151, 3],
[152, 3],
[153, 3],
[154, 7],
[155, 7],
[156, 6],
[157, 2],
[158, 2],
[159, 7],
[160, 3],
[161, 2],
[162, 6],
[163, 3],
[164, 6],
[165, 6],
[166, 8],
[167, 6],
[168, 4],
[169, 8],
[170, 4],
[171, 3],
[172, 3],
[173, 6],
[174, 5],
[175, 2],
[176, 6],
[177, 6],
[178, 4],
[179, 2],
[180, 3],
[181, 6],
[182, 6],
[183, 4],
[184, 6],
[185, 6],
[186, 7],
[187, 6],
[188, 6],
[189, 7],
[190, 4],
[191, 4],
[192, 4],
[193, 4],
[194, 6],
[195, 7],
[196, 2]
]


trees_order_to_family = [
[1, 53, 54, 65, 69, 70, 74, 75, 86, 87, 90, 91, 106, 111, 113, 114, 122, 124, 125], 
[2, 8, 10, 12, 21, 27, 31, 36, 38, 39, 45, 55, 59, 77, 80, 81, 100, 102, 103, 107, 123, 157, 158, 161, 175, 179, 196], 
[3, 6, 9, 11, 13, 14, 15, 22, 25, 28, 34, 42, 43, 46, 56, 57, 66, 72, 93, 99, 101, 104, 112, 128, 141, 144, 150, 151, 152, 153, 160, 163, 171, 172, 180], 
[4, 7, 19, 98, 130, 139, 168, 170, 178, 183, 190, 191, 192, 193], 
[5, 60, 78, 85, 108, 126, 127, 174], 
[6, 2, 3, 4, 5, 16, 17, 20, 23, 24, 26, 29, 35, 40, 41, 44, 47, 49, 51, 61, 63, 67, 73, 79, 96, 97, 105, 115, 117, 129, 134, 135, 136, 137, 138, 140, 156, 162, 164, 165, 167, 173, 176, 177, 181, 182, 184, 185, 187, 188, 194], 
[7, 1, 32, 33, 37, 48, 50, 52, 58, 62, 68, 76, 89, 94, 95, 109, 110, 118, 120, 121, 131, 132, 133, 142, 143, 145, 146, 147, 148, 149, 154, 155, 159, 186, 189, 195], 
[8, 64, 71, 88, 119, 166, 169], 
[9, 18, 30, 82, 83, 84, 92, 116]
]




def get_order_family_target(targets):


    order_target_list = []
    # family_target_list = []


    for i in range(targets.size(0)):

        order_target_list.append(trees[targets[i]][1]-1)
        # family_target_list.append(trees[targets[i]][2]-1)



    order_target_list = Variable(torch.from_numpy(np.array(order_target_list)).cuda())   
    # family_target_list = Variable(torch.from_numpy(np.array(family_target_list)).cuda())

    # return order_target_list, family_target_list
    return order_target_list


# family_target_list = [] # family_target_list.append(trees[targets[i]][2]-1) # family_target_list = Variable(torch.from_numpy(np.array(family_target_list)).cuda()) # return order_target_list, family_target_list