# %%
import gradio as gr
from sklearn import preprocessing
import collections
import SimpleITK as sitk
from radiomics import featureextractor
import os
import numpy as np
from sklearn.preprocessing import StandardScaler

# %%


def resize_image_itk(itkimage, newSize, resamplemethod=sitk.sitkLinear):

    resampler = sitk.ResampleImageFilter()
    originSize = itkimage.GetSize()  # 原来的体素块尺寸
    # print(originSize)
    originSpacing = itkimage.GetSpacing()
    newSize = np.array(newSize, float)
    factor = originSize / newSize
    newSpacing = originSpacing * factor
    newSize = newSize.astype(np.int0)  # spacing肯定不能是整数
    resampler.SetReferenceImage(itkimage)  # 需要重新采样的目标图像
    resampler.SetSize(newSize.tolist())
    resampler.SetOutputSpacing(newSpacing.tolist())
    resampler.SetTransform(sitk.Transform(3, sitk.sitkIdentity))
    resampler.SetInterpolator(resamplemethod)
    itkimgResampled = resampler.Execute(itkimage)  # 得到重新采样后的图像
    # print(itkimgResampled.GetSize())

    return itkimgResampled


def load_image(img):

    # imgfilePath = os.path.join(img_path, img_name)
    imgfilePath = img
    if not os.path.isfile(imgfilePath):
        print('Error: IMAGE DIR READ FAILED!')

    # mskfilePath = imgfilePath + '.img'
    '''
    # 读取.dcm Series
    reader = sitk.ImageSeriesReader()
    dcm_series = reader.GetGDCMSeriesFileNames(imgfilePath)
    reader.SetFileNames(dcm_series)
    image = reader.Execute()
    mask = sitk.ReadImage(mskfilePath)
    '''
    # 读取.png或.jpeg图像——可在此处更换读取IO适应不同格式图像
    if imgfilePath[-4:] == '.png':
        image = sitk.ReadImage(
            imgfilePath, imageIO="PNGImageIO", outputPixelType=sitk.sitkInt16)
    else:
        image = sitk.ReadImage(
            imgfilePath, imageIO="JPEGImageIO", outputPixelType=sitk.sitkInt16)

    # resize_img = resize_image_itk(image, newSize=(512, 512))

    # img_array = sitk.GetArrayFromImage(resize_img)

    return image  # np.asarray(image)


def feature_extractor(img):

    image = load_image(img)
    params = 'settings.yaml'
    if os.path.isfile(params):
        extractor = featureextractor.RadiomicsFeatureExtractor(
            params)

    else:
        settings = {}
        settings['binWidth'] = 25
        settings['resampledPixelSpacing'] = [3, 3, 3]
        settings['interpolator'] = sitk.sitkBSpline
        settings['enableCExtensions'] = True

        extractor = featureextractor.RadiomicsFeatureExtractor(**settings)

    headers = None
    featureVector = collections.OrderedDict()

    # image = sitk.GetImageFromArray(img_array)
    img_array = sitk.GetArrayFromImage(image)
    mask = sitk.GetImageFromArray(np.ones_like(img_array), isVector=True)

    mask.CopyInformation(image)

    row = []
    # try:
    featureVector.update(extractor.execute(
        image, mask, label=1))  # 提取特征ROI区域 label=1,2,3,4
    if headers is None:
        headers = list(featureVector.keys())  # [22:]

    for h in headers:
        row.append(featureVector.get(h, "N/A"))
    # except Exception:
    #     print('Error: FEATURE EXTRACTION FAILED!')

    return np.array(row[22:])


def tansig(x):
    return (2/(1+np.exp(-2*x)))-1


def to_proba(output):
    output = output-np.min(output)
    ret = (output-np.min(output))/(np.max(output)-np.min(output))
    return ret


def softmax(x):
    x = x-np.min(x)
    ex = np.exp(x)
    return ex/ex.sum()


labels = ['covid', 'nofindings', 'pneumonia']


def predict(test_img):

    params = np.load('params.npz')
    N1 = params['N1']
    N2 = params['N2']
    Beta1OfEachWindow = params['Beta1OfEachWindow']
    ymax = params['ymax']
    ymin = params['ymin']
    minOfEachWindow = params['minOfEachWindow']
    distOfMaxAndMin = params['distOfMaxAndMin']
    weightOfEnhanceLayer = params['weightOfEnhanceLayer']
    parameterOfShrink = params['parameterOfShrink']
    OutputWeight = params['OutputWeight']

    test_x = feature_extractor(test_img).reshape(1, -1)
    scaler = StandardScaler()
    test_x = scaler.fit_transform(test_x)
    print(test_x.shape)
    # test_x = np.expand_dims(test_x, axis=0)
    # print(test_x.shape)
    test_x = preprocessing.scale(test_x, axis=1)

    FeatureOfInputDataWithBiasTest = np.hstack(
        [test_x, 0.1 * np.ones((test_x.shape[0], 1))])
    OutputOfFeatureMappingLayerTest = np.zeros(
        [test_x.shape[0],  N2 * N1])

    for i in range(N2):
        outputOfEachWindowTest = np.dot(
            FeatureOfInputDataWithBiasTest,  Beta1OfEachWindow[i])
        OutputOfFeatureMappingLayerTest[:,  N1*i: N1*(i+1)] = (ymax - ymin)*(
            outputOfEachWindowTest - minOfEachWindow[i]) / distOfMaxAndMin[i] - ymin

    InputOfEnhanceLayerWithBiasTest = np.hstack(
        [OutputOfFeatureMappingLayerTest, 0.1 * np.ones((OutputOfFeatureMappingLayerTest.shape[0], 1))])
    tempOfOutputOfEnhanceLayerTest = np.dot(
        InputOfEnhanceLayerWithBiasTest,  weightOfEnhanceLayer)

    OutputOfEnhanceLayerTest = tansig(
        tempOfOutputOfEnhanceLayerTest * parameterOfShrink)

    InputOfOutputLayerTest = np.hstack(
        [OutputOfFeatureMappingLayerTest, OutputOfEnhanceLayerTest])

    OutputOfTest = np.dot(InputOfOutputLayerTest,  OutputWeight)
    # print(OutputOfTest.shape)
    OutputOfTest = np.squeeze(OutputOfTest)
    proba = OutputOfTest  # softmax(OutputOfTest)
    confidences = {labels[i]: float(proba[i]) for i in range(3)}
    # np.argmax(OutputOfTest, axis=1)  # np.array(OutputOfTest)
    return confidences


# %%
# pth = '01E392EE-69F9-4E33-BFCE-E5C968654078.jpeg'
# predict(pth)

# %%

demo = gr.Interface(fn=predict,
                    inputs=gr.Image(type="filepath"),
                    outputs=gr.Label()
                    )
demo.launch(share=False)

# %%