import os
import numpy as np
import matplotlib.pyplot as plt

import cv2 as cv
import rasterio as rio

RASTERS_DIR = './data/data/train_features'
IMAGES_DIR = './images/'
FC_DIR = "fc"


def load_img(path, expand_dim=False):
    img = cv.imread(path)
    img = img / 255.0
    if expand_dim:
        img = np.expand_dims(img, axis=0)

    return img


def create_folder(name, path):
    n = os.path.join(path, name)
    if not os.path.exists(n):
        os.makedirs(n)
    return n


def scale(band):
    return band/np.max(band)


def convert_raster_to_image(rasters_dir,
                            rgb_path=None,
                            false_color_path=None):

    b2 = scale(rio.open(rasters_dir+'/B02.tif').read().reshape(512, 512, 1))
    b3 = scale(rio.open(rasters_dir+'/B03.tif').read().reshape(512, 512, 1))
    b4 = scale(rio.open(rasters_dir+'/B04.tif').read().reshape(512, 512, 1))
    b8 = scale(rio.open(rasters_dir+'/B08.tif').read().reshape(512, 512, 1))

    file_name = rasters_dir.split(os.sep)[-1]

    rgb = np.dstack([b4, b3, b2])

    plt.imsave(fname=rgb_path + f'/{file_name}.jpeg',
               arr=rgb)

    fc = np.dstack([b8, b3, b2])

    plt.imsave(fname=false_color_path + f'/{file_name}.jpeg',
               arr=fc)


def transform_photo(raster_dir):
    dp = create_folder(raster_dir, IMAGES_DIR)
    fc = create_folder(FC_DIR, dp)
    convert_raster_to_image(os.path.join(RASTERS_DIR, raster_dir), dp, fc)


if __name__ == "__main__":
    for d in os.listdir(RASTERS_DIR):
        transform_photo(d)