import numpy as np

import pandas as pd

from PIL import Image

from tqdm import tqdm

from preprocessingutils import pwr_transform

import os

from pathlib import Path

from argparse import ArgumentParser

parser = ArgumentParser()

parser.add_argument('--dir', default='', type=str,

                    help='location of files, should contain subdirs with splits with subdirs imgs and masks; if empty, starts from parent dir of this folder')

def measure_slice(x, mask=None):

    '''

    calculate slice size and variance

    '''

    assert isinstance(x, np.ndarray)

    if mask is None:

        mask = np.ones_like(x)

    assert isinstance(mask, np.ndarray)

    size     = mask.sum()

    variance = x[np.nonzero(mask)].std() 

    img_max  = x[np.nonzero(mask)].max()

    img_min  = x[np.nonzero(mask)].min()

    img_mean = x[np.nonzero(mask)].mean()

    return size, variance, img_max, img_min, img_mean

args = parser.parse_args()

if args.dir == '':

    args.dir = Path(Path.cwd().parent / 'data' / 'nodules2d')

dfs = {}

imgs = os.listdir(os.path.join(args.dir, "imgs"))

imgs = [x for x in imgs if x.endswith('.png')]

df = pd.DataFrame({'img_name': imgs, 

                    'size': np.zeros((len(imgs, ))), 

                    'variance': np.zeros((len(imgs,))),

                    'max': np.zeros((len(imgs,))),

                    'min': np.zeros((len(imgs,))),

                    'mean': np.zeros((len(imgs,)))})

for i, img_name in tqdm(enumerate(imgs)):

    img  = np.array(Image.open(os.path.join(args.dir, "imgs", img_name))) / 255

    mask = (np.array(Image.open(os.path.join(args.dir, "masks", img_name))) / 255).astype(np.int16)

    size, variance, img_max, img_min, img_mean = measure_slice(img, mask)

    df.iloc[i, 1] = size

    df.iloc[i, 2] = variance

    df.iloc[i, 3] = img_max

    df.iloc[i, 4] = img_min

    df.iloc[i, 5] = img_mean

df = df.rename(columns={'img_name': 'name'})

# df["name"] = df.name.apply(lambda x: os.path.join("imgs", x))

# normalize values

# scalar_vars = ["size", "variance", "img_min", "img_max", "img_mean"]

# df[scalar_vars] = df[scalar_vars].apply(pwr_transform)

# print(df.head())

df.to_csv(os.path.join(args.dir, "measurements.csv"), index=False)