import os

import click

from tqdm import tqdm

import numpy as np

import pandas as pd

import cv2

@click.command()

@click.option('--path_root_in', help='E.g. data/31_OKOA_full_meta')

@click.option('--spacing_in', nargs=2, default=(0.5859375, 0.5859375))

@click.option('--path_root_out', help='E.g. data/32_OKOA_full_meta_rescaled')

@click.option('--spacing_out', nargs=2, default=(0.36458333, 0.36458333))

@click.option('--dirname_images', default='images')

@click.option('--dirname_masks', default='masks')

@click.option('--num_threads', default=12, type=click.IntRange(-1, 12))

@click.option('--margin', default=0, type=int)

@click.option('--update_meta', is_flag=True)

def main(**config):

    # Get the index of image files and the corresponding metadata

    path_meta = os.path.join(config['path_root_in'], 'meta_base.csv')

    if os.path.exists(path_meta):

        pass

    else:

        path_meta = os.path.join(config['path_root_in'], 'meta_dynamic.csv')

    df_meta = pd.read_csv(path_meta,

                          dtype={'patient': str,

                                 'release': str,

                                 'prefix_var': str,

                                 'sequence': str,

                                 'side': str,

                                 'slice_idx': int,

                                 'pixel_spacing_0': float,

                                 'pixel_spacing_1': float,

                                 'slice_thickness': float,

                                 'KL': int,

                                 'has_mask': int},

                          index_col=False)

    df_in = df_meta.sort_values(['patient', 'release', 'sequence', 'side', 'slice_idx'])

    ratio = (np.asarray(config['spacing_in']) /

             np.asarray(config['spacing_out']))

    groupers_stack = ['patient', 'release', 'sequence', 'side', 'slice_idx']

    # Resample images

    if config['dirname_images'] is not None:

        for name_gb, df_gb in tqdm(df_in.groupby(groupers_stack), desc='Resample images'):

            patient, release, sequence, side, slice_idx = name_gb

            fn_base = f'{slice_idx:03d}.png'

            dir_in = os.path.join(config['path_root_in'],

                                  patient, release, sequence,

                                  config['dirname_images'])

            dir_out = os.path.join(config['path_root_out'],

                                   patient, release, sequence,

                                   config['dirname_images'])

            os.makedirs(dir_out, exist_ok=True)

            path_in = os.path.join(dir_in, fn_base)

            path_out = os.path.join(dir_out, fn_base)

            img_in = cv2.imread(path_in, cv2.IMREAD_GRAYSCALE)

            if config['margin'] == 0:

                tmp = img_in

            else:

                tmp = img_in[config['margin']:-config['margin'],

                             config['margin']:-config['margin']]

            shape_out = tuple(np.floor(tmp.shape * ratio).astype(np.int))[::-1]

            tmp = cv2.resize(tmp, shape_out)

            img_out = tmp

            cv2.imwrite(path_out, img_out)

    # Resample masks

    if config['dirname_masks'] is not None:

        for name_gb, df_gb in tqdm(df_in.groupby(groupers_stack), desc='Resample masks'):

            patient, release, sequence, side, slice_idx = name_gb

            fn_base = f'{slice_idx:03d}.png'

            dir_in = os.path.join(config['path_root_in'],

                                  patient, release, sequence,

                                  config['dirname_masks'])

            dir_out = os.path.join(config['path_root_out'],

                                   patient, release, sequence,

                                   config['dirname_masks'])

            os.makedirs(dir_out, exist_ok=True)

            path_in = os.path.join(dir_in, fn_base)

            if not os.path.exists(path_in):

                print(f'No mask found for {name_gb}')

                continue

            path_out = os.path.join(dir_out, fn_base)

            mask_in = cv2.imread(path_in, cv2.IMREAD_GRAYSCALE)

            if config['margin'] == 0:

                tmp = mask_in

            else:

                tmp = mask_in[config['margin']:-config['margin'],

                              config['margin']:-config['margin']]

            shape_out = tuple(np.floor(tmp.shape * ratio).astype(np.int))[::-1]

            tmp = cv2.resize(tmp, shape_out, interpolation=cv2.INTER_NEAREST)

            mask_out = tmp

            cv2.imwrite(path_out, mask_out)

    if config['update_meta']:

        df_out = (df_in.assign(pixel_spacing_0=config['spacing_out'][0])

                       .assign(pixel_spacing_1=config['spacing_out'][1]))

        df_out.to_csv(os.path.join(config['path_root_out'], 'meta_base.csv'), index=False)

if __name__ == '__main__':

    main()