import glob
import numpy as np

class Data:
    
    def __init__(self, path, random=False):
        """
        input:
            path: path to the folder with subfolders: DSM, PAN, LABEL
            max_num: int, num of samples
            random: bool, to load samples randomly or from 0 to num_max
        """
        self.DSM = sorted(glob.glob(path+"/DSM/*.tif"))
        self.PAN = sorted(glob.glob(path+"/PAN/*.tif"))
        self.LABEL = sorted(glob.glob(path+"/LABEL/*.tif"))
        if len(self.DSM) != len(self.PAN) or len(self.LABEL) != len(self.PAN):
            raise ValueError('DSM, PAN or LABEL do not match')
      
    def get_data(self, start=0, num=10, as_arr=True, random=False):
        """
        function: load max_num of XY into lists
        output: list of numpy arrays, X (images) and Y (labels)
        """
        DSM_out = []
        PAN_out = []
        LABEL_out = []
      
        if random:
            idx = np.random.choice(list(range(len(self.X))), num, replace=False)
            print('randomly loading {0} tiles from {1} tiles'.format(num, len(self.DSM))) 
        else:
            idx = list(range(start, start+num))
            print('loading {0} - {1} image tiles'.format(start, start+num-1))

        for i in idx:
            DSM_out.append(np.moveaxis(rasterio.open(self.DSM[i]).read(),0,2))
            PAN_out.append(np.moveaxis(rasterio.open(self.PAN[i]).read(),0,2))
            LABEL_out.append(np.moveaxis(rasterio.open(self.LABEL[i]).read(),0,2))
        
        DSM_remove = [self.DSM[i] for i in idx]
        PAN_remove = [self.PAN[i] for i in idx]
        LABEL_remove = [self.LABEL[i] for i in idx]
        
        for i in range(len(DSM_remove)):
            self.DSM.remove(DSM_remove[i])
            self.PAN.remove(PAN_remove[i])
            self.LABEL.remove(LABEL_remove[i])
        
        if as_arr:
            return np.asarray(DSM_out), np.asarray(PAN_out), np.asarray(LABEL_out)
        else:
            return DSM_out, PAN_out, LABEL_out
           
    def split_trn_vld_tst(self, vld_rate=0.2, tst_rate=0.0, random=True, seed=10):
        np.random.seed(seed)

        num = len(self.DSM)
        vld_num = int(num*vld_rate)
        tst_num = int(num*tst_rate)
        
        print('split into {0} train, {1} validation, {2} test samples'.format(num-vld_num-tst_num, vld_num, tst_num))
        idx = np.arange(num)
        if random:
            np.random.shuffle(idx)
        DSM_tst, PAN_tst, LABEL_tst = [self.DSM[k] for k in idx[:tst_num]], [self.PAN[k] for k in idx[:tst_num]], [self.LABEL[k] for k in idx[:tst_num]]
        DSM_vld, PAN_vld, LABEL_vld = [self.DSM[k] for k in idx[tst_num:tst_num+vld_num]], [self.PAN[k] for k in idx[tst_num:tst_num+vld_num]], [self.LABEL[k] for k in idx[tst_num:tst_num+vld_num]]
        DSM_trn, PAN_trn, LABEL_trn = [self.DSM[k] for k in idx[tst_num+vld_num:]], [self.PAN[k] for k in idx[tst_num+vld_num:]], [self.LABEL[k] for k in idx[tst_num+vld_num:]]
        
        
        return DSM_trn, PAN_trn, LABEL_trn, DSM_vld, PAN_vld, LABEL_vld, DSM_tst, PAN_tst, LABEL_tst

[((55, 8, 55, 28), 'numpy.random.seed', 'np.random.seed', ({(55, 23, 55, 27): 'seed'}, {}), '(seed)', True, 'import numpy as np\n'), ((62, 14, 62, 28), 'numpy.arange', 'np.arange', ({(62, 24, 62, 27): 'num'}, {}), '(num)', True, 'import numpy as np\n'), ((13, 26, 13, 54), 'glob.glob', 'glob.glob', ({(13, 36, 13, 53): "path + '/DSM/*.tif'"}, {}), "(path + '/DSM/*.tif')", False, 'import glob\n'), ((14, 26, 14, 54), 'glob.glob', 'glob.glob', ({(14, 36, 14, 53): "path + '/PAN/*.tif'"}, {}), "(path + '/PAN/*.tif')", False, 'import glob\n'), ((15, 28, 15, 58), 'glob.glob', 'glob.glob', ({(15, 38, 15, 57): "path + '/LABEL/*.tif'"}, {}), "(path + '/LABEL/*.tif')", False, 'import glob\n'), ((64, 12, 64, 34), 'numpy.random.shuffle', 'np.random.shuffle', ({(64, 30, 64, 33): 'idx'}, {}), '(idx)', True, 'import numpy as np\n'), ((50, 19, 50, 38), 'numpy.asarray', 'np.asarray', ({(50, 30, 50, 37): 'DSM_out'}, {}), '(DSM_out)', True, 'import numpy as np\n'), ((50, 40, 50, 59), 'numpy.asarray', 'np.asarray', ({(50, 51, 50, 58): 'PAN_out'}, {}), '(PAN_out)', True, 'import numpy as np\n'), ((50, 61, 50, 82), 'numpy.asarray', 'np.asarray', ({(50, 72, 50, 81): 'LABEL_out'}, {}), '(LABEL_out)', True, 'import numpy as np\n')]