import os, shutil, sys

from pathlib import Path

import numpy as np

import pandas as pd

from tqdm import tqdm

import re

from PIL import Image

sys.path.append("./")

from preprocessingutils import pwr_transform

from argparse import ArgumentParser

parser = ArgumentParser()

parser.add_argument('--out-dir', default='', help='place to put the slices (leave empty for default)')

parser.add_argument('--in-dir', default='', help='input directory (leave empty for default)')

parser.add_argument('--min-size', default=20, type=int, help='minimal size of noduls in mm')

parser.add_argument('--splits', default='train,valid', help='which spits to create, separate with ,')

parser.add_argument('--manual-seed', default=12345, help='seed for generating splits')

parser.add_argument('--silent', action='store_false', dest='verbose', help='dont print stuff')

parser.add_argument('--valid_min', default=int(1000), type=int, help='minimal number of validation and test imgs, otherwise do 20%')

parser.add_argument('--no-imgs', action='store_false', dest='copy_imgs', help='dont copy, only update dataframe')

parser.add_argument('--out-size', default=70, type=int, help='out size, uses center crop, if None, no resizing will be done')

parser.add_argument('--test', action='store_true', help='do only a few imgs for testing the code')

parser.set_defaults(verbose=True, copy_imgs=True, test=False)

def get_center_crop_bbox(in_size, out_size):

    """

    get bounding box for center crop

    in_size is (width, height for PIL.Image)

    """

    center = np.array(in_size) / 2

    left   = int(center[0] - out_size / 2)

    right  = int(left + out_size)

    upper  = int(center[0] - out_size / 2)

    lower  = int(upper + out_size)

    return (left, upper, right, lower)

def main(args):

    # find location for resources

    resourcedir = Path.cwd().parent / 'resources'

    # load dataframe with annotation data per nodule, made in the step 'lidc-preprocessing'

    df_ann = pd.read_csv(resourcedir / "annotation_df.csv")

    # show source files

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

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

    if args.verbose:

        print(f"found {len(imgs)} files")

    # img files are like 0001n01a2s086.png

    # imgs = imgs[:10]

    pids      = [x.split("n")[0] for x in imgs]

    nods      = [re.search(r"(?<=n)\d+", x).group() for x in imgs]

    anns      = [re.search(r"(?<=a)\d", x).group() for x in imgs]

    zvals     = [re.search(r"(?<=s)\d+", x).group() for x in imgs]

    ann_ids   = [x.split("s")[0] for x in imgs]

    nod_ids   = [x.split("a")[0] for x in imgs]

    slice_ids = [x.split(".png")[0] for x in imgs]

    nodule_slice_ids = [f"{nod_id}s{zval}" for nod_id, zval in zip(nod_ids, zvals)]

    slice_df = pd.DataFrame({

        'in_name': imgs,

        'pid': pids,

        'nodule_idx': nods,

        "annotation_idx": anns,

        "annotation_id": ann_ids,

        "nodule_id": nod_ids,

        "zval": zvals,

        "slice_id": slice_ids,

        "nodule_slice_id": nodule_slice_ids 

    })

    # add max number of annotations per nodule

    annotation_counts = df_ann.groupby('nodule_id').nodule_id.count().reset_index(name="annotation_count")

    slice_df = pd.merge(slice_df, annotation_counts, on="nodule_id")

    max_annotation_count_pid = slice_df.groupby("pid").annotation_count.max().reset_index(name='max_ann_count_per_pid')

    slice_df = pd.merge(slice_df, max_annotation_count_pid, on="pid")

    slice_counts = slice_df.groupby(["nodule_id", "zval"]).size().reset_index(name="slice_count")

    slice_df = pd.merge(slice_df, slice_counts, on=["nodule_id", "zval"])

    slice_df["all_anns_agree"] = slice_df.slice_count == slice_df.max_ann_count_per_pid

    slice_df.to_csv(resourcedir / "slice_df.csv", index=False)

    # keep only those slices where all annotators included the slice in their segmentation

    df = slice_df[(slice_df.all_anns_agree)]

    # import measurements

    measurements = pd.read_csv(os.path.join(args.in_dir, "measurements.csv"))

    df = pd.merge(df, measurements, left_on="in_name", right_on="name")

    # keep only slices greater than the cutoff

    df = df[df["size"] > args.min_size]

    print(f"number of slices left: {len(df)}")

    slices_per_pid    = len(df) / len(df.pid.unique())

    # divide by 4 because only 1 of the annotations gets selected

    slices_per_nodule = (len(df) / len(df.nodule_id.unique())) / 4

    np.random.seed(args.manual_seed)

    VALID_PROP = 0.3

    TEST_PROP  = 0.0

    df["uid"] = df.nodule_id

    # df.set_index("slice_id", drop=False, inplace=True)

    uids = df['uid'].unique().tolist()

    # valid_size  = int(min(args.valid_min, int(len(uids) * VALID_PROP * slices_per_pid / 4)) / (slices_per_pid / 4))

    # test_size   = int(min(args.valid_min, int(len(uids) * TEST_PROP * slices_per_pid / 4)) / (slices_per_pid / 4))

    valid_size  = int(min(args.valid_min, int(len(uids) * VALID_PROP * slices_per_nodule)) / (slices_per_nodule))

    test_size   = int(min(args.valid_min, int(len(uids) * TEST_PROP * slices_per_nodule)) / (slices_per_nodule))

    test_uids  = list(np.random.choice(uids, replace = False, size = test_size))

    valid_uids = list(np.random.choice(list(set(uids) - set(test_uids)), size = valid_size))

    train_uids = list(set(uids) - (set(valid_uids +  test_uids)))

    split_dict = dict(zip(train_uids + valid_uids + test_uids,

                        ["train"] *len(train_uids) + ["valid"]*len(valid_uids) + ["test"] * len(test_uids)))

    df["split"] = df.uid.map(split_dict)

    # normalize continuous variables

    cont_vars = ["size", "variance", "min", "max", "mean"]

    train_idxs = np.where(df.uid.isin(train_uids))

    df[cont_vars] = df[cont_vars].apply(pwr_transform, train_ids=train_idxs)

    # average measurements over annotations, pick single slice per measurement

    df = df.groupby("nodule_slice_id").agg({

        'size': 'mean', 

        'variance': 'mean',

        "min": 'mean',

        "max": 'mean',

        "mean": 'mean',

        'in_name': 'first',

        'split': 'first',

    })

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

    if args.test:

        df = df.iloc[:10,]

    if args.out_size:

        print("resizing and saving images")

        # create output directories

        for split in args.splits.split(","):

            for subdir in ["imgs", "masks"]:

                if not os.path.isdir(os.path.join(args.out_dir, split, subdir)):

                    os.makedirs(os.path.join(args.out_dir, split, subdir))

        # crop and copy images

        for slice_id, row in tqdm(df.iterrows()):

            img = Image.open(os.path.join(args.in_dir, 'imgs', row['in_name']), 'r')

            img_crop = img.crop(get_center_crop_bbox(img.size, args.out_size))

            img_crop.save(os.path.join(args.out_dir, row["split"], "imgs", row["in_name"]))

            mask = Image.open(os.path.join(args.in_dir, 'masks', row['in_name']), 'r')

            mask_crop = mask.crop(get_center_crop_bbox(mask.size, args.out_size))

            mask_crop.save(os.path.join(args.out_dir, row["split"], "masks", row["in_name"]))

    else:

        if args.copy_imgs:

            print("copying images")

            for split in args.splits.split(","):

                for subdir in ["imgs", "masks"]:

                    if not os.path.isdir(os.path.join(args.out_dir, split, subdir)):

                        os.makedirs(os.path.join(args.out_dir, split, subdir))

            for slice_id, row in tqdm(df.iterrows()):

                shutil.copy(os.path.join(args.in_dir, 'imgs', row["in_name"]),

                            os.path.join(args.out_dir, row["split"], "imgs", row["in_name"]))

                shutil.copy(os.path.join(args.in_dir, 'masks', row["in_name"]),

                            os.path.join(args.out_dir, row["split"], 'masks', row["in_name"]))

    df.to_csv(os.path.join(args.out_dir, "labels.csv"), index=False)

    print(df.split.value_counts())

if __name__ == "__main__":

    args = parser.parse_args()

    if args.out_dir == '':

        args.out_dir = (Path.cwd().parent) / 'data' / 'slices'

    if args.in_dir == '':

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

    main(args)