# DeepLabV3+ Data Loader
import numpy as np
from glob import glob
from tqdm import tqdm
import tensorflow as tf
import tensorflow.data as tfd
import tensorflow.image as tfi
from tensorflow.keras.utils import load_img, img_to_array

from .config import IMAGE_SIZE, BATCH_SIZE, SHUFFLE_BUFFER


def load_image(path, image_size=IMAGE_SIZE):
    """단일 이미지 로드 및 전처리"""
    # Load Image
    image = load_img(path)

    # Convert to Array
    image = img_to_array(image)

    # Resize
    image = tfi.resize(image, (image_size, image_size))

    # Convert to float32
    image = tf.cast(image, tf.float32)

    # Standardization
    image = image / 255.0

    return image


def load_images(paths, image_size=IMAGE_SIZE):
    """여러 이미지 로드"""
    # Create Space
    images = np.zeros(shape=(len(paths), image_size, image_size, 3))

    # Iterate Paths
    for i, path in tqdm(enumerate(paths), desc="Loading", total=len(paths)):
        # Load Image
        images[i] = load_image(path, image_size=image_size)

    return images


def load_data(root_path, trim=None, image_size=IMAGE_SIZE, batch_size=BATCH_SIZE, return_arrays=False):
    """
    데이터셋 로드 및 tf.data.Dataset 생성

    Args:
        root_path: 이미지 폴더 경로 (Images 폴더)
        trim: [start, end] 인덱스로 데이터 범위 지정
        image_size: 이미지 크기
        batch_size: 배치 크기
        return_arrays: True면 (dataset, images, masks) 반환

    Returns:
        tf.data.Dataset 또는 (dataset, images, masks) 튜플
    """
    # Collect Image Paths
    image_paths = glob(root_path + "*.jpg")

    if not image_paths:
        raise ValueError(f"No images found in {root_path}")

    # Trim if needed
    if trim is not None:
        image_paths = image_paths[trim[0]:trim[1]]

    # Update mask paths
    mask_paths = []
    for path in image_paths:
        mask_path = path.replace("Images", "Masks")
        mask_paths.append(mask_path)

    # Load Data
    images = load_images(image_paths, image_size=image_size)
    masks = load_images(mask_paths, image_size=image_size)

    # Convert To tf.data.Dataset
    data = tfd.Dataset.from_tensor_slices((images, masks))
    data = data.batch(batch_size, drop_remainder=True)
    data = data.shuffle(SHUFFLE_BUFFER).prefetch(tfd.AUTOTUNE)

    if return_arrays:
        return data, images, masks
    return data


def create_dataset_from_paths(image_paths, mask_paths, image_size=IMAGE_SIZE, batch_size=BATCH_SIZE):
    """
    경로 리스트에서 직접 데이터셋 생성

    Args:
        image_paths: 이미지 경로 리스트
        mask_paths: 마스크 경로 리스트
        image_size: 이미지 크기
        batch_size: 배치 크기

    Returns:
        tf.data.Dataset
    """
    images = load_images(image_paths, image_size=image_size)
    masks = load_images(mask_paths, image_size=image_size)

    data = tfd.Dataset.from_tensor_slices((images, masks))
    data = data.batch(batch_size, drop_remainder=True)
    data = data.shuffle(SHUFFLE_BUFFER).prefetch(tfd.AUTOTUNE)

    return data
