import numpy as np

import pandas as pd

import os

import click

import glob

import cv2

import pydicom

from tqdm import tqdm

from utils import get_windowing, window_image

from joblib import delayed, Parallel

@click.group()

def cli():

    print("CLI")

windows_range = {

    'brain': [40, 80],

    'bone': [600, 2800],

    'subdual': [75, 215]

}

def convert_dicom_to_jpg(dicomfile, outputdir):

    try:

        data = pydicom.read_file(dicomfile)

        image = data.pixel_array

        window_center, window_width, intercept, slope = get_windowing(data)

        id = dicomfile.split("/")[-1].split(".")[0]

        images = []

        for k, v in windows_range.items():

            image_windowed = window_image(image, v[0], v[1], intercept, slope)

            images.append(image_windowed)

        images = np.asarray(images).transpose((1, 2, 0))

        output_image = os.path.join(outputdir, id + ".jpg")

        cv2.imwrite(output_image, images)

    except:

        print(dicomfile)

@cli.command()

@click.option('--inputdir', type=str)

@click.option('--outputdir', type=str)

def extract_images(

    inputdir,

    outputdir,

):

    os.makedirs(outputdir, exist_ok=True)

    files = glob.glob(inputdir + "/*.dcm")

    Parallel(n_jobs=8)(delayed(convert_dicom_to_jpg)(file, outputdir) for file in tqdm(files, total=len(files)))

def split_by_patient(

    train_csv,

    train_meta_csv,

    n_folds,

    outdir

):

    os.makedirs(outdir, exist_ok=True)

    train_df = pd.read_csv(train_csv)

    train_meta_df = pd.read_csv(train_meta_csv)

    train_meta_df['ID'] = train_meta_df['ID'].apply(lambda x: "_".join(x.split("_")[:2]))

    train_meta_df = train_meta_df[['ID', 'PatientID']]

if __name__ == '__main__':

    cli()