import numpy as np

import pandas as pd

import random

from glob import glob

import os, shutil

from tqdm import tqdm

tqdm.pandas()

import time

from PIL import Image

import sys

def prepare_label_df(label_csv_path):

    df = pd.read_csv(label_csv_path)

    df["absent"] = df["segmentation"].map(lambda x: int(pd.isna(x))) # 1 means the organ is absent, 0 means its present

    df["case"] = df["id"].str.split('_').str[0]

    df["day"] = df["id"].str.split('_').str[1]

    df["slice_id"] = df["id"].str.split('_').str[3]

    return df

def prepare_image_df(dataset_folder_path, path_split):

    images = glob(dataset_folder_path + '/*/*/*/*.png')

    image_df = pd.DataFrame(images, columns=["image_path"])

    image_df["case"] = image_df["image_path"].str.split(path_split).str[1]

    image_df["day"] = image_df["image_path"].str.split(path_split).str[2].str.split('_').str[1]

    image_df["slice_id"] = image_df["image_path"].str.split(path_split).str[4].str.split('_').str[1]

    image_df["pic_info"] = image_df["image_path"].str.split(path_split).str[4]

    image_df["slice_height"] = image_df["pic_info"].str.split("_").str[2].astype(int)

    image_df["slice_width"] = image_df["pic_info"].str.split("_").str[3].astype(int)

    image_df["pixel_height"] = image_df["pic_info"].str.split("_").str[4].astype(float)

    image_df["pixel_width"] = image_df["pic_info"].str.split("_").str[5].str.split('.png').str[0].astype(float)

    return image_df

if __name__ == '__main__':

    # usage: python csv_process.py [label_csv_path] [dataset_folder_path] [output_csv_name] [path_splitter]

    label_csv_path = sys.argv[1]

    dataset_folder_path = sys.argv[2]

    output_csv_path = sys.argv[3]

    path_split = sys.argv[4]

    label_df = prepare_label_df(label_csv_path)

    image_df = prepare_image_df(dataset_folder_path, path_split)

    combined_df = pd.merge(label_df, image_df, how='left', on=['case','day','slice_id'])

    combined_df.to_csv(output_csv_path + ".csv")

    print("Combined CSV generated!")