import os

import subprocess

from distutils import dir_util

from pathlib import Path

import cv2 as cv

import numpy as np

from PIL import Image

from matplotlib import pyplot as plt

from tqdm import tqdm

from features import NucleiFeatures

Image.MAX_IMAGE_PIXELS = None

def get_x_and_y(name):

    x, y = os.path.splitext(name)[0].split('_')[-2:]

    return int(x), int(y)

def split_image(img, x_tiles_cnt=None, y_tiles_cnt=None, x_tile_size=None, y_tile_size=None, base='img'):

    """

    Splits an image array to smaller tiles for further segmentation.

    Specify tiles count OR tiles size.

    X axis means the arr.shape[1] coordinate, be careful!

    Tile names are used to restore the initial image after segmentation.

    Parameters

    ----------

    img : numpy ndarray

        The input image.

    x_tiles_cnt : integer

        Number of tiles along the x axis of img.

    y_tiles_cnt : integer

        Number of tiles along the y axis of img.

    x_tile_size : integer

        Size of tile along x axis.

    y_tile_size : integer

        Size of tile along y axis.

    base : str

        Base for tile names.

    Returns

    -------

    tiles : list

        List of tiles.

    tile_names : list

        List of tile names.

    """

    if (x_tile_size is not None) and (y_tile_size is not None):

        x_tiles_cnt = img.shape[1] // x_tile_size

        y_tiles_cnt = img.shape[0] // y_tile_size

    if (x_tiles_cnt is not None) and (y_tiles_cnt is not None):

        x_ticks = np.linspace(0, img.shape[1], x_tiles_cnt + 1).astype(int)

        y_ticks = np.linspace(0, img.shape[0], y_tiles_cnt + 1).astype(int)

    else:

        raise Exception('Specify tiles count OR tiles size.')

    tiles = []

    tile_names = []

    for x_num, x in enumerate(zip(x_ticks[:-1], x_ticks[1:])):

        for y_num, y in enumerate(zip(y_ticks[:-1], y_ticks[1:])):

            tiles.append(img[y[0]:y[1], x[0]:x[1]])

            tile_names.append(f'{base}_{x_num}_{y_num}')

    return tiles, tile_names

def prepare_test_data(tiles, tile_names, base_dir, force=False):

    """

    Saves data in the proper way.

    Parameters

    ----------

    tiles : list

        List of tiles.

    tile_names : list

        List of tile names.

    base_dir : str

        Full path to base directory, 'full/path/../data_test' in normal case.

    force : bool

        Rewrite existing files

    Returns

    -------

    None

    """

    base_dir = Path(base_dir)

    if not os.path.exists(base_dir):

        os.makedirs(base_dir)

    if not force and len(os.listdir(base_dir)) > 0:

        raise ValueError(f'base_dir {base_dir} is not empty, use force=True option if you want to rewrite files')

    elif len(os.listdir(base_dir)):

        dir_util.remove_tree(str(base_dir))

        os.makedirs(base_dir)

    for tile, name in zip(tiles, tile_names):

        if tile.max() <= 1:

            tile = (tile * 255).astype(np.uint8)

        os.mkdir(base_dir / name)

        os.mkdir(base_dir / name / 'images')

        cv.imwrite(str(base_dir / name / 'images' / f'{name}.png'), tile)

def restore_image(work_dir, tiff=False):

    """

    Restores the initial image.

        Parameters

    ----------

    work_dir : str

        Full path to directory with files.

    tiff : bool

        Is the target image a multilayer tiff or not

    Returns

    -------

    img : numpy ndarray

        Initial image

    """

    work_dir = Path(work_dir)

    file_names = sorted(os.listdir(work_dir), key=lambda x: get_x_and_y(x)[::-1])

    coords = np.array([get_x_and_y(n) for n in file_names])

    x_max, y_max = coords.max(axis=0)

    if tiff:

        tiles = {}

        max_number = int(0)

        for n in file_names:

            tmp = cv.imread(str(work_dir / n), -1)

            tmp = (tmp + max_number) * (tmp > 0)

            max_number = tmp.max()

            tiles[get_x_and_y(n)] = tmp.copy()

    else:

        tiles = {get_x_and_y(n): cv.imread(str(work_dir / n), -1) for n in file_names}

    long_tiles = []

    for y in range(y_max + 1):

        long_tiles.append([])

        for x in range(x_max + 1):

            long_tiles[-1].append(tiles[(x, y)])

    long_tiles = [np.hstack(i) for i in long_tiles]

    return np.vstack(long_tiles)

def perform_segmentation(full_img_path, sample_dir, network_dir, force=False, features=None):

    network_dir = Path(network_dir)

    try:

        full_img = cv.imread(full_img_path, -1)

    except cv.error:

        full_img = plt.imread(full_img_path)

    tiles, tile_names = split_image(img=full_img, x_tile_size=1000, y_tile_size=1000)

    prepare_test_data(tiles, tile_names, sample_dir, force=force)

    try:

        dir_util.remove_tree(str(network_dir / 'data_test'))

    except:

        pass

    os.mkdir(str(network_dir / 'data_test'))

    dir_util.copy_tree(sample_dir, str(network_dir / 'data_test'));

    try:

        dir_util.remove_tree(str(network_dir / 'predictions'))

    except:

        pass

    try:

        dir_util.remove_tree(str(network_dir / 'albu/results_test'))

    except:

        pass

    result_dir = str(Path(sample_dir)) + '_segmented'

    subprocess.run(f"cd {network_dir} && bash 'predict_test.sh'", shell=True)

    dir_util.copy_tree(str(network_dir / 'predictions'), result_dir);

    if features is not None:

        NucleiFeatures(f'{result_dir}/lgbm_test_sub2', sample_dir,

                       features=features).df().to_csv(f'{result_dir}/{os.path.split(sample_dir)[1]}.csv',

                                                      index=False)

def color_tiff(img, n=60):

    img = img % n

    seg_color = np.zeros((*img.shape, 3), dtype=np.uint8)

    for i in tqdm(range(1, img.max() + 1)):

        seg_color[img == i] = np.random.randint(0, 255, 3)

    return seg_color