from torch.utils.data import Dataset
from libAIF import aif
import os
import json
import numpy as np
import cv2
from label import RPPGLabel
import random
import torch
from torchvision import transforms
from feature import feat_normalize

# captured compact AIF dataset
class iframe_feeder(Dataset):
    def __init__(self, featdfile, featifile, labfile, normdfile, normifile, test_ratio=0.1):
        self.test_ratio = test_ratio
        self.pic_size = 49

        #feature access
        self.feature_aif = aif(featdfile, featifile)
        self.ftdata_info = self.feature_aif.rbegin()
        self.norm_aif = aif(normdfile, normifile)
        self.norm_info = self.norm_aif.rbegin()

        #label access
        labfh = open(labfile)
        self.all_labels = json.load(labfh)
        labfh.close()

        #other
        self.datalen = 0
        self.datalen_train = 0
        self.datalen_test = 0
        self.valid_indeces = []
        self.valid_train_indeces = []
        self.valid_test_indeces = []
        self.mode = None
        #screen invalid data, randomization
        self._prepare()

    def finish(self):
        self.feature_aif.rfinish()
        self.norm_aif.rfinish()
        
    def set_mode(self, mode):
        self.mode = mode

    def _prepare(self):
        raw_len = len(self.ftdata_info)
        for i in range(0, raw_len):
            if self.all_labels[i] != None:# and len(self.ftps[i]) > 0:
                self.valid_indeces.append(i)
                self.datalen += 1
        # randomization
        test_size = int(self.datalen * self.test_ratio)
        test_beg = random.randint(0, self.datalen - test_size)
        test_end = test_beg + test_size
        #feats, labs
        for x in range(0, len(self.valid_indeces)):
            if test_beg < x < test_end:
                self.valid_test_indeces.append(self.valid_indeces[x])
            else:
                self.valid_train_indeces.append(self.valid_indeces[x])

    
    def __getitem__(self, index):
        # feat_buf = []
        f_tensor = None
        l_tensor = None
        actual_idx = -1
        if self.mode == "train":
            actual_idx = self.valid_train_indeces[index]
        elif self.mode == "test":
            actual_idx = self.valid_test_indeces[index]
        else:
            return None

        xfeat = self.feature_aif.extract(actual_idx)
        xfeat = xfeat.reshape(self.pic_size,self.pic_size,3)
        xnorm = self.norm_aif.extract_raw(actual_idx, 'float32')
        norm_feat = feat_normalize(xfeat, xnorm, m_offset=0)

        pre_proc = transforms.Compose([
            transforms.ToTensor(),
            #transforms.Normalize(mean=[0.5,0.5,0.5], std=[0.05,0.05,0.05]),
        ])

        f_tensor = pre_proc(norm_feat).to(torch.float32)
        xlab = np.array(self.all_labels[actual_idx])
        l_tensor = torch.from_numpy(xlab).to(torch.int64)
        
        return f_tensor, l_tensor

    def __len__(self):
        # return self.datalen
        if self.mode == "train":
            return len(self.valid_train_indeces)
        elif self.mode == "test":
            return len(self.valid_test_indeces)
        else:
            return None

if __name__ == "__main__":
    x = iframe_feeder("F:\\ecg")
    x._read_featlab("F:\\ecg")