from radiomics import featureextractor
import SimpleITK as sitk
import os
import numpy as np
import json
from memory_profiler import profile
import gc
import tensorflow as tf
from subprocess import call, Popen
import math
from sklearn.decomposition import IncrementalPCA
from functools import reduce
import pydoc

from modules.models.data_transform import DataTransformer
from src.test_retest.test_retest_base import EvaluateEpochsBaseTF
from src.test_retest.test_retest_base import linear_trafo
from src.test_retest.test_retest_base import regularizer
from src.test_retest.test_retest_base import mnist_input_fn
from src.test_retest.mri.model_components import name_to_hidden_regularization
from src.train_hooks import BatchDumpHook, RobustnessComputationHook, \
    SumatraLoggingHook, PredictionHook, PredictionRobustnessHook
from src.train_hooks import HookFactory

from .model_components import MultiLayerPairEncoder, Conv3DEncoder
from .model_components import MultiLayerPairDecoder, Conv3DDecoder

class PyRadiomicsFeatures(DataTransformer):
    def __init__(self, streamer):
        # Initialize streamer
        _class = streamer["class"]
        self.streamer = _class(**streamer["params"])

    def get_extractor(self):
        # Initialize extractor
        extractor = featureextractor.RadiomicsFeaturesExtractor()
        extractor.enableAllImageTypes()
        extractor.enableAllFeatures()

        return extractor

    def transform(self, X, y=None):
        out_path = os.path.join(self.save_path, "features")
        os.mkdir(out_path)
        # Stream image one by one
        batches = self.streamer.get_batches()
        for batch in batches:
            for group in batch:
                for file_id in group.get_file_ids():
                    path = self.streamer.get_file_path(file_id)

                    sitk_im = sitk.ReadImage(path)
                    all_ones = np.ones(sitk_im.GetSize())
                    sitk_mask = sitk.GetImageFromArray(all_ones)

                    extractor = self.get_extractor()
                    features = extractor.computeFeatures(sitk_im, sitk_mask, "brain")

                    with open(
                        os.path.join(out_path, str(file_id) + ".json"),
                        "w"
                    ) as f:
                        json.dump(features, f, indent=2, ensure_ascii=False)

class PyRadiomicsFeaturesSpawn(DataTransformer):
    def __init__(self, streamer, out_dir, n_processes):
        # Initialize streamer
        _class = streamer["class"]
        self.streamer = _class(**streamer["params"])
        self.out_dir = out_dir
        self.n_processes = n_processes

    def get_extractor(self):
        # Initialize extractor
        extractor = featureextractor.RadiomicsFeaturesExtractor()
        extractor.enableAllImageTypes()
        extractor.enableAllFeatures()

        return extractor

    def transform(self, X, y=None):
        out_path = self.out_dir
        if not os.path.exists(self.out_dir):
            os.makedirs(self.out_dir)

        # Stream image one by one
        batches = self.streamer.get_batches()
        processes = []
        for batch in batches:
            for group in batch:
                for file_id in group.get_file_ids():
                    image_label = self.streamer.get_image_label(file_id)
                    path_in = self.streamer.get_file_path(file_id)
                    path_out = os.path.join(out_path, str(image_label) + ".json")

                    cmd = 'python -m src.test_retest.mri.run_py_radiomics_transformer '
                    cmd += "{} {}".format(path_in, path_out)
                    print(cmd)
                    # call(cmd, shell=True)
                    proc = Popen(cmd, shell=True)
                    processes.append(proc)

                    if len(processes) >= self.n_processes:
                        for p in processes:
                            p.wait()
                        processes = []

        self.streamer = None

class PyRadiomicsSingleFileTransformer(DataTransformer):
    def __init__(self, in_path, out_path):
        self.in_path = in_path
        self.out_path = out_path

    def get_extractor(self):
        # Initialize extractor
        extractor = featureextractor.RadiomicsFeaturesExtractor()
        extractor.enableAllImageTypes()
        extractor.enableAllFeatures()

        return extractor

    def transform(self, X, y=None):
        sitk_im = sitk.ReadImage(self.in_path)
        all_ones = np.ones(sitk_im.GetSize())
        sitk_mask = sitk.GetImageFromArray(all_ones)

        extractor = self.get_extractor()
        features = extractor.computeFeatures(sitk_im, sitk_mask, "brain")

        with open(
            os.path.join(self.out_path),
            "w"
        ) as f:
            json.dump(features, f, indent=2, ensure_ascii=False)

class MriIncrementalPCA(DataTransformer):
    """
    batch_size has to be larger than n_components
    """
    def __init__(self, streamer, n_components):
        _class = streamer["class"]
        self.streamer = _class(**streamer["params"])
        self.pca = IncrementalPCA(
            n_components=n_components,
            batch_size=self.streamer.batch_size
        )

    def transform(self, X, y=None):
        batches = self.streamer.get_batches()

        for batch in batches:
            file_ids = [fid for group in batch for fid in group.file_ids]
            X = []
            for fid in file_ids:
                path = self.streamer.get_file_path(fid)
                im = self.streamer.load_sample(path)
                X.append(im.ravel())

            X = np.array(X)
            print(X.shape)
            self.pca = self.pca.partial_fit(X)

        self.streamer = None

class PCAAutoEncoder(EvaluateEpochsBaseTF):
    def model_fn(self, features, labels, mode, params):
        input_dim = params["input_dim"]
        X = tf.reshape(
            features["X_0"],
            [-1, input_dim]
        )

        hidden_dim = params["hidden_dim"]
        w = tf.get_variable(
            name="weights",
            shape=[input_dim, hidden_dim],
            dtype=X.dtype,
            initializer=tf.contrib.layers.xavier_initializer(seed=43)
        )

        hidden = tf.matmul(X, w, name="hidden_rep")
        # hidden = tf.nn.sigmoid(hidden)
        w_T = tf.transpose(w)
        reconstruction = tf.matmul(
            hidden,
            w_T,
            name="reconstruction"
        )

        predictions = {
            "hidden_rep": hidden,
            "reconstruction": reconstruction
        }

        if mode == tf.estimator.ModeKeys.PREDICT:
            return tf.estimator.EstimatorSpec(
                mode=mode,
                predictions=predictions
            )

        # Compute loss
        loss = tf.losses.mean_squared_error(X, reconstruction)

        optimizer = tf.train.AdamOptimizer(
            learning_rate=params["learning_rate"]
        )
        train_op = optimizer.minimize(loss, tf.train.get_global_step())

        dump_hook_train, dump_hook_test = \
            self.get_batch_dump_hook(hidden, features["file_name_0"])

        if mode == tf.estimator.ModeKeys.TRAIN:
            return tf.estimator.EstimatorSpec(
                mode=mode,
                loss=loss,
                train_op=train_op,
                training_hooks=[dump_hook_train]
            )

        return tf.estimator.EstimatorSpec(
            mode=mode,
            loss=loss,
            evaluation_hooks=[dump_hook_test]
        )

    def gen_input_fn(self, X, y=None, train=True, input_fn_config={}):
        return self.streamer.get_input_fn(train)

class PCAAutoEncoderTuples(EvaluateEpochsBaseTF):
    def model_fn(self, features, labels, mode, params):
        encoder = MultiLayerPairEncoder(
            features=features,
            params=params,
            streamer=self.streamer
        )

        decoder = MultiLayerPairDecoder(
            features=features,
            params=params,
            encoder=encoder
        )

        hidden_0, hidden_1 = encoder.get_encodings()

        # Make sure embedinngs have the correct dimension
        dim = hidden_0.get_shape().as_list()[-1]
        assert dim == params["hidden_dim"]

        rec_0, rec_1 = decoder.get_nodes()

        predictions = {
            "input": encoder.get_nodes()[0],
            "encoding": hidden_0,
            "decoding": rec_0,
            "file_name": features["file_name_0"]
        }

        if mode == tf.estimator.ModeKeys.PREDICT:
            return tf.estimator.EstimatorSpec(
                mode=mode,
                predictions=predictions
            )

        # Compute loss
        reg_loss = decoder.get_regularization_loss()
        reconstruction_loss = decoder.get_reconstruction_loss()
        loss = decoder.get_total_loss()

        optimizer = tf.train.AdamOptimizer(
            learning_rate=params["learning_rate"]
        )
        train_op = optimizer.minimize(loss, tf.train.get_global_step())

        # Set up hooks
        train_hooks, eval_hooks = self.get_mri_ae_hooks(
            reg_loss=reg_loss,
            rec_loss=reconstruction_loss,
            enc_0=hidden_0,
            enc_1=hidden_1,
            features=features,
            params=params
        )

        if mode == tf.estimator.ModeKeys.TRAIN:
            return tf.estimator.EstimatorSpec(
                mode=mode,
                loss=loss,
                train_op=train_op,
                training_hooks=train_hooks
            )

        return tf.estimator.EstimatorSpec(
            mode=mode,
            loss=loss,
            evaluation_hooks=eval_hooks
        )

    def gen_input_fn(self, X, y=None, train=True, input_fn_config={}):
        return self.streamer.get_input_fn(train)

class MnistPCAAutoEncoder(PCAAutoEncoder):
    def gen_input_fn(self, X, y=None, train=True, input_fn_config={}):
        return mnist_input_fn(
            X,
            self.data_params,
            input_fn_config=input_fn_config
        )

class RandomImageAutoEncoder(PCAAutoEncoder):
    def gen_input_fn(self, X, y=None, train=True, input_fn_config={}):
        # images = np.random.rand(2000, 28, 28)
        images = np.zeros((2000, 28, 28))
        labels = np.ones((2000, 1))

        return tf.estimator.inputs.numpy_input_fn(
            x={"X_0": images},
            y=labels,
            **input_fn_config,
        )

# Source: https://github.com/pkmital/tensorflow_tutorials/blob/master/python/09_convolutional_autoencoder.py
class Conv2DAutoEncoder(EvaluateEpochsBaseTF):
    def model_fn(self, features, labels, mode, params):
        input_dim = params["input_dim"]
        x = features["X_0"]

        n_filters = params["n_filters"]
        filter_sizes = params["filter_sizes"]

        x = tf.reshape(x, [-1, input_dim, input_dim, 1])
        current_input = x

        # Build the encoder
        encoder = []
        shapes = []
        for layer_i, n_output in enumerate(n_filters):
            n_input = current_input.get_shape().as_list()[3]
            shapes.append(current_input.get_shape().as_list())
            W_shape = [
                filter_sizes[layer_i],
                filter_sizes[layer_i],
                n_input,
                n_output
            ]
            W = tf.get_variable(
                name="filters_layer_" + str(layer_i),
                shape=W_shape,
                initializer=tf.contrib.layers.xavier_initializer(seed=40)
            )

            b = tf.get_variable(
                name="bias_layer_" + str(layer_i),
                shape=[n_output],
                initializer=tf.initializers.zeros
            )
            encoder.append(W)
            output = tf.nn.relu(
                tf.add(
                    tf.nn.conv2d(
                        current_input, W, strides=[1, 2, 2, 1], padding='SAME'
                    ),
                    b
                )
            )
            current_input = output

        z = current_input
        encoder.reverse()
        shapes.reverse()
        # Build the decoder
        for layer_i, shape in enumerate(shapes):
            # deconv
            if not params["tied_weights"]:
                W_shape = encoder[layer_i].get_shape().as_list()
                print(shape)
                print(W_shape)
                W = tf.get_variable(
                    name="filters_deconv_layer_" + str(layer_i),
                    shape=[W_shape[0], W_shape[1], W_shape[2], W_shape[3]],
                    initializer=tf.contrib.layers.xavier_initializer(seed=40)
                )
            else:
                W = encoder[layer_i]
            b = tf.get_variable(
                name="bias_deconv_layer_" + str(layer_i),
                shape=W.get_shape().as_list()[2],
                initializer=tf.initializers.zeros
            )

            input_shapes = current_input.get_shape().as_list()
            print("input shape {}".format(input_shapes))
            print("W {}".format(W))
            output = tf.nn.relu(
                tf.add(
                    tf.nn.conv2d_transpose(
                        current_input, W,
                        tf.stack([tf.shape(x)[0], shape[1],
                                  shape[2], shape[3]]),
                        strides=[1, 2, 2, 1], padding='SAME'
                    ),
                    b
                )
            )
            print("out shape {}".format(output.get_shape().as_list()))
            current_input = output

        y = current_input
        print("y {}".format(y))
        predictions = {
            "hidden_rep": z,
            "reconstruction": y
        }

        if mode == tf.estimator.ModeKeys.PREDICT:
            return tf.estimator.EstimatorSpec(
                mode=mode,
                predictions=predictions
            )

        loss = tf.losses.mean_squared_error(x, y)

        optimizer = tf.train.AdamOptimizer(
            learning_rate=params["learning_rate"]
        )
        train_op = optimizer.minimize(loss, tf.train.get_global_step())

        if mode == tf.estimator.ModeKeys.TRAIN:
            return tf.estimator.EstimatorSpec(
                mode=mode,
                loss=loss,
                train_op=train_op
            )

        return tf.estimator.EstimatorSpec(
            mode=mode,
            loss=loss
        )

class Conv3DAutoEncoder(EvaluateEpochsBaseTF):
    def model_fn(self, features, labels, mode, params):
        encoder = Conv3DEncoder(features, params, self.streamer)
        decoder = Conv3DDecoder(features, params, encoder)

        z = encoder.get_encoding()
        y = decoder.get_reconstruction()
        x = encoder.get_reconstruction_target()

        flattened_z = tf.contrib.layers.flatten(z)
        dim = flattened_z.get_shape().as_list()[-1]
        assert dim == params["encoding_dim"]

        predictions = {
            "hidden_rep": flattened_z,
            "reconstruction": y
        }

        if mode == tf.estimator.ModeKeys.PREDICT:
            return tf.estimator.EstimatorSpec(
                mode=mode,
                predictions=predictions
            )

        loss = tf.losses.mean_squared_error(x, y)

        optimizer = tf.train.AdamOptimizer(
            learning_rate=params["learning_rate"]
        )
        train_op = optimizer.minimize(loss, tf.train.get_global_step())

        y_hook_train, y_hook_test = \
            self.get_batch_dump_hook(z, features["file_name_0"])

        if mode == tf.estimator.ModeKeys.TRAIN:
            return tf.estimator.EstimatorSpec(
                mode=mode,
                loss=loss,
                train_op=train_op,
                training_hooks=[y_hook_train]
            )

        return tf.estimator.EstimatorSpec(
            mode=mode,
            loss=loss,
            evaluation_hooks=[y_hook_test]
        )

    def gen_input_fn(self, X, y=None, train=True, input_fn_config={}):
        return self.streamer.get_input_fn(train)

class Conv3DTupleAE(EvaluateEpochsBaseTF):
    def model_fn(self, features, labels, mode, params):
        encoder_class = params["encoder_class"]
        decoder_class = params["decoder_class"]

        with tf.variable_scope("conv_3d_encoder", reuse=tf.AUTO_REUSE):
            encoder_0 = encoder_class(
                input_key="X_0",
                features=features,
                params=params,
                streamer=self.streamer
            )
        with tf.variable_scope("conv_3d_encoder", reuse=tf.AUTO_REUSE):
            encoder_1 = encoder_class(
                input_key="X_1",
                features=features,
                params=params,
                streamer=self.streamer
            )

        target_0 = encoder_0.get_reconstruction_target()
        target_1 = encoder_1.get_reconstruction_target()
        if params["asymmetric"]:
            target_0 = encoder_1.get_reconstruction_target()
            target_1 = encoder_0.get_reconstruction_target()

        with tf.variable_scope("conv_3d_decoder", reuse=tf.AUTO_REUSE):
            decoder_0 = decoder_class(features, params, encoder_0, target_0)

        with tf.variable_scope("conv_3d_decoder", reuse=tf.AUTO_REUSE):
            decoder_1 = decoder_class(features, params, encoder_1, target_1)

        z_0 = encoder_0.get_encoding()
        y_0 = decoder_0.get_reconstruction()
        x_0 = encoder_0.get_reconstruction_target()

        z_1 = encoder_1.get_encoding()
        y_1 = decoder_1.get_reconstruction()
        x_1 = encoder_1.get_reconstruction_target()

        if params["asymmetric"]:
            tmp = y_0
            y_0 = y_1
            y_1 = tmp

        flattened_z_0 = tf.contrib.layers.flatten(z_0)
        flattened_z_1 = tf.contrib.layers.flatten(z_1)
        dim = flattened_z_0.get_shape().as_list()[-1]

        assert dim == encoder_0.get_encoding_dim()

        predictions = {
            "input": x_0,
            "encoding": flattened_z_0,
            "decoding": y_0,
            "file_name": features["file_name_0"]
        }

        if mode == tf.estimator.ModeKeys.PREDICT:
            return tf.estimator.EstimatorSpec(
                mode=mode,
                predictions=predictions
            )

        rec_loss_0 = decoder_0.get_reconstruction_loss()
        rec_loss_1 = decoder_1.get_reconstruction_loss()
        rec_loss = rec_loss_0 / 2 + rec_loss_1 / 2

        # Regularization
        reg_loss = tf.constant(0, dtype=rec_loss.dtype)
        to_reg_0 = z_0
        to_reg_1 = z_1

        diagnose_dim = params["diagnose_dim"]
        hidden_dim = encoder_0.get_encoding_dim()
        if diagnose_dim > 0:
            patient_dim = hidden_dim - diagnose_dim
            patient_encs_0, diag_encs_0 = tf.split(
                z_0,
                [patient_dim, diagnose_dim],
                axis=1
            )

            patient_encs_1, diag_encs_1 = tf.split(
                z_1,
                [patient_dim, diagnose_dim],
                axis=1
            )

            to_reg_0 = diag_encs_0
            to_reg_1 = diag_encs_1

        reg_lambda = params["hidden_lambda"]
        if reg_lambda != 0:
            reg_name = params["hidden_regularizer"]
            reg_loss = name_to_hidden_regularization(
                0,
                reg_name,
                to_reg_0,
                to_reg_1
            )
            reg_loss *= reg_lambda

        loss = rec_loss + reg_loss

        optimizer = tf.train.AdamOptimizer(
            learning_rate=params["learning_rate"]
        )

        train_op = optimizer.minimize(loss, tf.train.get_global_step())

        train_hooks, eval_hooks = self.get_mri_ae_hooks(
            reg_loss=reg_loss,
            rec_loss=rec_loss,
            enc_0=flattened_z_0,
            enc_1=flattened_z_1,
            features=features,
            params=params
        )

        if mode == tf.estimator.ModeKeys.TRAIN:
            return tf.estimator.EstimatorSpec(
                mode=mode,
                loss=loss,
                train_op=train_op,
                training_hooks=train_hooks
            )

        return tf.estimator.EstimatorSpec(
            mode=mode,
            loss=loss,
            evaluation_hooks=eval_hooks
        )

    def gen_input_fn(self, X, y=None, train=True, input_fn_config={}):
        return self.streamer.get_input_fn(train)

class MnistConv2DAutoEncoder(Conv2DAutoEncoder):
    def gen_input_fn(self, X, y=None, train=True, input_fn_config={}):
        return mnist_input_fn(
            X,
            self.data_params,
            input_fn_config=input_fn_config
        )

class MnistConv3DAutoEncoder(Conv3DAutoEncoder):
    def gen_input_fn(self, X, y=None, train=True, input_fn_config={}):
        return mnist_input_fn(
            X,
            self.data_params,
            input_fn_config=input_fn_config
        )


import kivy
from kivy.app import App
from kivy.uix.button import Button
from kivy.uix.boxlayout import BoxLayout
import random
import Encounter

class RootDisplay(BoxLayout):

    def __init__(self):
        super(RootDisplay, self).__init__()
        self.listOfPlayers = []

    def add_player(self):
        level = self.player_level.text
        self.listOfPlayers.append(level)
        self.playerList.text = str(self.listOfPlayers)

    def remove_player(self):
        if not len(self.listOfPlayers) == 0:
            self.listOfPlayers.pop()
            self.playerList.text = str(self.listOfPlayers)

    def update_num_monsters(self):
        self.numMonsters.text = "Number of Monsters: " + str(self.slider1.value)

    def update_difficulty(self):
        diff = int(self.slider2.value)
        output = ''
        if diff == 0:
            output = "Easy"
        elif diff == 1:
            output = "Medium"
        elif diff == 2:
            output = "Hard"
        else:
            output = "Deadly"

        self.difficulty.text = "Difficulty: " + output

    def call_encounter_builder(self):
            myEncounter = Encounter.Encounter()
            # levels = ["8", "8", "10", "6"]
            difficulty = 3   # 0-3 Easy, Medium, Hard, Deadly
            size = self.slider1.value          # Exact number of desired monsters
            onlyOfType = []
            output = myEncounter.builder(self.listOfPlayers, int(self.slider2.value), size, onlyOfType)
            self.random_label.text = str(output)

class EncounterBuilder(App):

    def build(self):
        return RootDisplay()

myEncounter = EncounterBuilder()
myEncounter.run()


# 1번 > 이미지 편집
# `crop` 메소드는 잘라낼 영역의 좌측 상단 좌표$(x_1, y_1)$과 우측 하단 좌표$(x_2, y_2)$값 4개를 튜플로 요구합니다. 즉, `crop` 메소드는 아래와 같은 형태로 사용해야 합니다.
# img.crop((x1, y1, x2, y2))

#  2번 > CNN 행렬곱
# CNN 행렬곱 -> 그자리 그대로 곱하기
# https://untitledtblog.tistory.com/150 참고

# 3번 > RNN 설명

# 4번문제 > CNN 모델로 CIFAR-10 분류하기

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

from elice_utils import EliceUtils

elice_utils = EliceUtils()

import tensorflow as tf
from tensorflow.keras import layers, Sequential, Input
from tensorflow.keras.optimizers import Adam, SGD
from tensorflow.keras.utils import to_categorical

import numpy as np
import matplotlib.pyplot as plt

def load_cifar10():
    # CIFAR-10 데이터셋을 불러옵니다.
    X_train = np.load("cifar10_train_X.npy")
    y_train = np.load("cifar10_train_y.npy")
    X_test = np.load("cifar10_test_X.npy")
    y_test = np.load("cifar10_test_y.npy")

    # TODO: [지시사항 1번] 이미지의 각 픽셀값을 0에서 1 사이로 정규화하세요.
    X_train = X_train/255.0
    X_test = X_test/255.0

    # 정수 형태로 이루어진 라벨 데이터를 one-hot encoding으로 바꿉니다.
    y_train = to_categorical(y_train, num_classes=10)
    y_test = to_categorical(y_test, num_classes=10)

    return X_train, X_test, y_train, y_test

def build_cnn_model(num_classes, input_shape):
    model = Sequential()

    # TODO: [지시사항 2번] 지시사항 대로 CNN 모델을 만드세요.
    model.add(layers.Conv2D(16, kernel_size=(3,3), padding="same", activation="relu", input_shape=input_shape))
    model.add(layers.Conv2D(32, kernel_size=(3,3), padding="same", activation="relu"))
    model.add(layers.MaxPool2D(2))
    model.add(layers.Conv2D(64, kernel_size=(3,3), padding="same", activation="relu"))
    model.add(layers.Conv2D(128, kernel_size=(3,3), padding="same", activation="relu"))
    model.add(layers.MaxPool2D(2))
    model.add(layers.Conv2D(256, kernel_size=(3,3), padding="same", activation="relu"))
    model.add(layers.Conv2D(512, kernel_size=(3,3), padding="same", activation="relu"))
    model.add(layers.MaxPool2D(2))
    model.add(layers.GlobalAveragePooling2D())
    model.add(layers.Dense(100, activation="relu"))
    model.add(layers.Dense(num_classes, activation="softmax"))

    return model

def plot_loss(hist):
    # hist 객체에서 train loss와 valid loss를 불러옵니다.
    train_loss = hist.history["loss"]
    val_loss = hist.history["val_loss"]
    epochs = np.arange(1, len(train_loss) + 1)

    fig, ax = plt.subplots(figsize=(12, 8))
    ax.set_xticks(list(epochs))

    # ax를 이용하여 train loss와 valid loss를 plot 합니다..
    ax.plot(epochs, train_loss, marker=".", c="blue", label="Train Loss")
    ax.plot(epochs, val_loss, marker=".", c="red", label="Valid Loss")

    ax.legend(loc="upper right")
    ax.grid()
    ax.set_xlabel("Epoch")
    ax.set_ylabel("Loss")

    fig.savefig("loss.png")

def plot_accuracy(hist):
    # hist 객체에서 train accuracy와 valid accuracy를 불러옵니다..
    train_acc = hist.history["accuracy"]
    val_acc = hist.history["val_accuracy"]
    epochs = np.arange(1, len(train_acc) + 1)

    fig, ax = plt.subplots(figsize=(12, 8))
    ax.set_xticks(list(epochs))
    # ax를 이용하여 train accuracy와와 valid accuracy와를 plot 합니다.
    ax.plot(epochs, val_acc, marker=".", c="red", label="Valid Accuracy")
    ax.plot(epochs, train_acc, marker=".", c="blue", label="Train Accuracy")

    ax.legend(loc="lower right")
    ax.grid()
    ax.set_xlabel("Epoch")
    ax.set_ylabel("Accuracy")

    fig.savefig("accuracy.png")

def get_topk_accuracy(y_test, y_pred, k=1):
    # one-hot encoding으로 이루어진(y_test를 다시 정수 라벨 형식으로 바꿉니다.
    true_labels = np.argmax(y_test, axis=1)

    # y_pred를 확률값이 작은 것에서 큰 순서로 정렬합니다.
    pred_labels = np.argsort(y_pred, axis=1)

    correct = 0
    for true_label, pred_label in zip(true_labels, pred_labels):
        # TODO: [지시사항 3번] 현재 pred_label에서 확률값이 가장 큰 라벨 k개를 가져오세요
        cur_preds = pred_label[-k:]

        if true_label in cur_preds:
            correct += 1

    # TODO: [지시사항 3번] Top-k accuarcy를 구하세요.
    topk_accuracy = correct / len(true_labels)

    return topk_accuracy

def main(model=None, epochs=5):
    # 시드 고정을 위한 코드입니다. 수정하지 마세요!
    tf.random.set_seed(2022)

    X_train, X_test, y_train, y_test = load_cifar10()
    cnn_model = build_cnn_model(len(y_train[0]), X_train[0].shape)
    cnn_model.summary()

    # TODO: [지시사항 4번] 지시사항 대로 모델의 optimizer, loss, metrics을 설정하세요.
    optimizer = SGD(learning_rate=0.01, momentum=0.9, nesterov=True)
    cnn_model.compile(optimizer=optimizer, loss="categorical_crossentropy", metrics=["accuracy"])

    # TODO: [지시사항 5번] 지시사항 대로 hyperparameter를 설정하여 모델을 학습하세요.
    hist = cnn_model.fit(X_train, y_train, epochs=epochs, batch_size=32, validation_split= 0.2, shuffle= True, verbose=2)

    # Test 데이터를 적용했을 때 예측 확률을 구합니다.
    y_pred = cnn_model.predict(X_test)
    top1_accuracy = get_topk_accuracy(y_test, y_pred)
    top3_accuracy = get_topk_accuracy(y_test, y_pred, k=3)

    print("Top-1 Accuracy: {:.3f}%".format(top1_accuracy * 100))
    print("Top-3 Accuracy: {:.3f}%".format(top3_accuracy * 100))

    # Test accuracy를 구합니다.
    _, test_accuracy = cnn_model.evaluate(X_test, y_test, verbose=0)

    # Tensorflow로 구한 test accuracy와 top1 accuracy는 같아야 합니다.
    # 다만 부동 소수점 처리 문제로 완전히 같은 값이 나오지 않는 경우도 있어서
    # 소수점 셋째 자리까지 반올림하여 비교합니다.
    assert round(test_accuracy, 3) == round(top1_accuracy, 3)

    plot_loss(hist)
    plot_accuracy(hist)

    return optimizer, hist

if __name__ == '__main__':
    main()

# 5번문제 >


import logging
import time

import redis

class RedisConnector:
    def __init__(self, host, port):
        self.host = host
        self.port = port
        self.logger = logging.getLogger("eventgen_server")
        self.members_connection = redis.Redis(host=self.host, port=int(self.port), db=0)
        self.message_connection = redis.Redis(host=self.host, port=int(self.port), db=1)
        self.pubsub = self.message_connection.pubsub()
        self.logger.info("Initialized RedisConnector")
        self.servers_channel = "servers_channel"
        self.controller_channel = "controller_channel"
        self.retry_time_list = [5, 10, 20, 30, 60, 0]

    def register_myself(self, hostname, role="server"):
        for retry_time in self.retry_time_list:
            try:
                if role == "server":
                    self.members_connection.sadd("servers", hostname)
                    self.pubsub.subscribe(self.servers_channel)
                else:
                    self.members_connection.set("controller", hostname)
                    self.pubsub.subscribe(self.controller_channel)
                self.logger.info(
                    "Registered as {} and subscribed the channel.".format(role)
                )
                return
            except:
                self.logger.warning(
                    "Could not connect to Redis at {}:{}. Retrying in {} seconds.".format(
                        self.host, self.port, retry_time
                    )
                )
                if not retry_time:
                    raise Exception("Failed to connect to Redis.")
                time.sleep(retry_time)
                continue

    def get_registered_servers(self):
        for retry_time in self.retry_time_list:
            try:
                servers = list(self.members_connection.smembers("servers"))
                self.logger.info("Registered Servers: {}".format(servers))
                return servers
            except:
                self.logger.warning(
                    "Could not connect to Redis at {}:{}. Retrying in {} seconds.".format(
                        self.host, self.port, retry_time
                    )
                )
                if not retry_time:
                    raise Exception("Failed to connect to Redis.")
                time.sleep(retry_time)
                continue


from __future__ import division
__author__ = 'Calle Svensson <calle.svensson@zeta-two.com>'
from builtins import bytes, str, filter, map, range, zip
from future.standard_library import install_aliases
from functools import reduce
from six import byte2int, int2byte
import operator
install_aliases()

import string, operator, itertools, scipy.stats
from scipy.stats.distributions import chi2
from collections import Counter

from . import utility, mathtools, INF, BYTE_MAX
from zetacrypt.ciphers import xor_seq_byte
from zetacrypt.conversions import *

# The relative frequency of alphabet letters in the English language
# FREQ_ENGLISH = {'e': 0.12575645, 't': 0.9085226, 'a': 0.8000395, 'o': 0.7591270, 'i': 0.6920007, 'n': 0.6903785,
#                's': 0.6340880, 'h': 0.6236609, 'r': 0.5959034, 'd': 0.4317924, 'l': 0.4057231, 'u': 0.2841783,
#                'c': 0.2575785, 'm': 0.2560994, 'f': 0.2350463, 'w': 0.2224893, 'g': 0.1982677, 'y': 0.1900888,
#                'p': 0.1795742, 'b': 0.1535701, 'v': 0.0981717, 'k': 0.0739906, 'x': 0.0179556, 'j': 0.0145188,
#                'q': 0.0117571, 'z': 0.0079130}
FREQ_ENGLISH = {'e': 0.12702, 't': 0.09056, 'a': 0.08167, 'o': 0.07507, 'i': 0.06966, 'n': 0.06749, 's': 0.06327,
                'h': 0.06094, 'r': 0.05987, 'd': 0.04253, 'l': 0.04025, 'c': 0.02782, 'u': 0.02758, 'm': 0.02406,
                'w': 0.02361, 'f': 0.02228, 'g': 0.02015, 'y': 0.01974, 'p': 0.01929, 'b': 0.01492, 'v': 0.00978,
                'k': 0.00772, 'j': 0.00153, 'x': 0.00150, 'q': 0.00095, 'z': 0.00074}
IC_ENGLISH = 1.73

def get_expected_freq(message_len, freq):
    return dict((k, v * message_len) for k, v in freq.items())

def count_printable(seq):
    """Returns the number of printable and the total number of ASCII characters in seq as a tuple"""
    printable_list = map(lambda c: 32 <= c < 127 or c == ord('\n'), seq)
    printable_total_list = map(lambda c: (1 if c else 0, 1), printable_list)
    return tuple(reduce(lambda x, y: (x[0]+y[0], x[1]+y[1]), printable_total_list, (0,0)))

def is_printable(seq):
    """Returns true is seq consists solely of printable ASCII characters"""
    printable, total = count_printable(seq)
    return printable == total

def letter_frequency(seq):
    """Returns a dictionary with the frequencies of letters in the sequence"""
    freq = filter(lambda x: x in string.ascii_letters, seq.lower())
    freq = dict(Counter(freq).most_common())
    freq.update(dict((x, 0) for x in filter(lambda x: x not in freq, string.ascii_lowercase)))

    return freq

def letter_frequency_rel(seq):
    """Returns a dictionary with the relative frequency of letters in the sequence"""
    freq = letter_frequency(seq)
    total = len(seq)
    if total == 0:
        return freq
    else:
        return {k: v / total for k, v in freq.items()}

def index_coincidence(seq):
    """Returns the index of coincidence for a sequence"""
    freq = letter_frequency(seq)
    return len(string.ascii_lowercase) * sum(map(lambda x: x * (x - 1), freq.values())) / (len(seq) * (len(seq) - 1))

def chi_square_letter_freq(freq, expected):
    obs = list(map(operator.itemgetter(1), sorted(freq.items(), key=operator.itemgetter(0))))
    expected = list(map(operator.itemgetter(1), sorted(expected.items(), key=operator.itemgetter(0))))
    chival, _ = scipy.stats.chisquare(obs, expected)
    return chival

#@profile
def find_single_byte_xor_key(seq):
    """Find the most probable single byte XOR key used to encrypt seq"""
    best_dist = INF
    best_key = 0
    best = "FAIL"
    for key in range(BYTE_MAX):
        key = int2byte(key)
        xor = xor_seq_byte(seq, key)
        #mb = iterator_to_bytes(xor)
        xor1, xor2 = itertools.tee(xor)
        xor2, a = itertools.tee(xor2)
        if not is_printable(xor1):
            continue
        mb = iterator_to_bytes(xor2)

        m = bytes_to_ascii(mb)

        freq = letter_frequency(m)
        dist = chi_square_letter_freq(freq, get_expected_freq(len(m), FREQ_ENGLISH))

        # If better, save
        if dist < best_dist:
            best = m
            best_key = key
            best_dist = dist

    return best, best_key, best_dist

def find_vigenere_key_len(cipher, mink, maxk):
    """Returns the most probable keylength in a vigenere based cipher"""
    best_dist = INF
    best_keysize = 0
    for keysize in range(mink, maxk):
        # Average hamming weight over BLOCKS
        dist = []
        blocks = utility.chunks(cipher, keysize)

        # Moving pairwise distance
        prev_block = next(blocks)
        for block in blocks:
            dist.append(float(mathtools.hamming_distance_bit(prev_block, block)) / keysize)
            prev_block = block
        dist = sum(dist) / len(dist)

        # If better, save
        if dist < best_dist:
            best_dist = dist
            best_keysize = keysize
    return best_keysize

def find_vigenere_key(cipher, keylen):
    """Breaks a vigenere cipher with known keylen"""
    key = []
    for block in utility.transpose(cipher, keylen):
        _, k, _ = find_single_byte_xor_key(block)
        key.append(k)
    return bytes(key)

def count_repeated_blocks(ciphertext, block_size):
    blocks = Counter(utility.chunks(ciphertext, block_size))
    return blocks.most_common(1)[0][1]

def detect_ecb(ciphertext, block_size):
    return count_repeated_blocks(ciphertext, block_size) > 1

def encryption_detection_oracle_ecb_cbc(oracle, blocklen, answer=False):
    plaintext = ascii_to_bytes("A" * (16 * 3))
    if answer:  # If black box supports it, leak real answer
        c, ans = oracle(plaintext, True)
        return detect_ecb(c, blocklen), ans
    else:  # Otherwise, just return guess
        c = oracle(plaintext)
        return detect_ecb(c, blocklen)

def find_ecb_block_length(blackbox):
    """Finds out the block length of a ECB encryption function."""
    # Find start of new block
    secret_len = len(blackbox(""))
    for block_start_cand in range(256):
        newlen = len(blackbox(b"A" * block_start_cand))
        if newlen > secret_len:
            baseline = newlen
            block_start = block_start_cand
            break

    # Find block len
    for block_len_cand in range(256):
        newlen = len(blackbox(b"A" * (block_len_cand + block_start)))
        if newlen > baseline:
            return block_len_cand

def decrypt_ecb_postfix(blackbox, block_size):
    """Decrypts the postfix part of an ECB like encryption function"""

    # Number of blocks to decrypt
    secret_len = len(blackbox(b""))
    secret_blocks = (secret_len + block_size - 1) // block_size

    # For each block
    message = bytes()
    for block in range(secret_blocks):
        # For each element in the block
        for element in range(block_size):
            # Calculate target block
            block_base = (block_size - element - 1) * b"A"
            block_cipher_base = blackbox(block_base)[:(block + 1) * block_size]

            # Try all bytes to get matching block
            for b in range(BYTE_MAX):
                guess_byte = bytes([b])
                block_guess = block_base + message + guess_byte
                block_cipher_guess = blackbox(block_guess)[:(block + 1) * block_size]

                # Append to message and move to next element
                if block_cipher_base == block_cipher_guess:
                    message += guess_byte
                    break
    return message


'''
This main file uses imported functions and classes to create different room and
evaluate them with the base fall model. The base fall model does not have the
trajectpries implemented yet
'''
#!/usr/bin/env python
import numpy as np
from matplotlib import pyplot as plt
#from functions_collision_detection import plot_lights, plot_polygon # for plotting the rooms
from class_CMAESRoom import CMAES_Room

#
#############################################################################
#                                 MAIN
#############################################################################

if __name__ == "__main__": # Run when it is called, not when it is imported

    ''' ______________________ ROOM GENERATOR ___________________________'''

    ''' WHICH FURNITURE, DOOR AND LIGHTS ARE USED '''
    mainFurnitureCodeArray = [8, 14, 16]
#    mainFurnitureCodeArray = [23,24, 28, 30, 16]
    mainDoorsCodeArray = [10] #*** MAIN DOORS ONLY TAKES IN THE DOORS THAT ARE NOT THE DOOR BETWEEN MAIN ROOM AND BATHROOM ****
#    mainDoorsCodeArray = []
    mainLightCodeArray = [2, 30]
#    mainLightCodeArray = []

    bathFurnitureCodeArray = [5,9]
#    bathFurnitureCodeArray = []
    bathDoorsCodeArray = [11] # *** BATHROOMDOORS ARE ONLY THE DOOR THAT IS BETWEEN THE MAIN ROOM AND THE BATHROOM
#    bathDoorsCodeArray = []
    bathLightCodeArray = [2]
    bathLightCodeArray = []

    mainRoom = [mainFurnitureCodeArray, mainDoorsCodeArray, mainLightCodeArray, 'main_room']
    bathRoom = [bathFurnitureCodeArray, bathDoorsCodeArray, bathLightCodeArray, 'bath_room']

    ''' ____________________________ FALL MODEL ___________________________'''

# intention_set = {'start': ['Bed', 'Main Door', 'Bed', 'Toilet', 'Sink-Bath', 'Chair-Patient', 'Toilet', 'Bed', 'Chair-Patient', 'Bed', 'Chair-Visitor'], 'end': ['Main Door', 'Bed', 'Toilet', 'Sink-Bath', 'Bed', 'Toilet', 'Chair-Patient', 'Chair-Patient', 'Bed', 'Chair-Visitor', 'Bed'], 'frequency': [3, 3, 6, 6, 6, 3, 3, 2, 2, 1, 1]} # for P22 & Room-4 & Room-2 designs
#    intention_set = {'start': ['Bed', 'Main Door', 'Bed', 'Toilet', 'Sink-Bath', 'Chair-Patient', 'Toilet', 'Bed', 'Chair-Patient', 'Bed', 'Sofa'], 'end': ['Main Door', 'Bed', 'Toilet', 'Sink-Bath', 'Bed', 'Toilet', 'Chair-Patient', 'Chair-Patient', 'Bed', 'Sofa', 'Bed'], 'frequency': [3, 3, 6, 6, 6, 3, 3, 2, 2, 1, 1]} # for A-K & S-B & J-M & J-C & J-G & B-L & B-JH & Room-1 & Room-3 designs
    intention_set = {'start': ['Bed', 'Bed'], 'end': ['Sofa', 'Toilet'], 'frequency': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]} # for test
#    intention_set = {'start': ['Sofa', 'Sofa'], 'end': ['Cabinet', 'Chair-Patient'], 'frequency': [2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1]} # for test

    ''' _______________________   OPTIMIZATION __________________________'''
    ''' The size of elite population and consequently the size of the possible solution
    should be relatively large to avoid small covariance matrix. With small covariance matrices
    there is a chance the program gets stuck in an infinite loop of object placement without collision
    '''
    nBest =  5  # number of elite solutions. Size of the elite population
    numPossibleSolution = 10 # Number of possible solutions, size of the candid solutions
    numOfIteration = 70 # number of iteration for optimization
    optimization = CMAES_Room(nBest, numPossibleSolution, numOfIteration, mainRoom, bathRoom, intention_set)

#    plt.plot(optimization.m0x); plt.ylabel('object 0, x-axis'); plt.xlabel('# of iterations'); plt.show()
#    plt.plot(optimization.m1x); plt.ylabel('object 1, x-axis'); plt.xlabel('# of iterations'); plt.show()
#
#    plt.plot(optimization.m0y); plt.ylabel('object 0, y-axis'); plt.xlabel('# of iterations'); plt.show()
#    plt.plot(optimization.m1y); plt.ylabel('object 1, y-axis'); plt.xlabel('# of iterations'); plt.show()
#
#    plt.plot(optimization.m0theta); plt.ylabel('object 0, orientation'); plt.xlabel('# of iterations'); plt.show()
#    plt.plot(optimization.m1theta); plt.ylabel('object 1, orientation'); plt.xlabel('# of iterations'); plt.show()
#
    summed_list = [optimization.bestSolutionMean[i] + optimization.bestSolutionStd[i] for i in range(len(optimization.bestSolutionMean))]
    subtracted_list = [optimization.bestSolutionMean[i] - optimization.bestSolutionStd[i] for i in range(len(optimization.bestSolutionMean))]
    plt.plot(np.linspace(0,numOfIteration-1, numOfIteration,numOfIteration), optimization.bestSolutionMean, color = 'crimson'); plt.fill_between(np.linspace(0,numOfIteration-1, numOfIteration,numOfIteration), subtracted_list, summed_list, facecolor = 'coral', alpha = 0.5)
#    plt.ylabel('Best Solution Value'); plt.xlabel('# of iterations'); plt.show()

    plt.plot(optimization.bestSol)
    plt.ylabel('Best Solution Value'); plt.xlabel('# of iterations'); plt.show()


import pathlib
import tempfile

from async_service import background_trio_service
import pytest
import trio

from alexandria.abc import ContentBundle
from alexandria.rpc import RPCServer
from alexandria.payloads import Advertise, FindNodes, Locate, Ping, Retrieve
from alexandria.tools.factories import ApplicationFactory
from alexandria.w3 import get_w3

@pytest.fixture
def ipc_path():
    with tempfile.TemporaryDirectory() as temp_dir:
        yield pathlib.Path(temp_dir) / 'jsonrpc.ipc'

@pytest.fixture
async def peers(ipc_path):
    bootnodes = []
    peer_a = ApplicationFactory()
    async with background_trio_service(peer_a):
        bootnodes.append(peer_a.client.local_node)
        peer_b = ApplicationFactory(bootnodes=bootnodes)
        async with background_trio_service(peer_b):
            bootnodes.append(peer_b.client.local_node)
            peer_c = ApplicationFactory()
            async with background_trio_service(peer_c):
                bootnodes.append(peer_c.client.local_node)
                peer_d = ApplicationFactory(bootnodes=bootnodes)
                async with background_trio_service(peer_d):
                    bootnodes.append(peer_d.client.local_node)
                    peer_e = ApplicationFactory(bootnodes=bootnodes)
                    async with background_trio_service(peer_e):
                        yield (peer_a, peer_b, peer_c, peer_d, peer_e)

@pytest.fixture
async def rpc_node(ipc_path, peers):
    bootnodes = tuple(
        peer.client.local_node
        for peer in peers
    )
    alice = ApplicationFactory(bootnodes=bootnodes)
    async with background_trio_service(alice):
        json_rpc_server = RPCServer(
            ipc_path=ipc_path,
            client=alice.client,
            network=alice.network,
            kademlia=alice.kademlia,
            routing_table=alice.routing_table,
        )
        async with background_trio_service(json_rpc_server):
            await json_rpc_server.wait_serving()
            for _ in range(32):
                await trio.hazmat.checkpoint()
            yield alice

@pytest.fixture
async def w3(ipc_path, rpc_node):
    await trio.sleep(1)
    return get_w3(ipc_path=ipc_path)

@pytest.mark.trio
async def test_rpc_get_local_node(w3, rpc_node):
    node = await trio.to_thread.run_sync(w3.alexandria.get_local_node_id)
    assert node == rpc_node.client.local_node

@pytest.mark.trio
async def test_rpc_get_routing_table_stats(w3, rpc_node):
    stats = await trio.to_thread.run_sync(w3.alexandria.get_routing_table_stats)
    assert stats == rpc_node.routing_table.get_stats()

@pytest.mark.trio
async def test_rpc_get_routing_table_bucket_info(w3, rpc_node):
    info = await trio.to_thread.run_sync(w3.alexandria.get_routing_table_bucket_info, 255)
    assert info == rpc_node.routing_table.get_bucket_info(255)
    assert all(isinstance(value, int) for value in info.nodes)
    assert all(isinstance(value, int) for value in info.replacement_cache)

@pytest.mark.trio
async def test_rpc_ping(w3, rpc_node, peers):
    bob = peers[0]
    async with bob.client.message_dispatcher.subscribe(Ping) as subscription:
        elapsed = await trio.to_thread.run_sync(w3.alexandria.ping, bob.client.local_node)
        with trio.fail_after(1):
            await subscription.receive()
        assert elapsed > 0

@pytest.mark.trio
async def test_rpc_find_nodes(w3, rpc_node, peers):
    bob = peers[0]
    async with bob.client.message_dispatcher.subscribe(FindNodes) as subscription:
        found_nodes = await trio.to_thread.run_sync(
            w3.alexandria.find_nodes,
            bob.client.local_node,
            256,
        )
        with trio.fail_after(1):
            await subscription.receive()
        assert len(found_nodes) == len(bob.routing_table.buckets[255])

@pytest.mark.trio
async def test_rpc_advertise(w3, rpc_node, peers):
    bob = peers[0]

    async with bob.client.message_dispatcher.subscribe(Advertise) as subscription:
        elapsed = await trio.to_thread.run_sync(
            w3.alexandria.advertise,
            bob.client.local_node,
            b'key',
            rpc_node.client.local_node,
        )
        with trio.fail_after(1):
            await subscription.receive()
        assert elapsed > 0

@pytest.mark.trio
async def test_rpc_locate(w3, rpc_node, peers):
    bob = peers[0]
    bob.kademlia.content_manager.ingest_content(ContentBundle(
        b'key-a',
        b'value-a',
        bob.client.local_node_id,
    ))
    async with bob.client.message_dispatcher.subscribe(Locate) as subscription:
        found_nodes = await trio.to_thread.run_sync(
            w3.alexandria.locate,
            bob.client.local_node,
            b'key-a',
        )
        with trio.fail_after(1):
            await subscription.receive()
        assert len(found_nodes) >= 1
        node = found_nodes[0]
        assert node.node_id == bob.client.local_node_id

@pytest.mark.trio
async def test_rpc_retrieve(w3, rpc_node, peers):
    bob = peers[0]
    bob.kademlia.content_manager.ingest_content(ContentBundle(
        b'key-a',
        b'value-a',
        bob.client.local_node_id,
    ))
    async with bob.client.message_dispatcher.subscribe(Retrieve) as subscription:
        data = await trio.to_thread.run_sync(
            w3.alexandria.retrieve,
            bob.client.local_node,
            b'key-a',
        )
        with trio.fail_after(1):
            await subscription.receive()
        assert data == b'value-a'

@pytest.mark.trio
async def test_rpc_contentStats(w3, rpc_node):
    stats = await trio.to_thread.run_sync(w3.alexandria.get_content_stats)
    assert stats == rpc_node.kademlia.content_manager.get_stats()

@pytest.mark.trio
async def test_rpc_addContent_durable(w3, rpc_node):
    assert b'key' not in set(rpc_node.kademlia.content_manager.iter_content_keys())
    await trio.to_thread.run_sync(w3.alexandria.add_content, b'key', b'value', False)
    assert b'key' in set(rpc_node.kademlia.content_manager.iter_content_keys())
    assert b'key' in set(rpc_node.kademlia.content_manager.durable_db.keys())
    assert b'key' not in set(rpc_node.kademlia.content_manager.ephemeral_db.keys())
    assert b'key' not in set(rpc_node.kademlia.content_manager.cache_db.keys())

@pytest.mark.trio
async def test_rpc_addContent_ephemeral(w3, rpc_node):
    assert b'key' not in set(rpc_node.kademlia.content_manager.iter_content_keys())
    await trio.to_thread.run_sync(w3.alexandria.add_content, b'key', b'value', True)
    assert b'key' in set(rpc_node.kademlia.content_manager.iter_content_keys())
    assert b'key' not in set(rpc_node.kademlia.content_manager.durable_db.keys())
    assert b'key' in set(rpc_node.kademlia.content_manager.ephemeral_db.keys())
    assert b'key' in set(rpc_node.kademlia.content_manager.cache_db.keys())


import functools
from enum import Enum
from typing import Any, List, Optional, Sequence, TYPE_CHECKING, Tuple, Union

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow.python.ops import array_ops
from typing_extensions import Final

from .config import default_float, default_summary_fmt

if TYPE_CHECKING:
    from IPython.lib import pretty

DType = Union[np.dtype, tf.DType]
VariableData = Union[List, Tuple, np.ndarray, int, float]  # deprecated
Transform = Union[tfp.bijectors.Bijector]
Prior = Union[tfp.distributions.Distribution]

TensorType = Union[tf.Tensor, tf.Variable, "Parameter"]
"""
Type alias for tensor-like types that are supported by most TensorFlow and GPflow operations.

NOTE: Union types like this do not work with the `register` method of `multipledispatch`'s
`Dispatcher` class. Instead use `TensorLike`.
"""

# We've left this as object until we've tested the performance consequences of using the full set
# (np.ndarray, tf.Tensor, tf.Variable, Parameter), see https://github.com/GPflow/GPflow/issues/1434
TensorLike: Final[Tuple[type, ...]] = (object,)
"""
:var TensorLike: Collection of tensor-like types for registering implementations with
    `multipledispatch` dispatchers.
"""

_NativeScalar = Union[int, float]
_Array = Sequence[Any]  # a nested array of int, float, bool etc. kept simple for readability
TensorData = Union[_NativeScalar, _Array, TensorType]

def _IS_PARAMETER(o: object) -> bool:
    return isinstance(o, Parameter)

def _IS_TRAINABLE_PARAMETER(o: object) -> bool:
    return isinstance(o, Parameter) and o.trainable

class Module(tf.Module):
    @property
    def parameters(self) -> Tuple["Parameter", ...]:
        return tuple(self._flatten(predicate=_IS_PARAMETER))

    @property
    def trainable_parameters(self) -> Tuple["Parameter", ...]:
        return tuple(self._flatten(predicate=_IS_TRAINABLE_PARAMETER))

    def _representation_table(self, object_name: str, tablefmt: Optional[str]) -> str:
        from .utilities import leaf_components, tabulate_module_summary

        repr_components = [object_name]
        if leaf_components(self):
            repr_components.append(tabulate_module_summary(self, tablefmt=tablefmt))
        return "\n".join(repr_components)

    def _repr_html_(self) -> str:
        """ Nice representation of GPflow objects in IPython/Jupyter notebooks """
        from html import escape

        return self._representation_table(escape(repr(self)), "html")

    def _repr_pretty_(self, p: "pretty.RepresentationPrinter", cycle: bool) -> None:
        """ Nice representation of GPflow objects in the IPython shell """
        repr_str = self._representation_table(repr(self), default_summary_fmt())
        p.text(repr_str)

class PriorOn(Enum):
    CONSTRAINED = "constrained"
    UNCONSTRAINED = "unconstrained"

class Parameter(tf.Module):
    def __init__(
        self,
        value: TensorData,
        *,
        transform: Optional[Transform] = None,
        prior: Optional[Prior] = None,
        prior_on: Union[str, PriorOn] = PriorOn.CONSTRAINED,
        trainable: bool = True,
        dtype: Optional[DType] = None,
        name: Optional[str] = None,
    ):
        """
        A parameter retains both constrained and unconstrained
        representations. If no transform is provided, these two values will be the same.
        It is often challenging to operate with unconstrained parameters. For example, a variance cannot be negative,
        therefore we need a positive constraint and it is natural to use constrained values.
        A prior can be imposed either on the constrained version (default) or on the unconstrained version of the parameter.
        """
        super().__init__()

        self._transform = transform
        self.prior = prior
        self.prior_on = prior_on  # type: ignore  # see https://github.com/python/mypy/issues/3004

        if isinstance(value, tf.Variable):
            self._unconstrained = value
        else:
            unconstrained_value = self.validate_unconstrained_value(value, dtype)
            self._unconstrained = tf.Variable(
                unconstrained_value, dtype=dtype, name=name, trainable=trainable
            )

    def log_prior_density(self) -> tf.Tensor:
        """ Log of the prior probability density of the constrained variable. """

        if self.prior is None:
            return tf.convert_to_tensor(0.0, dtype=self.dtype)

        y = self.read_value()

        if self.prior_on == PriorOn.CONSTRAINED:
            # evaluation is in same space as prior
            return tf.reduce_sum(self.prior.log_prob(y))

        else:
            # prior on unconstrained, but evaluating log-prior in constrained space
            x = self._unconstrained
            log_p = tf.reduce_sum(self.prior.log_prob(x))

            if self.transform is not None:
                # need to include log|Jacobian| to account for coordinate transform
                log_det_jacobian = self.transform.inverse_log_det_jacobian(y, y.shape.ndims)
                log_p += tf.reduce_sum(log_det_jacobian)

            return log_p

    @property
    def prior_on(self) -> PriorOn:
        return self._prior_on

    @prior_on.setter
    def prior_on(self, value: Union[str, PriorOn]) -> None:
        self._prior_on = PriorOn(value)

    def value(self) -> tf.Tensor:
        return _to_constrained(self._unconstrained.value(), self.transform)  # type: ignore  # assumes _to_constrained returns a tf.Tensor

    def read_value(self) -> tf.Tensor:
        return _to_constrained(self._unconstrained.read_value(), self.transform)  # type: ignore  # assumes _to_constrained returns a tf.Tensor

    def experimental_ref(self) -> "Parameter":
        return self

    def deref(self) -> "Parameter":
        return self

    @property
    def unconstrained_variable(self) -> tf.Variable:
        return self._unconstrained

    @property
    def transform(self) -> Optional[Transform]:
        return self._transform

    @transform.setter
    def transform(self, new_transform: Optional[Transform]) -> None:
        constrained_value = self.read_value()
        self._transform = new_transform
        self.assign(constrained_value)

    @property
    def trainable(self) -> bool:
        """
        `True` if this instance is trainable, else `False`.

        This attribute cannot be set directly. Use :func:`gpflow.set_trainable`.
        """
        return self._unconstrained.trainable

    @property
    def initial_value(self) -> tf.Tensor:
        return self._unconstrained.initial_value

    def validate_unconstrained_value(self, value: TensorData, dtype: DType) -> tf.Tensor:
        value = _cast_to_dtype(value, dtype)
        unconstrained_value = _to_unconstrained(value, self.transform)
        message = (
            "gpflow.Parameter: the value to be assigned is incompatible with this parameter's "
            "transform (the corresponding unconstrained value has NaN or Inf) and hence cannot be "
            "assigned."
        )
        return tf.debugging.assert_all_finite(unconstrained_value, message=message)

    def assign(
        self,
        value: TensorData,
        use_locking: bool = False,
        name: Optional[str] = None,
        read_value: bool = True,
    ) -> tf.Tensor:
        """
        Assigns constrained `value` to the unconstrained parameter's variable.
        It passes constrained value through parameter's transform first.

        Example:
            ```
            a = Parameter(2.0, transform=tfp.bijectors.Softplus())
            b = Parameter(3.0)

            a.assign(4.0)               # `a` parameter to `2.0` value.
            a.assign(tf.constant(5.0))  # `a` parameter to `5.0` value.
            a.assign(b)                 # `a` parameter to constrained value of `b`.
            ```

        :param value: Constrained tensor-like value.
        :param use_locking: If `True`, use locking during the assignment.
        :param name: The name of the operation to be created.
        :param read_value: if True, will return something which evaluates to the new
            value of the variable; if False will return the assign op.
        """
        unconstrained_value = self.validate_unconstrained_value(value, self.dtype)
        return self._unconstrained.assign(
            unconstrained_value, use_locking=use_locking, name=name, read_value=read_value
        )

    @property
    def is_tensor_like(self) -> bool:
        """
        This method means that TensorFlow's `tensor_util.is_tensor` function
        will return `True`
        """
        return True

    @property
    def name(self) -> str:
        return self._unconstrained.name

    @property
    def initializer(self):  # type unknown
        return self._unconstrained.initializer

    @property
    def device(self) -> Optional[str]:
        return self._unconstrained.device

    @property
    def dtype(self) -> tf.DType:
        return self._unconstrained.dtype

    @property
    def op(self) -> tf.Operation:
        return self._unconstrained.op

    @property
    def shape(self) -> tf.TensorShape:
        if self.transform is not None:
            return self.transform.forward_event_shape(self._unconstrained.shape)
        return self._unconstrained.shape

    def numpy(self) -> np.ndarray:
        return self.read_value().numpy()

    def get_shape(self) -> tf.TensorShape:
        return self.shape

    def _should_act_as_resource_variable(self):  # type unknown
        # needed so that Parameters are correctly identified by TensorFlow's
        # is_resource_variable() in resource_variable_ops.py
        pass  # only checked by TensorFlow using hasattr()

    @property
    def handle(self):  # type unknown
        return self._unconstrained.handle

    def __repr__(self) -> str:
        unconstrained = self.unconstrained_variable
        constrained = self.read_value()
        if tf.executing_eagerly():
            info = (
                f"unconstrained-shape={unconstrained.shape} "
                f"unconstrained-value={unconstrained.numpy()} "
                f"constrained-shape={constrained.shape} "
                f"constrained-value={constrained.numpy()}"
            )
        else:
            if unconstrained.shape == constrained.shape:
                info = f"shape={constrained.shape}"
            else:
                info = (
                    f"unconstrained-shape={unconstrained.shape} "
                    f"constrained-shape={constrained.shape}"
                )

        return f"<gpflow.Parameter {self.name!r} dtype={self.dtype.name} {info}>"

    # Below
    # TensorFlow copy-paste code to make variable-like object to work

    @classmethod
    def _OverloadAllOperators(cls):  # pylint: disable=invalid-name
        """Register overloads for all operators."""
        for operator in tf.Tensor.OVERLOADABLE_OPERATORS:
            cls._OverloadOperator(operator)
        # For slicing, bind getitem differently than a tensor (use SliceHelperVar
        # instead)
        # pylint: disable=protected-access
        setattr(cls, "__getitem__", array_ops._SliceHelperVar)

    @classmethod
    def _OverloadOperator(cls, operator):  # pylint: disable=invalid-name
        """Defer an operator overload to `ops.Tensor`.

        We pull the operator out of ops.Tensor dynamically to avoid ordering issues.

        Args:
            operator: string. The operator name.
        """
        tensor_oper = getattr(tf.Tensor, operator)

        def _run_op(a, *args, **kwargs):
            # pylint: disable=protected-access
            return tensor_oper(a.read_value(), *args, **kwargs)

        functools.update_wrapper(_run_op, tensor_oper)
        setattr(cls, operator, _run_op)

    # NOTE(mrry): This enables the Variable's overloaded "right" binary
    # operators to run when the left operand is an ndarray, because it
    # accords the Variable class higher priority than an ndarray, or a
    # numpy matrix.
    # TODO(mrry): Convert this to using numpy's __numpy_ufunc__
    # mechanism, which allows more control over how Variables interact
    # with ndarrays.
    __array_priority__ = 100

Parameter._OverloadAllOperators()
tf.register_tensor_conversion_function(Parameter, lambda x, *args, **kwds: x.read_value())

def _cast_to_dtype(
    value: TensorData, dtype: Optional[DType] = None
) -> Union[tf.Tensor, tf.Variable]:
    if dtype is None:
        dtype = default_float()

    if tf.is_tensor(value):
        # NOTE(awav) TF2.2 resolves issue with cast.
        # From TF2.2, `tf.cast` can be used alone instead of this auxiliary function.
        # workaround for https://github.com/tensorflow/tensorflow/issues/35938
        return tf.cast(value, dtype)
    else:
        return tf.convert_to_tensor(value, dtype=dtype)

def _to_constrained(value: TensorType, transform: Optional[Transform]) -> TensorType:
    if transform is not None:
        return transform.forward(value)
    return value

def _to_unconstrained(value: TensorType, transform: Optional[Transform]) -> TensorType:
    if transform is not None:
        return transform.inverse(value)
    return value


import time
import torch
import torch.nn as nn
import numpy as np
from torch import Tensor
from torch.nn.init import kaiming_normal_, constant_
from torchvision.transforms import Resize
from torch_scatter import scatter_mean
from typing import Callable, Tuple
from fea_transfer import poolfeat, upfeat

TIME_WINDOW = 24
PRED_LEN = 6

def _count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

class Model(nn.Module):
    # 原 station_mask 中 False 的地方代表站点，现在我们需要 True 的地方代表站点，所以取反（~）
    station_mask: np.ndarray = ~np.load("/home/guoshuai/coding/PRD-MSGNN/final_dataset/prd_station_mask.npy")
    in_resizer = Resize(size=(128, 128))
    out_resizer = Resize(size=(163, 137))

    def __init__(self, gnn_h, cnn_h, rnn_h, rnn_l, group_em, pm25_em, gnn_layer, input_fea_num, edge_h, group_num, device) -> None:
        super(Model, self).__init__()
        self.rnn_l = rnn_l
        self.rnn_h = rnn_h
        self.gnn_layer = gnn_layer
        self.group_num = group_num
        self.device = device
        # self.superpixel_size = superpixel_size
        # self.group_num_w = group_num_w
        self.aod_encoder = CNNEncoder(7)
        self.wea_encoder = CNNEncoder(11)
        self.land_encoder = CNNEncoder(28)
        self.aod_tans = Pattern_trans()
        self.wea_tans = Pattern_trans()
        self.land_tans = Pattern_trans()
        self.cnn_predict = CNNPredict()
        # self.group_gnn: Callable[..., Tensor] = GroupGNN_Layer(group_em, 1, gnn_h)
        self.aod_group_gnn = nn.ModuleList([GroupGNN_Layer(group_em, edge_h, gnn_h)])
        for i in range(self.gnn_layer - 1):
            self.aod_group_gnn.append(GroupGNN_Layer(gnn_h, edge_h, gnn_h))

        self.wea_group_gnn = nn.ModuleList([GroupGNN_Layer(group_em, edge_h, gnn_h)])
        for i in range(self.gnn_layer - 1):
            self.wea_group_gnn.append(GroupGNN_Layer(gnn_h, edge_h, gnn_h))

        self.land_group_gnn = nn.ModuleList([GroupGNN_Layer(group_em, edge_h, gnn_h)])
        for i in range(self.gnn_layer - 1):
            self.land_group_gnn.append(GroupGNN_Layer(gnn_h, edge_h, gnn_h))
        # self.lstm_encoder = LSTMEncoder(
        #     input_size=gnn_h + cnn_h, hidden_size=rnn_h, num_layers=rnn_l
        # )
        # self.lstm_decoder: Callable[..., Tuple[Tensor, Tensor, Tensor]] = LSTMDecoder(
        #     input_size=pm25_em, hidden_size=rnn_h, num_layers=rnn_l
        # )
        self.aod_superpixel_encode = SuperPixel_encoder(1, 8, device, self.group_num)
        self.wea_superpixel_encode = SuperPixel_encoder(5, 8, device, self.group_num)
        self.land_superpixel_encode = SuperPixel_encoder(22, 8, device, self.group_num)
        self.aod_group_embed: Callable[..., Tensor] = nn.Sequential(
            nn.Linear(7, group_em), nn.ReLU()
        )
        self.land_group_embed: Callable[..., Tensor] = nn.Sequential(
            nn.Linear(28, group_em), nn.ReLU()
        )
        self.wea_group_embed: Callable[..., Tensor] = nn.Sequential(
            nn.Linear(11, group_em), nn.ReLU()
        )
        self.aod_edge_inf = nn.Sequential(
            nn.Linear(group_em * 2, edge_h), nn.ReLU(inplace=True)
        )
        self.wea_edge_inf = nn.Sequential(
            nn.Linear(group_em * 2, edge_h), nn.ReLU(inplace=True)
        )
        self.land_edge_inf = nn.Sequential(
            nn.Linear(group_em * 2, edge_h), nn.ReLU(inplace=True)
        )
        self.attention = nn.Sequential(
            nn.Linear(32, 1), nn.Sigmoid()
        )

    @staticmethod
    def batch_input(x: Tensor, edge_w: Tensor, edge_conn: Tensor):
        sta_num = x.shape[1]
        x = x.reshape(-1, x.shape[-1])
        edge_w = edge_w.reshape(-1, edge_w.shape[-1])
        edge_index = torch.clone(edge_conn)
        for i in range(edge_conn.size(0)):
            edge_index[i, :] = torch.add(edge_index[i, :], i * sta_num)
        edge_index = edge_index.transpose(0, 1)
        edge_index = edge_index.reshape(2, -1)
        return x, edge_w, edge_index

    # def edge_encode(self, g_x:Tensor, g_edge_index):
    #     g_edge_index = g_edge_index.clone().transpose(0, 1)
    #     row, col = g_edge_index
    #     g_x_concat = torch.concat((g_x[:, row], g_x[:, col]), dim=2)
    #     g_edge_w = self.edge_inf(g_x_concat)
    #     g_edge_index = g_edge_index.unsqueeze(dim=0)
    #     g_edge_index = g_edge_index.repeat_interleave(g_x.shape[0], dim=0)
    #     # g_edge_index = g_edge_index.transpose(1, 2)
    #     return g_edge_index, g_edge_w

    @staticmethod
    def edge_encode(g_x:Tensor, g_edge_index, edge_inf):
        g_edge_index = g_edge_index.clone().transpose(0, 1)
        row, col = g_edge_index
        g_x_concat = torch.concat((g_x[:, row], g_x[:, col]), dim=2)
        g_edge_w = edge_inf(g_x_concat)
        g_edge_index = g_edge_index.unsqueeze(dim=0)
        g_edge_index = g_edge_index.repeat_interleave(g_x.shape[0], dim=0)
        # g_edge_index = g_edge_index.transpose(1, 2)
        return g_edge_index, g_edge_w

    def score_com(self, new_x_aod):
        # shape:[32, 3*32, 128, 128]
        new_x = new_x_aod.clone()
        new_x = new_x.permute(0, 2, 3, 1)
        atte_score: Tensor = self.attention(new_x)

        # 形状变换: [32, 128, 128, 1]->[32, 1 128, 128]
        atte_score = atte_score.permute(0, 3, 1, 2)

        return atte_score

    def forward(
        self,
        x: Tensor,
        land_use,
        device,
    ):
        # 转换输入 x 的shape：[32,156,151,7]->[32,7,156,151]
        x = x.permute(0, 3, 1, 2)
        air = x[:, :6, :, :]

        ''' 取出 aod 值，计算 assi_matrix，将空气质量比较类似的 pixel 聚合到一起'''
        # aod 的 shape：[32,156,151,1], aod_assi_matrix shape:[32, 156, 151, 16, 16]
        aod = x[:, 6, :, :].unsqueeze(1)

        # 计算在 AOD superpixel 下的 group
        # 先将原始输入作一下 resize，[32,7,156,151]->[32,7,128,128]
        aod_x = torch.cat([air, aod], dim=1)

        # 计算下层特征
        aod_x_spatial_fea = self.aod_encoder(aod_x)

        # 计算 aod_assi_matrix，aod_assi_matrix:[b, 256, 128, 128]
        aod_assi_matrix = self.aod_superpixel_encode(aod_x_spatial_fea, device)

        resize_aodx = self.in_resizer(aod_x)

        # 更新上层特征
        # 新传变换:[32, 256, 128, 128]->[32, 128, 128, 256]
        aod_assi_matrix = aod_assi_matrix.permute(0, 2, 3, 1)
        g_x_aod = torch.einsum('abcd,acde->abe', resize_aodx, aod_assi_matrix)
        # g_x_aod = poolfeat(resize_aodx, aod_assi_matrix, device, self.superpixel_size, self.superpixel_size)

        # 形状转换 [32, 7, 256]->[32, 256, 7]
        g_x_aod = g_x_aod.permute(0, 2, 1)

        # 计算每个 superpixel 的特征表示
        g_x_aod = self.aod_group_embed(g_x_aod)

        # edge_index shape:[batch_size, edge_num, 2]->[batch_size, 2, edge_num]
        g_edge_index = torch.tensor(
            [
                [i, j]
                for i in range(self.group_num)
                for j in range(self.group_num)
                if i != j
            ],
            dtype=int,
            device=self.device,
        )
        # 建模 superpixel 之间的关系
        g_edge_index_aod, g_edge_w_aod = self.edge_encode(g_x_aod, g_edge_index, self.aod_edge_inf)

        # 输入 batchinput 将数据处理成适合 GNN 的输入，并输入 GNN 网络
        g_x_aod, g_edge_w_aod, g_edge_index_aod = self.batch_input(g_x_aod, g_edge_w_aod, g_edge_index_aod)
        for i in range(self.gnn_layer):
            g_x_aod = self.aod_group_gnn[i](g_x_aod, g_edge_index_aod, g_edge_w_aod)
        # shape [batch*256, 32]->[32, 256, 32]->[32(b), 32, 16, 16]
        gx_spatial_fea_aod = g_x_aod
        gx_spatial_fea_aod = gx_spatial_fea_aod.reshape(
            -1, self.group_num, gx_spatial_fea_aod.shape[-1]
        )

        # new_x_aod = upfeat(gx_spatial_fea_aod, aod_assi_matrix, self.superpixel_size, self.superpixel_size)
        new_x_aod = torch.einsum('abcd,ade->abce', aod_assi_matrix, gx_spatial_fea_aod)

        # 形状变换
        new_x_aod = new_x_aod.permute(0, 3, 1, 2)

        # 此时 new_x shape：[32, 32, 128, 128]
        # 形状变化:[32, 32, 128, 128]->[32, 32, 156, 151]
        new_x_aod = self.out_resizer(new_x_aod)

        ''' 取出 weather 值，计算 assi_matrix, 将天气情况比较类似的 pixel 聚合到一起'''
        # wea 的 shape：[32,156,151,5], aod_assi_matrix shape:[32, 156, 151, 16, 16]
        weather = x[:, 7:12, :, :]

        # 先将原始输入作一下 resize，[32,7,156,151]->[32,7,128,128]
        wea_x = torch.cat([air, weather], dim=1)

        # 计算下层特征
        wea_x_spatial_fea = self.wea_encoder(wea_x)

        # 计算 wea_assi_matrix，wea_assi_matrix:[b, 9, 128, 128]
        wea_assi_matrix = self.wea_superpixel_encode(wea_x_spatial_fea, device)

        resize_weax = self.in_resizer(wea_x)

        # 计算在 AOD superpixel 下的 group
        wea_assi_matrix = wea_assi_matrix.permute(0, 2, 3, 1)
        g_x_wea = torch.einsum('abcd,acde->abe', resize_weax, wea_assi_matrix)

        # g_x_wea = poolfeat(resize_weax, wea_assi_matrix, device, self.superpixel_size, self.superpixel_size)

        # 形状转换 [32, 7, 256]->[32, 256, 7]
        g_x_wea = g_x_wea.permute(0, 2, 1)

        # 计算每个 superpixel 的特征表示
        g_x_wea = self.wea_group_embed(g_x_wea)

        # edge_index shape:[batch_size, edge_num, 2]->[batch_size, 2, edge_num]
        # 建模 superpixel 之间的关系
        g_edge_index_wea, g_edge_w_wea = self.edge_encode(g_x_wea, g_edge_index, self.wea_edge_inf)

        # 输入 batchinput 将数据处理成适合 GNN 的输入，并输入 GNN 网络
        g_x_wea, g_edge_w_wea, g_edge_index_wea = self.batch_input(g_x_wea, g_edge_w_wea, g_edge_index_wea)
        for i in range(self.gnn_layer):
            g_x_wea = self.wea_group_gnn[i](g_x_wea, g_edge_index_wea, g_edge_w_wea)
        # shape [batch*256, 32]->[32, 16, 16, 32]->[32(b), 32, 16, 16]
        gx_spatial_fea_wea = g_x_wea
        gx_spatial_fea_wea = gx_spatial_fea_wea.reshape(
            -1, self.group_num, gx_spatial_fea_wea.shape[-1]
        )

        new_x_wea = torch.einsum('abcd,ade->abce', wea_assi_matrix, gx_spatial_fea_wea)

        # 形状变换
        new_x_wea = new_x_wea.permute(0, 3, 1, 2)

        # 此时 new_x shape：[32, 32, 128, 128]
        # 形状变化:[32, 32, 128, 128]->[32, 32, 156, 151]
        new_x_wea = self.out_resizer(new_x_wea)

        ''' 取出 land_assi_matrix, 将 land_use 比较类似的 pixel 聚合到一起'''
        # land 的 shape：[32,156,151,22], aod_assi_matrix shape:[32, 156, 151, 16, 16]
        # 形状变化
        land_use = land_use.unsqueeze(dim=0)
        land_use = land_use.permute(0, 3, 1, 2)
        # 先将原始输入作一下 resize，[32,7,156,151]->[32,7,128,128]
        land_use = land_use.repeat_interleave(x.shape[0], dim=0)
        land_x = torch.cat([air, land_use], dim=1)

        # 计算下层特征
        land_x_spatial_fea = self.land_encoder(land_x)

        # 计算 wea_assi_matrix，wea_assi_matrix:[b, 9, 128, 128]
        land_assi_matrix = self.land_superpixel_encode(land_x_spatial_fea, device)

        resize_landx = self.in_resizer(land_x)

        # 更新上层特征
        # 新传变换:[32, 256, 128, 128]->[32, 128, 128, 256]
        land_assi_matrix = land_assi_matrix.permute(0, 2, 3, 1)
        g_x_land = torch.einsum('abcd,acde->abe', resize_landx, land_assi_matrix)

        # 形状转换 [32, 7, 256]->[32, 256, 7]
        g_x_land = g_x_land.permute(0, 2, 1)

        # 计算每个 superpixel 的特征表示
        g_x_land = self.land_group_embed(g_x_land)

        # edge_index shape:[batch_size, edge_num, 2]->[batch_size, 2, edge_num]
        # 建模 superpixel 之间的关系
        g_edge_index_land, g_edge_w_land = self.edge_encode(g_x_land, g_edge_index, self.land_edge_inf)

        # 输入 batchinput 将数据处理成适合 GNN 的输入，并输入 GNN 网络
        g_x_land, g_edge_w_land, g_edge_index_land = self.batch_input(g_x_land, g_edge_w_land, g_edge_index_land)
        for i in range(self.gnn_layer):
            g_x_land = self.land_group_gnn[i](g_x_land, g_edge_index_land, g_edge_w_land)
        # shape [batch*256, 32]->[32, 16, 16, 32]->[32(b), 32, 16, 16]
        gx_spatial_fea_land = g_x_land
        gx_spatial_fea_land = gx_spatial_fea_land.reshape(
            -1, self.group_num, gx_spatial_fea_land.shape[-1]
        )

        # new_x_aod = upfeat(gx_spatial_fea_aod, aod_assi_matrix, self.superpixel_size, self.superpixel_size)
        new_x_land = torch.einsum('abcd,ade->abce', land_assi_matrix, gx_spatial_fea_land)

        # 形状变换
        new_x_land = new_x_land.permute(0, 3, 1, 2)

        # 此时 new_x shape：[32, 32, 128, 128]
        # 形状变化:[32, 32, 128, 128]->[32, 32, 156, 151]
        new_x_land = self.out_resizer(new_x_land)

        # 形状变换
        aod_x_spatial_fea = self.out_resizer(aod_x_spatial_fea)
        wea_x_spatial_fea = self.out_resizer(wea_x_spatial_fea)
        land_x_spatial_fea = self.out_resizer(land_x_spatial_fea)

        # 此时 x_spatial_fea 的形状是：[32, 16, 156, 151]

        aod_group_cat = torch.cat([aod_x_spatial_fea, new_x_aod], dim=1)
        wea_group_cat = torch.cat([wea_x_spatial_fea, new_x_wea], dim=1)
        land_group_cat = torch.cat([land_x_spatial_fea, new_x_land], dim=1)

        # 形状变换一下，都给变到 32
        aod_group_trans =  self.aod_tans(aod_group_cat)
        wea_group_trans =  self.wea_tans(wea_group_cat)
        land_group_trans =  self.land_tans(land_group_cat)

        # 计算重要性
        aod_impor = self.score_com(aod_group_trans)
        wea_impor = self.score_com(wea_group_trans)
        land_impor = self.score_com(land_group_trans)

        # x_group_cat:[32, 32+32+32, 156, 151 ]
        x_group_cat = torch.cat([aod_group_trans * aod_impor, wea_group_trans * wea_impor, land_group_trans * land_impor], dim=1)

        # 输入 CNN 当中做 estimate
        # predict shape:[32, 1, 156, 151]
        outputs = self.cnn_predict(x_group_cat) # 改
        outputs = outputs.reshape(outputs.shape[0], outputs.shape[2], outputs.shape[3])
        # end = time.perf_counter_ns()
        # print(f"Model.forward() time: {(end - start) / 1000000 :,} ms")
        return outputs, aod_assi_matrix, wea_assi_matrix, land_assi_matrix

class SuperPixel_encoder(nn.Module):
    in_resizer = Resize(size=(128, 128))
    out_resizer = Resize(size=(163, 137))

    def __init__(self, in_channel, out_channel, device, group_num, batch_norm=True) -> None:
        super().__init__()
        self.assign_ch = group_num
        self.conv1 = self._conv(batch_norm, 16, 32)
        self.conv2 = self._conv(batch_norm, 32, 64)
        self.pred_mask0 = self.predict_mask(64, self.assign_ch)

        print(f"[SuperPixel_encoder] Number of params: {_count_parameters(self):,}")

        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
                kaiming_normal_(m.weight, 0.1)
                if m.bias is not None:
                    constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                constant_(m.weight, 1)
                constant_(m.bias, 0)

    # def expand_sp(self, sp_feat):
    #     sp_feat = nn.functional.pad(sp_feat, [1,1,1,1])
    #     b, c, h, w = sp_feat.shape
    #     output_list = []
    #     #this loop is acceptable due to the lower h, w
    #     for i in range(1,h-1):
    #         row_list = []
    #         for j in range(1,w-1):
    #             sp_patch = sp_feat[:,:, (i-1):(i+2), (j-1):(j+2)]
    #             sp_patch = sp_patch.repeat(1,1,8,8)
    #             row_list.append(sp_patch)

    #         output_list.append(torch.cat(row_list, dim=-1))

    #     output = torch.cat(output_list, dim=-2)

    #     return output

    # def expand_pixel(self, pixel_feat):
    #     b,c,h,w = pixel_feat.shape
    #     pixel_feat = pixel_feat.view(b,c,h,1,w,1)
    #     pixel_feat = pixel_feat.repeat(1,1,1,3,1,3)
    #     pixel_feat = pixel_feat.reshape(b,c,h*3,w*3)

    #     pixel_feat = pixel_feat * self.mask_select

    #     return pixel_feat

    def forward(self, x: Tensor, device):
        # 计算下采样所用时间
        conv1 = self.conv1(x)
        conv2 = self.conv2(conv1)
        mask0 = self.pred_mask0(conv2)
        prob0 = nn.functional.softmax(mask0, dim=1)

        # 计算上采样所用时间

        # print(f'上采样所用时间：{up_consume}')

        return prob0

    @staticmethod
    def predict_mask(in_planes, channel=9):
        return  nn.Conv2d(in_planes, channel, kernel_size=3, stride=1, padding=1, bias=True)

    @staticmethod
    def _conv(
        batch_norm, in_planes, out_planes, kernel_size=3, stride=1
    ) -> Callable[..., Tensor]:
        if batch_norm:
            return nn.Sequential(
                nn.Conv2d(
                    in_planes,
                    out_planes,
                    kernel_size=kernel_size,
                    stride=stride,
                    padding=(kernel_size - 1) // 2,
                    bias=False,
                ),
                nn.BatchNorm2d(out_planes),
                nn.LeakyReLU(0.1),
            )
        else:
            return nn.Sequential(
                nn.Conv2d(
                    in_planes,
                    out_planes,
                    kernel_size=kernel_size,
                    stride=stride,
                    padding=(kernel_size - 1) // 2,
                    bias=True,
                ),
                nn.LeakyReLU(0.1),
            )

    @staticmethod
    def _deconv(in_planes, out_planes) -> Callable[..., Tensor]:
        return nn.Sequential(
            nn.ConvTranspose2d(
                in_planes, out_planes, kernel_size=4, stride=2, padding=1, bias=True
            ),
            nn.LeakyReLU(0.1),
        )

class CNNPredict(nn.Module):
    # in_resizer = Resize(size=(128, 128))
    # out_resizer = Resize(size=(156, 151))

    def __init__(self, batch_norm=True) -> None:
        super().__init__()
        self._conv0a = self._conv(batch_norm, 96, 64)
        self._conv1a = self._conv(batch_norm, 64, 32)
        self.predict = self._conv(batch_norm, 32, 1)

        print(f"[CNNPredict] Number of params: {_count_parameters(self):,}")

        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
                kaiming_normal_(m.weight, 0.1)
                if m.bias is not None:
                    constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                constant_(m.weight, 1)
                constant_(m.bias, 0)

    def forward(self, x: Tensor):
        # x = self.in_resizer(x)
        out0 = self._conv0a(x)  # 5*5
        out1 = self._conv1a(out0)
        out = self.predict(out1)  # 11*11
        return out

    @staticmethod
    def _conv(
        batch_norm, in_planes, out_planes, kernel_size=3, stride=1
    ) -> Callable[..., Tensor]:
        if batch_norm:
            return nn.Sequential(
                nn.Conv2d(
                    in_planes,
                    out_planes,
                    kernel_size=kernel_size,
                    stride=stride,
                    padding=(kernel_size - 1) // 2,
                    bias=False,
                ),
                nn.BatchNorm2d(out_planes),
                nn.LeakyReLU(0.1),
            )
        else:
            return nn.Sequential(
                nn.Conv2d(
                    in_planes,
                    out_planes,
                    kernel_size=kernel_size,
                    stride=stride,
                    padding=(kernel_size - 1) // 2,
                    bias=True,
                ),
                nn.LeakyReLU(0.1),
            )

class Pattern_trans(nn.Module):
    # in_resizer = Resize(size=(128, 128))
    # out_resizer = Resize(size=(156, 151))

    def __init__(self, batch_norm=True) -> None:
        super().__init__()
        self._conv0a = self._conv(batch_norm, 48, 32)

        print(f"[Pattern_trans] Number of params: {_count_parameters(self):,}")

        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
                kaiming_normal_(m.weight, 0.1)
                if m.bias is not None:
                    constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                constant_(m.weight, 1)
                constant_(m.bias, 0)

    def forward(self, x: Tensor):
        # x = self.in_resizer(x)
        out0 = self._conv0a(x)  # 5*5

        return out0

    @staticmethod
    def _conv(
        batch_norm, in_planes, out_planes, kernel_size=1, stride=1
    ) -> Callable[..., Tensor]:
        if batch_norm:
            return nn.Sequential(
                nn.Conv2d(
                    in_planes,
                    out_planes,
                    kernel_size=kernel_size,
                    stride=stride,
                    padding=(kernel_size - 1) // 2,
                    bias=False,
                ),
                nn.BatchNorm2d(out_planes),
                nn.LeakyReLU(0.1),
            )
        else:
            return nn.Sequential(
                nn.Conv2d(
                    in_planes,
                    out_planes,
                    kernel_size=kernel_size,
                    stride=stride,
                    padding=(kernel_size - 1) // 2,
                    bias=True,
                ),
                nn.LeakyReLU(0.1),
            )

class CNNEncoder(nn.Module):
    in_resizer = Resize(size=(128, 128))
    out_resizer = Resize(size=(163, 137))

    def __init__(self, in_channel = 0, batch_norm=True) -> None:
        super().__init__()
        self._conv0a = self._conv(batch_norm, in_channel, 16)
        self._conv0b = self._conv(batch_norm, 16, 16)
        self._drop0 = nn.Dropout(0.2)

        self._conv1a = self._conv(batch_norm, 16, 32, stride=2)
        self._conv1b = self._conv(batch_norm, 32, 32)
        self._drop1 = nn.Dropout(0.4)

        self._conv2a = self._conv(batch_norm, 32, 64, stride=2)
        self._conv2b = self._conv(batch_norm, 64, 64)
        self._drop2 = nn.Dropout(0.5)

        self._conv3a = self._conv(batch_norm, 64, 128, stride=2)
        self._conv3b = self._conv(batch_norm, 128, 128)

        # self._conv4a = self._conv(batch_norm, 128, 256, stride=2)
        # self._conv4b = self._conv(batch_norm, 256, 256)

        # self._deconv3 = self._deconv(256, 128)
        # self._conv3_1 = self._conv(batch_norm, 256, 128)

        self._deconv2 = self._deconv(128, 64)
        self._conv2_1 = self._conv(batch_norm, 128, 64)

        self._deconv1 = self._deconv(64, 32)
        self._conv1_1 = self._conv(batch_norm, 64, 32)

        self._deconv0 = self._deconv(32, 16)
        self._conv0_1 = self._conv(batch_norm, 32, 16)

        print(f"[CNNEncoder] Number of params: {_count_parameters(self):,}")

        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
                kaiming_normal_(m.weight, 0.1)
                if m.bias is not None:
                    constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                constant_(m.weight, 1)
                constant_(m.bias, 0)

    def forward(self, x: Tensor):
        x = self.in_resizer(x)
        out1 = self._conv0b(self._conv0a(x))  # 5*5
        out1 = self._drop0(out1)
        out2 = self._conv1b(self._conv1a(out1))  # 11*11
        out2 = self._drop1(out2)
        out3 = self._conv2b(self._conv2a(out2))  # 23*23
        out3 = self._drop2(out3)
        out4 = self._conv3b(self._conv3a(out3))  # 47*47
        # out5 = self._conv4b(self._conv4a(out4))  # 95*95

        # out_deconv3 = self._deconv3(out5)
        # concat3 = torch.cat((out4, out_deconv3), 1)
        # out_conv3_1 = self._conv3_1(concat3)

        out_deconv2 = self._deconv2(out4)
        concat2 = torch.cat((out3, out_deconv2), 1)
        out_conv2_1 = self._conv2_1(concat2)

        out_deconv1 = self._deconv1(out_conv2_1)
        # out_deconv1 = self._drop2(out_deconv1)
        concat1 = torch.cat((out2, out_deconv1), 1)
        out_conv1_1 = self._conv1_1(concat1)
        # out_conv1_1 = self._drop2(out_conv1_1)

        out_deconv0 = self._deconv0(out_conv1_1)
        # out_deconv0 = self._drop2(out_deconv0)
        concat0 = torch.cat((out1, out_deconv0), 1)
        out_conv0_1 = self._conv0_1(concat0)
        # out_conv0_1 = self._drop1(out_conv0_1)

        # out_resized = self.out_resizer(out_conv0_1)

        return out_conv0_1

    @staticmethod
    def _conv(
        batch_norm, in_planes, out_planes, kernel_size=3, stride=1
    ) -> Callable[..., Tensor]:
        if batch_norm:
            return nn.Sequential(
                nn.Conv2d(
                    in_planes,
                    out_planes,
                    kernel_size=kernel_size,
                    stride=stride,
                    padding=(kernel_size - 1) // 2,
                    bias=False,
                ),
                nn.BatchNorm2d(out_planes),
                nn.LeakyReLU(0.1),
            )
        else:
            return nn.Sequential(
                nn.Conv2d(
                    in_planes,
                    out_planes,
                    kernel_size=kernel_size,
                    stride=stride,
                    padding=(kernel_size - 1) // 2,
                    bias=True,
                ),
                nn.LeakyReLU(0.1),
            )

    @staticmethod
    def _deconv(in_planes, out_planes) -> Callable[..., Tensor]:
        return nn.Sequential(
            nn.ConvTranspose2d(
                in_planes, out_planes, kernel_size=4, stride=2, padding=1, bias=True
            ),
            nn.LeakyReLU(0.1),
        )

class GroupGNN_Layer(nn.Module):
    def __init__(self, node_h, edge_h, gnn_h) -> None:
        super(GroupGNN_Layer, self).__init__()
        self.node_mlp_1 = nn.Sequential(
            nn.Linear(2 * node_h + edge_h, gnn_h), nn.ReLU(inplace=True)
        )
        self.node_mlp_2 = nn.Sequential(
            nn.Linear(node_h + gnn_h, gnn_h), nn.ReLU(inplace=True)
        )

        print(f"[GroupGNN] Number of params: {_count_parameters(self):,}")

    def forward(self, x, edge_index, edge_attr):
        # x: [N, F_x], where N is the number of nodes.
        # edge_index: [2, E] with max entry N - 1.
        # edge_attr: [E, F_e]
        row, col = edge_index
        out = torch.cat([x[row], x[col], edge_attr], dim=1)
        out = self.node_mlp_1(out)
        out = scatter_mean(out, col, dim=0, dim_size=x.size(0))
        out = torch.cat([x, out], dim=1)
        return self.node_mlp_2(out)


import re
import difflib
import enum
from typing import List

from functools import cached_property
from wcwidth import wcswidth
from . import processor
from .helper import RE_URL

sentence_sep_s = [
    r'\r', r'\n',
    r'\.', r',', r';', r'\?', r'!', r':', r'"', r"'", r'\[', r'\]', r'\(', r'\)',
    r'。', r'，', r'；', r'？', r'！', r'：', r'…', r'、',
    r'“', r'”', r'‘', r'’', r'【', r'】', r'《', r'》', r'（', r'）', r'〈', r'〉',
]

_SEP_S = r'(?:[\s\d]*(?:{})[\s\d]*)+'.format('|'.join(sentence_sep_s))
RE_SENTENCE_SEP = re.compile(r'(?:{})|(?:{})'.format(RE_URL.pattern, _SEP_S), re.I)

def split_sentences(text: str, keep_short: bool = False) -> List[str]:
    """
    中英文分句
    """
    if not text:
        return []
    sentences = RE_SENTENCE_SEP.split(text)
    sentences = list(filter(None, (x.strip() for x in sentences)))
    if sentences and not sentences[0]:
        sentences = sentences[1:]
    if sentences and not sentences[-1]:
        sentences = sentences[:-1]
    if not keep_short:
        sentences = [x for x in sentences if not is_short_sentence(x)]
    return sentences

def is_short_sentence(sentence: str) -> bool:
    """
    >>> is_short_sentence('')
    True
    >>> is_short_sentence('hello')
    True
    >>> is_short_sentence('你好呀')
    True
    >>> is_short_sentence('hello world')
    False
    >>> is_short_sentence('你好世界!')
    False
    >>> is_short_sentence('aaaa bbbb cccc dddd eeee ffff')
    False
    """
    return len(sentence) <= 16 and wcswidth(sentence) <= 8

def is_summary_prob(subtext: str, fulltext: str) -> float:
    """
    判断subtext是fulltext的摘要的概率。概率大于0.5可认为是。
    """
    sub_sentences = split_sentences(subtext)
    full_sentences = split_sentences(fulltext)
    if not sub_sentences:
        return 1.0 if full_sentences else 0.0
    elif not full_sentences:
        return 0.0
    num_sub = len(sub_sentences)
    num_full = len(full_sentences)
    if num_sub - num_full >= 0:
        return 0.0
    max_check = min(num_sub * 2 + 1, num_full)
    num_positive = 0
    prev_value = 0
    for delta in difflib.ndiff(sub_sentences, full_sentences[:max_check]):
        diff_type = delta[0]
        if diff_type == ' ':
            value = 1
        elif diff_type == '+':
            value = 0.9
        elif diff_type == '-':
            value = -2.5
        elif diff_type == '?':
            value = 1.5 - prev_value * 0.5
        else:
            value = 0
        num_positive += value
        prev_value = value
    return max(0.0, min(1.0, num_positive / max_check))

def is_summary(subtext: str, fulltext: str) -> bool:
    return is_summary_prob(subtext, fulltext) > 0.5

class StoryContentInfo:
    def __init__(self, html: str, *, sentence_count: int = None):
        self.html = html
        self._sentence_count = sentence_count

    def __bool__(self):
        return bool(self.html)

    @cached_property
    def text(self) -> str:
        return processor.story_html_to_text(self.html)

    @property
    def length(self) -> int:
        return len(self.html)

    @cached_property
    def link_count(self) -> int:
        return processor.story_link_count(self.html)

    @cached_property
    def image_count(self) -> int:
        return processor.story_image_count(self.html)

    @cached_property
    def url_count(self) -> int:
        return processor.story_url_count(self.html)

    @property
    def sentence_count(self) -> int:
        if self._sentence_count is None:
            self._sentence_count = len(split_sentences(self.text))
        return self._sentence_count

def is_fulltext_content(story_content_info: StoryContentInfo):
    """
    detect whether the full content is already in rss feed.
    """
    if not story_content_info:
        return False
    if story_content_info.length >= 2000:
        return True
    if story_content_info.link_count >= 5:
        return True
    if story_content_info.url_count >= 7:
        return True
    if story_content_info.image_count >= 2:
        return True
    return False

class FulltextAcceptStrategy(enum.Enum):
    REJECT = "REJECT"
    REPLACE = "REPLACE"
    APPEND = "APPEND"

def decide_accept_fulltext(new_info: StoryContentInfo, old_info: StoryContentInfo) -> FulltextAcceptStrategy:
    if not new_info:
        return FulltextAcceptStrategy.REJECT
    if is_fulltext_content(old_info):
        is_accept = (
            is_fulltext_content(new_info)
            and new_info.sentence_count > old_info.sentence_count
            and new_info.image_count >= old_info.image_count
            and new_info.link_count >= old_info.link_count
            and new_info.url_count >= old_info.url_count
        )
        if is_accept:
            return FulltextAcceptStrategy.REPLACE
        else:
            return FulltextAcceptStrategy.REJECT
    else:
        if is_summary(old_info.text, new_info.text):
            old_url_match = RE_URL.search(old_info.html)
            # eg: hackernews only has a comments link, but web page not contains the link
            if old_url_match:
                old_url = old_url_match.group(0)
                if old_url not in new_info.html:
                    return FulltextAcceptStrategy.APPEND
            if not (
                new_info.image_count >= old_info.image_count
                and new_info.link_count >= old_info.link_count
                and new_info.url_count >= old_info.url_count
            ):
                return FulltextAcceptStrategy.APPEND
            return FulltextAcceptStrategy.REPLACE
        elif is_fulltext_content(new_info):
            return FulltextAcceptStrategy.APPEND
        else:
            return FulltextAcceptStrategy.REJECT


from flask import Flask, request, g
from flask import send_file
from flask import jsonify

app = Flask(__name__, instance_relative_config=True)

@app.route('/count', methods=['get'])
def predict():
    f = open("tmp_count.txt", "r")
    people_count = f.read()
    return jsonify(people_count)

@app.route('/', methods=['get'])
def root():
    print('lol')
    return send_file('index.html')

if __name__ == "__main__":
    app.run(host='0.0.0.0')


__author__ = 'jobin'

from django.http import JsonResponse
from django.views.generic import View
from django.forms import Form, IntegerField, CharField

from .check_request import CheckRequest
from api.response_code import errorCode

class UpdateUserForm(Form):
    username = CharField(label=u'用户名：', max_length=50)
    sex = IntegerField()
    age = IntegerField(min_value=15, max_value=100)
    yearsOfWorking = IntegerField(min_value=0, max_value=60)
    school = CharField(max_length=40)
    education = IntegerField()

class Index(View):
    def get(self, request):
        check = CheckRequest(request);
        if not check.user:
            return JsonResponse({
                "status": "error",
                "code":10000,
                "msg": errorCode[10000]
            })
        return JsonResponse({"status" :  "success",
                "msg" :  '',
                "data" :  {
                    "username" :  check.user.username,
                    "qq" : check.user.qq,
                    'sex' : check.user.sex,
                    'age' : check.user.age,
                    'yearsOfWorking' : check.user.yearsOfWorking,
                    'school' : check.user.school,
                    'education' : check.user.education,
                    "addDate" :  check.user.addDate.strftime('%Y-%m-%d')
                    }
                })

    def put(self, request):
        check = CheckRequest(request);
        if not check.user:
            return JsonResponse({
                "status": "error",
                "msg": "User not logined"
            })
        uf = UpdateUserForm(check.jsonForm)
        if uf.is_valid():
            check.user.username = uf.cleaned_data['username']
            check.user.sex = uf.cleaned_data['sex']
            check.user.age = uf.cleaned_data['age']
            check.user.yearsOfWorking = uf.cleaned_data['yearsOfWorking']
            check.user.school = uf.cleaned_data['school']
            check.user.education = uf.cleaned_data['education']
            check.user.save()

            return JsonResponse({"status" : "success",
                    "msg" :  ''
                    })
        else:
            return JsonResponse({"status" : "error",
                    "msg" : 'Illegal put: %s' % uf.errors
                    })