import torch, numpy as np, math
from torch.utils.data import Dataset
import h5py
from example.nnet_reader import HCAS_Model


class My_H5Dataset(torch.utils.data.Dataset):
    def __init__(self, file_path):
        super(My_H5Dataset, self).__init__()
        h5_file = h5py.File(file_path, 'r')
        self.X = h5_file['X']
        self.y = h5_file['y']
        self.means = h5_file['means']
        self.min_inputs = h5_file['min_inputs']
        self.max_inputs = h5_file['max_inputs']

    def __getitem__(self, index):
        return (torch.from_numpy(self.X[index, :]),
                torch.from_numpy(self.y[index, :]))

    def __len__(self):
        return self.X.shape[0]


if __name__ == '__main__':
    for k in range(5):
        for j in ['00', '05', 10, 15, 20, 30, 40, 60]:
            print('-'*50, k, j ,'-'*50)
            model = HCAS_Model('TrainedNetworks/HCAS_rect_v6_pra{}_tau{}_25HU_3000.nnet'.format(k, j))
            h5_file = My_H5Dataset('TrainingData/HCAS_rect_TrainingData_v6_pra{}_tau{}.h5'.format(k, j))
            loader = torch.utils.data.DataLoader(h5_file, batch_size=64)
            pred_list = []
            for X, y in loader:
                pred = torch.argmax(model(X.float()), dim=1).data.numpy()
                # pred = torch.argmax(y, dim=1).data.numpy()
                pred_list += list(pred)
            from collections import Counter
            print(Counter(pred_list))

            batch_size, epoch = 64, 10000
            pred_list = []
            for i in range(epoch):
                inputs = torch.randint(-30000, 30000, [batch_size, 2], dtype=torch.float)
                theta = torch.randn([batch_size, 1])*2*math.pi - math.pi
                inputs[0], inputs[1], theta = inputs[0], inputs[1], theta
                random_x = torch.cat([inputs, theta], dim=-1)
                pred = torch.argmax(model(random_x), dim=1).data.numpy()
                pred_list += list(pred)

            from collections import Counter
            print(Counter(pred_list))
# -------------------------------------------------- 0 00 --------------------------------------------------
# Counter({0: 22626, 4: 12403, 3: 12085, 1: 3393, 2: 3285})
# Counter({1: 309349, 0: 287455, 2: 22699, 3: 14766, 4: 5731})
# -------------------------------------------------- 0 05 --------------------------------------------------
# Counter({0: 22964, 4: 11765, 3: 11736, 1: 4380, 2: 2947})
# Counter({0: 240814, 3: 213992, 2: 167550, 1: 17644})
# -------------------------------------------------- 0 10 --------------------------------------------------
# Counter({0: 23456, 4: 8587, 3: 7967, 1: 7205, 2: 6577})
# Counter({4: 432704, 3: 207296})
# -------------------------------------------------- 0 15 --------------------------------------------------
# Counter({0: 23343, 1: 9130, 2: 8125, 4: 6776, 3: 6418})
# Counter({4: 342865, 0: 297135})
# -------------------------------------------------- 0 20 --------------------------------------------------
# Counter({0: 24863, 1: 10094, 2: 9009, 4: 5036, 3: 4790})
# Counter({0: 456475, 3: 153207, 4: 17965, 1: 12353})
# -------------------------------------------------- 0 30 --------------------------------------------------
# Counter({0: 23658, 1: 14392, 2: 10843, 4: 3017, 3: 1882})
# Counter({2: 521662, 3: 118338})
# -------------------------------------------------- 0 40 --------------------------------------------------
# Counter({0: 28038, 1: 12554, 2: 12473, 3: 700, 4: 27})
# Counter({3: 333239, 2: 121565, 0: 110735, 4: 57313, 1: 17148})
# -------------------------------------------------- 0 60 --------------------------------------------------
# Counter({0: 27127, 1: 13915, 2: 12750})
# Counter({1: 362550, 2: 170833, 0: 106617})
# -------------------------------------------------- 1 00 --------------------------------------------------
# Counter({0: 23554, 3: 12974, 1: 8853, 4: 8411})
# Counter({3: 414355, 1: 148941, 4: 76704})
# -------------------------------------------------- 1 05 --------------------------------------------------
# Counter({0: 24048, 3: 12646, 1: 9028, 4: 8070})
# Counter({4: 610042, 0: 29958})
# -------------------------------------------------- 1 10 --------------------------------------------------
# Counter({0: 24751, 1: 12552, 3: 10565, 4: 5924})
# Counter({1: 471531, 0: 82643, 3: 51075, 2: 34751})
# -------------------------------------------------- 1 15 --------------------------------------------------
# Counter({0: 23116, 1: 18513, 3: 8202, 4: 3961})
# Counter({4: 323593, 1: 316407})
# -------------------------------------------------- 1 20 --------------------------------------------------
# Counter({0: 24668, 1: 19611, 3: 6647, 4: 2866})
# Counter({4: 361549, 1: 117187, 2: 99844, 3: 42806, 0: 18614})
# -------------------------------------------------- 1 30 --------------------------------------------------
# Counter({1: 24601, 0: 23835, 3: 3900, 4: 1456})
# Counter({4: 578474, 3: 61526})
# -------------------------------------------------- 1 40 --------------------------------------------------
# Counter({1: 27441, 0: 23131, 3: 3220})
# Counter({1: 640000})
# -------------------------------------------------- 1 60 --------------------------------------------------
# Counter({1: 29770, 0: 22458, 3: 1564})
# Counter({4: 640000})
# -------------------------------------------------- 2 00 --------------------------------------------------
# Counter({0: 24506, 4: 13005, 3: 8368, 2: 7913})
# Counter({0: 597773, 4: 42227})
# -------------------------------------------------- 2 05 --------------------------------------------------
# Counter({0: 23800, 4: 12525, 2: 9928, 3: 7539})
# Counter({4: 214339, 1: 164153, 0: 160360, 3: 65757, 2: 35391})
# -------------------------------------------------- 2 10 --------------------------------------------------
# Counter({0: 23989, 2: 13454, 4: 10380, 3: 5969})
# Counter({1: 608842, 3: 31158})
# -------------------------------------------------- 2 15 --------------------------------------------------
# Counter({0: 24905, 2: 16030, 4: 8662, 3: 4195})
# Counter({3: 422780, 4: 217220})
# -------------------------------------------------- 2 20 --------------------------------------------------
# Counter({0: 23274, 2: 20393, 4: 6927, 3: 3198})
# Counter({0: 586837, 3: 32071, 1: 21092})
# -------------------------------------------------- 2 30 --------------------------------------------------
# Counter({2: 24790, 0: 23906, 4: 3767, 3: 1329})
# Counter({4: 636387, 0: 3613})
# -------------------------------------------------- 2 40 --------------------------------------------------
# Counter({2: 28834, 0: 21934, 4: 3024})
# Counter({3: 623263, 1: 16737})
# -------------------------------------------------- 2 60 --------------------------------------------------
# Counter({2: 28430, 0: 24442, 4: 920})
# Counter({1: 276046, 3: 228533, 2: 135421})
# -------------------------------------------------- 3 00 --------------------------------------------------
# Counter({0: 24717, 3: 14753, 4: 8416, 1: 5906})
# Counter({4: 426461, 0: 149344, 1: 64195})
# -------------------------------------------------- 3 05 --------------------------------------------------
# Counter({0: 24967, 3: 14311, 4: 8214, 1: 6300})
# Counter({1: 638493, 3: 1507})
# -------------------------------------------------- 3 10 --------------------------------------------------
# Counter({0: 25056, 3: 13025, 1: 9962, 4: 5749})
# Counter({1: 505722, 0: 82834, 3: 43879, 4: 7565})
# -------------------------------------------------- 3 15 --------------------------------------------------
# Counter({0: 25183, 1: 13908, 3: 10518, 4: 4183})
# Counter({3: 640000})
# -------------------------------------------------- 3 20 --------------------------------------------------
# Counter({0: 25279, 1: 17063, 3: 8405, 4: 3045})
# Counter({1: 297677, 0: 146895, 3: 114178, 2: 81250})
# -------------------------------------------------- 3 30 --------------------------------------------------
# Counter({0: 25513, 1: 19338, 3: 7157, 4: 1784})
# Counter({0: 609948, 4: 19468, 3: 10584})
# -------------------------------------------------- 3 40 --------------------------------------------------
# Counter({0: 24602, 1: 23809, 3: 5381})
# Counter({3: 446822, 1: 130078, 0: 60634, 4: 2466})
# -------------------------------------------------- 3 60 --------------------------------------------------
# Counter({1: 26461, 0: 23492, 3: 3839})
# Counter({2: 312639, 4: 190717, 3: 131328, 0: 5316})
# -------------------------------------------------- 4 00 --------------------------------------------------
# Counter({0: 24699, 4: 14460, 3: 8417, 2: 6216})
# Counter({3: 640000})
# -------------------------------------------------- 4 05 --------------------------------------------------
# Counter({0: 24832, 4: 14381, 3: 8051, 2: 6528})
# Counter({4: 483128, 3: 156872})
# -------------------------------------------------- 4 10 --------------------------------------------------
# Counter({0: 25059, 4: 13269, 2: 9638, 3: 5826})
# Counter({3: 592242, 0: 47758})
# -------------------------------------------------- 4 15 --------------------------------------------------
# Counter({0: 25324, 2: 13618, 4: 10545, 3: 4305})
# Counter({2: 507447, 0: 80926, 1: 46563, 4: 5064})
# -------------------------------------------------- 4 20 --------------------------------------------------
# Counter({0: 24705, 2: 17217, 4: 8708, 3: 3162})
# Counter({4: 631661, 2: 8339})
# -------------------------------------------------- 4 30 --------------------------------------------------
# Counter({0: 25501, 2: 20316, 4: 6476, 3: 1499})
# Counter({0: 473344, 3: 166656})
# -------------------------------------------------- 4 40 --------------------------------------------------
# Counter({2: 24326, 0: 24270, 4: 5196})
# Counter({2: 598730, 0: 32938, 4: 8332})
# -------------------------------------------------- 4 60 --------------------------------------------------
# Counter({2: 25726, 0: 24529, 4: 3537})
# Counter({3: 608345, 4: 31655})
#
# Process finished with exit code 0

