import cv2
import os
import json
from libAIF import aif
import numpy as np
from label import RPPGLabel
from collections import deque
from multiprocessing import Process, Pool
import math

def make_feature_aif_databse(outputfile, basefolder, if_src_raw=False, expected_shape=(49,49,3)):
    aif_allftimg = aif(outputfile+"_ftimgd", outputfile+"_ftimgi")
    aif_allftimg.wbegin()

    all_labels = []
    fh_label = open(outputfile + "_label.json", 'w')
    for root, dirs, files in os.walk(basefolder):
        for dname in dirs:
            prev_path = os.path.join(root, dname) + "\\"
            # xsplit = prev_path.split("\\")
            dfilename = prev_path + "datafile"
            if os.path.exists(dfilename) and os.path.dirname(dfilename)[-1] == '1':# get only the first video
                print("on going: " + dfilename)

                #label src file
                src_label_fname = root + "\\" + "bvp_" + root.split("\\")[-1] + "_T" + os.path.dirname(dfilename)[-1] + ".csv"
                print("corresponding label: " + src_label_fname)
                rppg_lab = RPPGLabel(src_label_fname)

                #features & labels
                aif_inst = aif(dfilename, prev_path + "idxfile")
                dat_info = aif_inst.rbegin()
                for i in range(0, len(dat_info), 1):
                    if if_src_raw:
                        img0 = (aif_inst.extract_raw(i)).copy()
                    else:
                        img0 = aif_inst.extract(i)
                    # img0 = cv2.resize(img0, (144, 144))

                    if (i + 1) < len(dat_info):
                        if if_src_raw:
                            img1 = (aif_inst.extract_raw(i + 1)).copy()
                        else:
                            img1 = aif_inst.extract(i + 1)
                        if img1.shape == expected_shape and img0.shape == expected_shape:
                            # adjust to the median position
                            ftimg = img1 - img0 + 128
                            aif_allftimg.append(ftimg)# features
                            all_labels.append(rppg_lab.get_label_by_frame(i)) #labels
                        else:
                            aif_allftimg.append_any(np.array([]).tobytes())
                            all_labels.append(None)                            

                    else:# end of a file, append an empty node
                        aif_allftimg.append_any(np.array([]).tobytes())
                        all_labels.append(None)

                aif_inst.rfinish()
    json.dump(all_labels, fh_label)
    fh_label.close()
    aif_allftimg.wfinish()

#300 frames * 478 landmarks * 5 cross pixels * 3 channels
#every 5 pixel makes a cross
#dat shape=(300,49,49,3)
#return 478*3 (mean), 478*3 (std)
def get_norm(dat):
    ret_mean = []
    ret_std = []
    for i in range(0, 478): # cross
        c1=[]
        c2=[]
        c3=[]
        for f in range(0, len(dat)): # frames
            for j in range(0, 5): # pixels
                px_idx = (i * 5) + j
                px_ridx = px_idx // 49
                px_cidx = px_idx % 49
                c1.append(dat[f][px_ridx][px_cidx][0])
                c2.append(dat[f][px_ridx][px_cidx][1])
                c3.append(dat[f][px_ridx][px_cidx][2])
        ret_mean.append([np.mean(c1), np.mean(c2), np.mean(c3)])
        ret_std.append([np.std(c1), np.std(c2), np.std(c3)])
    return [ret_mean, ret_std]

#feat: 49*49*3, every 5 as a cross with 1 norm
#norm: 478*3 (mean), 478*3 (std)
def feat_normalize(feat, norm, m_offset=0):
    ret = []
    npixel = feat.shape[0] * feat.shape[1]
    norm = norm.reshape(-1, 478, 3)
    for i in range(0, npixel):
        irow = i // 49
        icol = i % 49
        if i < 478*5:
            tar_vec = [0,0,0]
            for z in range(0, 3):
                if norm[1][i//5][z] == 0:
                    if feat[irow][icol][z] >= 0:
                        tar_vec[z] = 4
                    else:
                        tar_vec[z] = -4
                else:
                    tar_vec[z] = (feat[irow][icol][z] - m_offset - norm[0][i//5][z]) / norm[1][i//5][z]
            ret.append(np.array(tar_vec))
        else:
            ret.append(np.array([0,0,0]))
    ret = np.array(ret, dtype='float32')
    return ret.reshape(49,49,3)

#prefix + digi + appendix
def merge_aifs_digiapdx(aiffolder, prefix,\
        appendix_d, appendix_i, begnum, endnum, destpath,\
        extract="image", dtype="float32"):
    merged_aif = aif(destpath + "_dfile", destpath + "_ifile")
    merged_aif.wbegin()
    for i in range(begnum, endnum + 1):
        src_d_path = aiffolder + prefix + str(i) + appendix_d
        src_i_path = aiffolder + prefix + str(i) + appendix_i
        if os.path.exists(src_d_path) and os.path.exists(src_i_path):
            print(src_i_path)
            src_aif = aif(src_d_path, src_i_path)
            info = src_aif.rbegin()
            buf = None
            for j in range(0, len(info)):
                if extract == "image":
                    buf = src_aif.extract(j)
                    merged_aif.append(buf)
                elif extract == "bytes":
                    buf = src_aif.extract_raw(j, dtype=dtype)
                    merged_aif.append_any(np.array(buf, dtype=dtype).tobytes())
            src_aif.rfinish()
    merged_aif.wfinish()

def merge_jsons(jsonfolder, prefix, appendix, begnum, endum, destpath):
    dat = []
    for i in range(begnum, endum + 1):
        filepath = jsonfolder + prefix + str(i) + appendix
        if os.path.exists(filepath):
            print(filepath)
            srcfh = open(filepath, 'r')
            dat.extend(json.load(srcfh))
            srcfh.close()
    fh = open(destpath, 'w')
    json.dump(dat, fh)
    fh.close()

def _make_norm_ldmk_MTCore(\
        outname_pre, subject_code, src_dfile, src_ifile, src_labfile,\
        norm_len, interval=5):
    if_src_raw = False
    expected_shape=(49,49,3)

    aif_dst_ftimg = aif(outname_pre + subject_code +"_ftimgd", outname_pre + subject_code +"_ftimgi")
    aif_dst_ftimg.wbegin()
    aif_dst_norm = aif(outname_pre + subject_code + "_normd", outname_pre + subject_code +"_normi")
    aif_dst_norm.wbegin()
    json_dst_labels = []
    fh_label = open(outname_pre + subject_code + "_label.json", 'w')
    print("on going: " + src_dfile)

    #label src file
    
    print("corresponding label: " + src_labfile)
    rppg_lab = RPPGLabel(src_labfile)

    #features & labels
    aif_inst = aif(src_dfile, src_ifile)
    dat_info = aif_inst.rbegin()

    # d_norm = []# [[[mean1,mean2,mean3],[std1,std2,std3]]...478]
    d_v300 = deque()

    total_none = 0

    for i in range(0, len(dat_info), interval):# frames within file
        if if_src_raw:
            img0 = (aif_inst.extract_raw(i)).copy()
        else:
            img0 = aif_inst.extract(i)
        # img0 = cv2.resize(img0, (144, 144))

        if (i + interval) < len(dat_info):
            # make norm
            # if i < norm_len and img0.shape == expected_shape:
            #     d_v300.append(img0) # 300 frames * 478 crosses
            # else:
            #     pass

            print('person: {}, frame num: {}'.format(subject_code,i))
            #one norm, one ft, one label
            if if_src_raw:
                img1 = (aif_inst.extract_raw(i + interval)).copy()
            else:
                img1 = aif_inst.extract(i + interval)

            if img1.shape == expected_shape and img0.shape == expected_shape:
                xdiff = img1 - img0 + 128
                if i < norm_len:
                    d_v300.append(xdiff) # 300 frames * 478 crosses
                else:
                    # make norm
                    cur_norm = get_norm(d_v300)
                    aif_dst_norm.append_any(np.array(cur_norm, dtype="float32").tobytes())
                    d_v300.popleft()
                    d_v300.append(xdiff)
                    # adjust to the median position
                    ftimg = xdiff
                    aif_dst_ftimg.append(ftimg)# features
                    xlab = rppg_lab.get_label_by_frame(i)
                    if xlab == None:
                        total_none += 1
                    json_dst_labels.append(xlab) #labels
            else:
                aif_dst_ftimg.append_any(np.array([]).tobytes())
                json_dst_labels.append(None)
                aif_dst_norm.append_any(np.array([]).tobytes())                        

        else:# end of a file, append an empty node
            aif_dst_ftimg.append_any(np.array([]).tobytes())
            json_dst_labels.append(None)
            aif_dst_norm.append_any(np.array([]).tobytes())
    aif_inst.rfinish()

    json.dump(json_dst_labels, fh_label)
    fh_label.close()
    aif_dst_ftimg.wfinish()
    aif_dst_norm.wfinish()
    print("total_none: " + str(total_none))

def _make_norm_pv_ldmk_MTCore(outname_pre, subject_code, src_dfile, src_ifile, src_labfile, norm_len):
    expected_shape=(49,49,3)

    aif_dst_ftimg = aif(outname_pre + subject_code +"_ftimgd", outname_pre + subject_code +"_ftimgi")
    aif_dst_ftimg.wbegin()
    aif_dst_norm = aif(outname_pre + subject_code + "_normd", outname_pre + subject_code +"_normi")
    aif_dst_norm.wbegin()
    json_dst_labels = []
    fh_label = open(outname_pre + subject_code + "_label.json", 'w')
    print("on going: " + src_dfile)

    #label src file
    
    print("corresponding label: " + src_labfile)
    rppg_lab = RPPGLabel(src_labfile)

    #features & labels
    aif_inst = aif(src_dfile, src_ifile)
    dat_info = aif_inst.rbegin()

    # loop here
    # make mixed label
    lab_coll = rppg_lab.get_frm_num_at_pvs()
    pv_frmnums_coll = []
    p_or_v = []
    niter = len(lab_coll[0]) if len(lab_coll[0]) > len(lab_coll[1]) else len(lab_coll[1])
    for xl in range(0, niter):
        if xl < len(lab_coll[0]):
            pv_frmnums_coll.append(lab_coll[0][xl])
            p_or_v.append(0)
        if xl < len(lab_coll[1]):
            pv_frmnums_coll.append(lab_coll[1][xl])
            p_or_v.append(1)
    # make featrue and norm
    loop_idx = 0
    for xpfrm in pv_frmnums_coll:
        # ret norm = []# [[[mean1,mean2,mean3],[std1,std2,std3]]...478]
        norm_buf = []
        if xpfrm < len(dat_info):
            img0 = aif_inst.extract(xpfrm)
            if xpfrm > norm_len and img0.shape == expected_shape:
                print(subject_code + ": " + str(xpfrm))
                # featlab
                aif_dst_ftimg.append(img0)
                json_dst_labels.append(p_or_v[loop_idx])
                # norm
                for i in range(xpfrm - norm_len, xpfrm):
                    norm_img = aif_inst.extract(xpfrm)
                    if norm_img.shape == expected_shape:
                        norm_buf.append(norm_img)
                cur_norm = get_norm(norm_buf)
                aif_dst_norm.append_any(np.array(cur_norm, dtype="float32").tobytes())
            loop_idx += 1
            
    aif_inst.rfinish()

    json.dump(json_dst_labels, fh_label)
    fh_label.close()
    aif_dst_ftimg.wfinish()
    aif_dst_norm.wfinish()

def make_normalized_landmarks_dbMT(outputfolder, basefolder, subject_list,\
                                   keyfunc, norm_len=60):
    pool = Pool(14)
    for root, dirs, files in os.walk(basefolder):
        for dname in dirs:
            prev_path = os.path.join(root, dname) + "\\"
            # xsplit = prev_path.split("\\")
            dfilename = prev_path + "datafile"
            difilename = prev_path + "idxfile"
            if os.path.exists(dfilename) and os.path.dirname(dfilename)[-1] == '1':# get only the first video
                src_label_fname = root + "\\" + "bvp_" + root.split("\\")[-1] + "_T" + os.path.dirname(dfilename)[-1] + ".csv"
                subject_code = root.split("\\")[-1]
                if len(subject_list) != 0 and subject_code not in subject_list:
                    break
                # MT
                pool.apply_async(func=keyfunc, args=(\
                    outputfolder, subject_code, dfilename,\
                    difilename, src_label_fname, norm_len,))
    pool.close()
    pool.join()

if __name__ == "__main__":
    # inter-frame plan
    # make_normalized_landmarks_dbMT(\
    #     "G:\\ecg\\AIFDatabase\\ldmk_intermediate\\interframe\\", "G:\\ecg\\raw",\
    #     subject_list=[], keyfunc=_make_norm_ldmk_MTCore)

    # _make_norm_ldmk_MTCore("G:/ecg/AIFDatabase/ldmk_intermediate/interframe/","s3",\
    #                     "G:/ecg/raw/s3/f_imgs1/datafile",\
    #                     "G:/ecg/raw/s3/f_imgs1/idxfile",\
    #                     "G:/ecg/raw/s3/bvp_s3_T1.csv",60)

    # only peaks and valleys ---- NO GOOD!
    # MT
    # make_normalized_landmarks_dbMT("G:\\ecg\\AIFDatabase\\ldmk_intermediate\\", "G:\\ecg\\raw",\
    #     subject_list=['s2'],\
    #     keyfunc=_make_norm_pv_ldmk_MTCore)
    # 
    # ST
    # _make_norm_pv_ldmk_MTCore("G:\\ecg\\AIFDatabase\\ldmk_intermediate\\peakvalley\\", "s2",\
    #                         "G:\\ecg\\raw\\s2\\f_imgs1\\datafile", "G:\\ecg\\raw\\s2\\f_imgs1\\idxfile",\
    #                         "G:\\ecg\\raw\\s2\\bvp_s2_T1.csv", 60)

    # image
    # merge_aifs_digiapdx("G:/ecg/AIFDatabase/ldmk_intermediate/interframe/",\
    #                     "s", "_ftimgd", "_ftimgi", 2, 38,\
    #                     "G:/ecg/AIFDatabase/ldmk_aggr/feature",dtype='uint8')
    # norm
    merge_aifs_digiapdx("G:/ecg/AIFDatabase/ldmk_intermediate/interframe/",\
                        "s", "_normd", "_normi", 2, 38, "G:/ecg/AIFDatabase/ldmk_aggr/norm",\
                        extract="bytes",dtype='float32')
    # # labels
    # merge_jsons("G:/ecg/AIFDatabase/ldmk_intermediate/interframe/",\
    #             "s", "_label.json", 2, 38, "G:/ecg/AIFDatabase/ldmk_aggr/label.json")