import os

import numpy as np

import pandas as pd

import pylidc as pl

import pydicom

import pickle

from sklearn.preprocessing import PowerTransformer

def annotation_to_dict(ann):

    '''

    Flattens annotations into a single row pandas DataFrame

    '''

    d = {

        'patient_id': ann.scan.patient_id,

        'nodule_id_lidc': ann._nodule_id,

        'annotation_id_lidc': ann.id,

        'scan_id': ann.scan_id,

        'volume': ann.volume,

        'surface_area': ann.surface_area,

        'diameter': ann.diameter

    }

    feature_names = ['sublety', 'internalstructure', 'calcification',

               'sphericity', 'margin', 'lobulation', 'spiculation',

               'texture', 'malignancy']

    for feature, value in zip(feature_names, ann.feature_vals()):

        d[feature] = value

    return d

def annotation_to_df(ann):

    try:

        d = annotation_to_dict(ann)

        d = {k: [v] for k, v in d.items()}

        df = pd.DataFrame.from_dict(d)

    except:

        df = None

    return df

def annotation_list_to_df(anns):

    assert isinstance(anns, list)

    dfs = []

    for ann in anns:

        dfs.append(annotation_to_df(ann))

    df = pd.concat(dfs, ignore_index=True)

    df["annotation_idx"] = range(1, df.shape[0]+1)

    return df

def get_intercept_and_slope(scan):

    ''' 

    scan is the results of a pydicom query

    returns the intercept and slope

    adapted from https://www.kaggle.com/gzuidhof/full-preprocessing-tutorial

    '''

    dcm_path = scan.get_path_to_dicom_files()

    dcm_files = [x for x in os.listdir(dcm_path) if x.endswith(".dcm")]

    slice0 = pydicom.read_file(os.path.join(dcm_path, dcm_files[0]), stop_before_pixels=True)

    intercept = slice0.RescaleIntercept

    slope = slice0.RescaleSlope

    return intercept, slope

def resample_and_crop_annotation(ann_id, ann, nodule_path, mask_path=None, scan=None, size_mm = 50, export_mask = True):

    '''

    take an annotation, crop and resample

    size is the length of the sides of the resulting cube in millimeters

    '''

    if scan is None:

        scan = ann.scan

    intercept, slope = get_intercept_and_slope(scan)

    try:

        vol, mask = ann.uniform_cubic_resample(side_length = size_mm, verbose = True)

        if slope != 1:

            vol = slope * vol.astype(np.float64)

        vol = vol.astype(np.int16)

        vol += np.int16(intercept)

        np.save(os.path.join(nodule_path, ann_id+".npy"), vol)

        if export_mask:

            assert mask_path != None

            np.save(os.path.join(mask_path, ann_id+".npy"), mask)

        print("")

    except:

        print("-failed")

def crop_nodule_tight_z(ann, volume=None, scan=None, scan_spacing=None, out_size_cm = 5):

    """

    Get nodule cropped tightly in z direction, but of minimum dimension in xy plane

    """

    # print(f"trying to crop")

    if volume is None:

        if scan is None:

            scan = ann.scan

        volume = scan.to_volume()

    if scan_spacing is None:

        scan_spacing = scan.pixel_spacing

    # print(f"scan_spacing: {scan_spacing}")

    padding = get_padding_tight_z(ann, scan_spacing=scan_spacing, out_size_cm=out_size_cm)

    # print(f"padding: {padding}")

    mask = ann.boolean_mask(pad=padding)

    bbox = ann.bbox(pad=padding)

    zvals= ann.contour_slice_indices

    arr  = volume[bbox]

    return arr, mask, zvals

def get_padding_tight_z(ann, scan=None, scan_spacing=None, out_size_cm = None):

    """

    Get bbox dimensions base on a minimal size, restricting to no padding in z direction

    """

    if scan_spacing is None:

        if scan is None:

            scan_spacing = ann.scan.pixel_spacing

        else:

            scan_spacing = scan.pixel_spacing

    # return tight bounding box

    if out_size_cm is None:

        padding = [(int(0), int(0))] * 3

    else:

        # if len(out_size_cm == 3):

        #     if out_size_cm[2] is None:

        #         padding_z = (0,0)        

        out_shape = (np.ceil((out_size_cm * 10) / scan_spacing) * np.ones((2,))).astype(int)

        bb_mat   = ann.bbox_matrix()

        bb_shape = bb_mat[:,1] - bb_mat[:,0]

        paddings = out_shape - bb_shape[:2]

        # print(f"paddings: {paddings}")

        padding_x = (int(np.ceil(paddings[0] / 2)), int(np.floor(paddings[0] / 2)))

        padding_y = (int(np.ceil(paddings[1] / 2)), int(np.floor(paddings[1] / 2)))

        padding = [padding_x, padding_y, (int(0),int(0))]

    return padding

def normalize(x, window = None, level = None, in_min = -1000.0, in_max = 600.0, center=0.0):

    """

    Normalize array to values between 0 and 1, possibly clipping

    """

    assert type(x) is np.ndarray

    if (not window is None) & (not level is None) :

        in_min = level - (window / 2)

        in_max = level + (window / 2)

    x = x - in_min                 # add zero point

    x = x / (in_max - in_min)      # scale to unit

    x = x + center                 # adjust white-balance

    x = np.clip(x, 0.0, 1.0)       # clip to (0,1)

    return x

def normalized_to_8bit(x):

    assert ((x.min() >= 0) & (x.max() <= 1))

    x = (255 * x)

    return x.astype(np.uint8)

def normalize_to_8bit(x, *args, **kwargs):

    return normalized_to_8bit(normalize(x, *args, **kwargs))

def pwr_transform(x, train_ids=None):

    x  = np.array(x).reshape(-1,1)

    pt = PowerTransformer(method="yeo-johnson")

    if train_ids is None:

        pt.fit(x)

    else:

        pt.fit(x[train_ids])

    y = pt.transform(x)

    return np.squeeze(y)