import os

import numpy as np
from skimage.io import imread


def get_file_count(paths, image_format='.tif'):
    total_count = 0
    for path in paths:
        try:
            path_list = [_ for _ in os.listdir(path) if _.endswith(image_format)]
            total_count += len(path_list)
        except OSError:
            print("Directory does not exist. Returned file count for this path will be 0")
    return total_count


# Function to load image
def load_image(img_path):
    img = imread(img_path)
    if img.shape[2] == 4:
        img = img[:, :, :-1]
    # img = np.roll(img, shift=1, axis=2)  # CHECK IMAGE FORMAT
    return img


# Function to load mask
def load_mask(mask_path):
    mask = imread(mask_path)
    return mask


def load_mask_from_img(cfg, img_path, img_name, suffixes):
    a_mask = imread(os.path.join(img_path, img_name + suffixes[0]))
    msk = np.zeros((a_mask.shape[0], a_mask.shape[1], len(suffixes) * cfg.NUMBER_MSK_CHANNELS))
    i = 0
    for suffix in suffixes:
        msk_channel = imread(os.path.join(img_path, img_name + suffix))
        if len(msk_channel.shape) == 2:
            msk_channel = np.expand_dims(msk_channel, axis=-1)
        if len(msk_channel.shape) != 3:
            raise ValueError("Mask must be 3-dim here. Does your mask have 1 or more than 3 dimensions? "
                             "Check the masks.")
        msk[:, :, i:i+cfg.NUMBER_MSK_CHANNELS] = msk_channel
        i += cfg.NUMBER_MSK_CHANNELS
    # print(msk, msk.shape)
    return msk


def load_weights(cfg, img_path, img_name, weight_suffixes):
    a_weights = np.load(os.path.join(img_path, img_name + weight_suffixes[0]))
    weights = np.zeros((a_weights.shape[0], a_weights.shape[1], len(weight_suffixes) * cfg.NUMBER_MSK_CHANNELS))
    i = 0
    for suffix in weight_suffixes:
        weights_channel = np.load(os.path.join(img_path, img_name + suffix))
        if len(weights_channel.shape) == 2:
            weights_channel = np.expand_dims(weights_channel, axis=-1)
        if len(weights_channel.shape) != 3:
            raise ValueError("Weights must be 3-dim here. Has your weights 1 or more than 3 dimensions? Check the weights.")
        weights[:, :, i:i+cfg.NUMBER_MSK_CHANNELS] = weights_channel
        i += cfg.NUMBER_MSK_CHANNELS
    return weights
[((20, 10, 20, 26), 'skimage.io.imread', 'imread', ({(20, 17, 20, 25): 'img_path'}, {}), '(img_path)', False, 'from skimage.io import imread\n'), ((29, 11, 29, 28), 'skimage.io.imread', 'imread', ({(29, 18, 29, 27): 'mask_path'}, {}), '(mask_path)', False, 'from skimage.io import imread\n'), ((34, 20, 34, 66), 'os.path.join', 'os.path.join', ({(34, 33, 34, 41): 'img_path', (34, 43, 34, 65): 'img_name + suffixes[0]'}, {}), '(img_path, img_name + suffixes[0])', False, 'import os\n'), ((51, 24, 51, 77), 'os.path.join', 'os.path.join', ({(51, 37, 51, 45): 'img_path', (51, 47, 51, 76): 'img_name + weight_suffixes[0]'}, {}), '(img_path, img_name + weight_suffixes[0])', False, 'import os\n'), ((38, 29, 38, 70), 'os.path.join', 'os.path.join', ({(38, 42, 38, 50): 'img_path', (38, 52, 38, 69): 'img_name + suffix'}, {}), '(img_path, img_name + suffix)', False, 'import os\n'), ((40, 26, 40, 62), 'numpy.expand_dims', 'np.expand_dims', (), '', True, 'import numpy as np\n'), ((55, 34, 55, 75), 'os.path.join', 'os.path.join', ({(55, 47, 55, 55): 'img_path', (55, 57, 55, 74): 'img_name + suffix'}, {}), '(img_path, img_name + suffix)', False, 'import os\n'), ((57, 30, 57, 70), 'numpy.expand_dims', 'np.expand_dims', (), '', True, 'import numpy as np\n'), ((11, 36, 11, 52), 'os.listdir', 'os.listdir', ({(11, 47, 11, 51): 'path'}, {}), '(path)', False, 'import os\n')]