import os
from typing import List, Tuple

import cv2
import numpy as np
from sklearn.model_selection import train_test_split
import rasterio

from pykelper.files_manager import DIRS

# Define paths
trainImagesPath = DIRS['TRAIN_FEATURES']
trainLabelsPath = DIRS['TRAIN_LABELS']


def read_tiff(file_path: str) -> np.ndarray:
    print('reading...')
    print(file_path)
    with rasterio.open(file_path) as src:
        image = src.read()
        # src.meta may hold interesting info not sure.
    return image


def load_files_as_ndarray(folder_path, limit) -> List[Tuple[str, np.ndarray]]:
    # Load all images and labels
    # TODO: add preproc here or no?
    return [
        (file_name,
         read_tiff(
             os.path.join(folder_path, file_name)
         ))for file_name in sorted(os.listdir(folder_path)[:limit])
    ]

def load_specific_file_as_ndarray(file_full_path):
    return read_tiff(file_full_path)