#! /usr/bin/env python
# _*_ coding: utf-8 _*_
"""
author:
date:
file:       custom_ann.py
version:
doc:        自定义的神经网络模型
"""
import numpy as np
# from ray.rllib.algorithms.dreamer.dreamer import Conv2d, ConvEncoder, DreamerModel
from ray.rllib.algorithms.sac.rnnsac_torch_model import RNNSACTorchModel
from ray.rllib.algorithms.sac.sac_tf_model import SACTFModel
from ray.rllib.algorithms.sac.sac_torch_model import SACTorchModel
from ray.rllib.models.modelv2 import ModelV2
from ray.rllib.models.preprocessors import get_preprocessor
from ray.rllib.models.tf.recurrent_net import RecurrentNetwork
from ray.rllib.models.tf.tf_modelv2 import TFModelV2
from ray.rllib.models.torch.misc import SlimFC
from ray.rllib.models.torch.recurrent_net import RecurrentNetwork as TorchRNN
from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
from ray.rllib.utils.annotations import override
from ray.rllib.utils.framework import try_import_tf, try_import_torch

tf1, tf, tfv = try_import_tf()
torch, nn = try_import_torch()


def get_ann_model(input_shape, out_shape):
    inputs = tf.keras.layers.Input(input_shape)
    hid_layer1 = tf.keras.layers.Dense(
        units=64, activation=tf.nn.relu, name="fc1")(inputs)
    batch_normalization_layer1 = tf.keras.layers.BatchNormalization()(hid_layer1)
    hid_layer2 = tf.keras.layers.Dense(
        units=64, activation=tf.nn.relu, name="fc2")(batch_normalization_layer1)
    batch_normalization_layer2 = tf.keras.layers.BatchNormalization()(hid_layer2)
    hid_layer3 = tf.keras.layers.Dense(
        units=64, activation=tf.nn.tanh, name="fc3")(batch_normalization_layer2)
    batch_normalization_layer3 = tf.keras.layers.BatchNormalization()(hid_layer3)
    last_layer = tf.keras.layers.Dense(
        units=64, activation=tf.nn.tanh, name="fc4")(batch_normalization_layer3)
    # flatten = tf.keras.layers.Flatten()(last_layer)
    output = tf.keras.layers.Dense(units=out_shape[-1], activation=None, name="fc_out")(last_layer)
    # output = tf.reshape(output, [-1, out_shape[0]])
    # if len(out_shape) > 1:
    #     output = tf.reshape(output, [-1, *out_shape])
    model = tf.keras.models.Model(inputs, output)
    model.compile('adam', 'mse')
    return model


# class ConvEncoderCustom(ConvEncoder):
#     """Standard Convolutional Encoder for Dreamer. This encoder is used
#     to encode images frm an enviornment into a latent state for the
#     RSSM model in PlaNET.
#     """
#
#     def __init__(
#             self, depth: int = 32, act=None, shape=(1, 33, 1908)
#     ):
#         """Initializes Conv Encoder
#
#         Args:
#             depth (int): Number of channels in the first conv layer
#             act (Any): Activation for Encoder, default ReLU
#             shape (List): Shape of observation input
#         """
#         super().__init__()
#         self.act = act
#         if not act:
#             self.act = nn.ReLU
#         self.depth = depth
#         self.shape = shape
#
#         init_channels = self.shape[0]
#         self.layers = [
#             Conv2d(init_channels, self.depth, 4, stride=2),
#             self.act(),
#             Conv2d(self.depth, 2 * self.depth, 4, stride=2),
#             self.act(),
#             Conv2d(2 * self.depth, 4 * self.depth, 4, stride=2),
#             self.act(),
#             Conv2d(4 * self.depth, 8 * self.depth, 4, stride=2),
#             self.act(),
#         ]
#         self.model = nn.Sequential(*self.layers)


# class DreamerTorch(DreamerModel):
#     def __init__(self, obs_space, action_space, num_outputs, model_config, name):
#         super().__init__(obs_space, action_space, num_outputs, model_config, name)
#         self.encoder = ConvEncoderCustom(self.depth)
#
#     pass


class SACTF(SACTFModel):
    def __init__(
        self,
        obs_space,
        action_space,
        num_outputs,
        model_config,
        name: str,
        policy_model_config=None,
        q_model_config=None,
        twin_q: bool = False,
        initial_alpha: float = 1.0,
        target_entropy=None,
    ):
        super().__init__(obs_space, action_space, num_outputs,
                         model_config, name, policy_model_config, q_model_config, twin_q, initial_alpha, target_entropy)
        # print(f'num_outputs = {num_outputs}')


class RNNSACTorch(RNNSACTorchModel):
    def __init__(
        self,
        obs_space,
        action_space,
        num_outputs,
        model_config,
        name: str,
        policy_model_config=None,
        q_model_config=None,
        twin_q: bool = False,
        initial_alpha: float = 1.0,
        target_entropy=None,
    ):
        super().__init__(obs_space, action_space, num_outputs,
                         model_config, name, policy_model_config, q_model_config, twin_q, initial_alpha, target_entropy)
        # print(f'num_outputs = {num_outputs}')
        self.obs_size = get_preprocessor(obs_space)(obs_space).size

        self.fc_size = 64
        self.lstm_state_size = 256

        # Build the Module from fc + LSTM + 2xfc (action + value outs).
        self.fc1 = nn.Linear(self.obs_size, self.fc_size)
        self.lstm = nn.LSTM(
            self.fc_size, self.lstm_state_size, batch_first=True)
        self.value_branch = nn.Linear(self.lstm_state_size, 1)
        self.action_branch = nn.Linear(self.lstm_state_size, num_outputs)
        # Holds the current "base" output (before logits layer).
        self._features = None

    @override(ModelV2)
    def get_initial_state(self):
        # Place hidden states on same device as model.
        h = [
            self.fc1.weight.new(1, self.lstm_state_size).zero_().squeeze(0),
            self.fc1.weight.new(1, self.lstm_state_size).zero_().squeeze(0)
        ]
        return h

    @override(ModelV2)
    def value_function(self):
        assert self._features is not None, "must call forward() first"
        return torch.reshape(self.value_branch(self._features), [-1])

    @override(TorchRNN)
    def forward_rnn(self, inputs, state, seq_lens):
        """Feeds `inputs` (B x T x ..) through the Gru Unit.

        Returns the resulting outputs as a sequence (B x T x ...).
        Values are stored in self._cur_value in simple (B) shape (where B
        contains both the B and T dims!).

        Returns:
            NN Outputs (B x T x ...) as sequence.
            The state batches as a List of two items (c- and h-states).
        """
        x = nn.functional.relu(self.fc1(inputs))
        self._features, [h, c] = self.lstm(
            x, [torch.unsqueeze(state[0], 0),
                torch.unsqueeze(state[1], 0)])
        action_out = self.action_branch(self._features)
        return action_out, [torch.squeeze(h, 0), torch.squeeze(c, 0)]


class SACTorch(SACTorchModel):
    def __init__(
        self,
        obs_space,
        action_space,
        num_outputs,
        model_config,
        name: str,
        policy_model_config=None,
        q_model_config=None,
        twin_q: bool = False,
        initial_alpha: float = 1.0,
        target_entropy=None,
    ):
        super().__init__(obs_space, action_space, num_outputs,
                         model_config, name, policy_model_config, q_model_config, twin_q, initial_alpha, target_entropy)
        # print(f'num_outputs = {num_outputs}')
        self.obs_size = get_preprocessor(obs_space)(obs_space).size

        self.fc_size = 64
        self.lstm_state_size = 256

        # Build the Module from fc + LSTM + 2xfc (action + value outs).
        self.fc1 = nn.Linear(self.obs_size, self.fc_size)
        self.lstm = nn.LSTM(
            self.fc_size, self.lstm_state_size, batch_first=True)
        self.value_branch = nn.Linear(self.lstm_state_size, 1)
        self.action_branch = nn.Linear(self.lstm_state_size, out_features=80)
        # Holds the current "base" output (before logits layer).
        self._features = None

    @override(ModelV2)
    def get_initial_state(self):
        # Place hidden states on same device as model.
        h = [
            self.fc1.weight.new(1, self.lstm_state_size).zero_().squeeze(0),
            self.fc1.weight.new(1, self.lstm_state_size).zero_().squeeze(0)
        ]
        return h

    @override(ModelV2)
    def value_function(self):
        assert self._features is not None, "must call forward() first"
        return torch.reshape(self.value_branch(self._features), [-1])

    @override(TorchRNN)
    def forward_rnn(self, inputs, state, seq_lens):
        """Feeds `inputs` (B x T x ..) through the Gru Unit.

        Returns the resulting outputs as a sequence (B x T x ...).
        Values are stored in self._cur_value in simple (B) shape (where B
        contains both the B and T dims!).

        Returns:
            NN Outputs (B x T x ...) as sequence.
            The state batches as a List of two items (c- and h-states).
        """
        x = nn.functional.relu(self.fc1(inputs))
        self._features, [h, c] = self.lstm(
            x, [torch.unsqueeze(state[0], 0),
                torch.unsqueeze(state[1], 0)])
        action_out = self.action_branch(self._features)
        return action_out, [torch.squeeze(h, 0), torch.squeeze(c, 0)]
    # def get_initial_state(self):
    #     """Get the initial recurrent state values for the model.
    #
    #     Returns:
    #         List of np.array objects containing the initial hidden state
    #         of an RNN, if applicable.
    #
    #     Examples:
    #         >>> import numpy as np
    #         >>> from ray.rllib.models.modelv2 import ModelV2
    #         >>> class MyModel(ModelV2): # doctest: +SKIP
    #         ...     # ...
    #         ...     def get_initial_state(self):
    #         ...         return [
    #         ...             np.zeros(self.cell_size, np.float32),
    #         ...             np.zeros(self.cell_size, np.float32),
    #         ...         ]
    #     """
    #     cell_size = getattr(self, 'cell_size', 13)
    #     a = [
    #         np.zeros(cell_size, np.float32),
    #         np.zeros(cell_size, np.float32),
    #     ]
    #     return a


TF2_GLOBAL_SHARED_LAYER = None


class TF2SharedWeightsModel(TFModelV2):
    """Example of weight sharing between two different TFModelV2s.

    NOTE: This will only work for tf2.x. When running with config.framework=tf,
    use SharedWeightsModel1 and SharedWeightsModel2 below, instead!

    The shared (single) layer is simply defined outside of the two Models,
    then used by both Models in their forward pass.
    """

    def __init__(
        self, observation_space, action_space, num_outputs, model_config, name
    ):
        super().__init__(
            observation_space, action_space, num_outputs, model_config, name
        )

        global TF2_GLOBAL_SHARED_LAYER
        # The global, shared layer to be used by both models.
        if TF2_GLOBAL_SHARED_LAYER is None:
            TF2_GLOBAL_SHARED_LAYER = tf.keras.layers.Dense(
                units=64, activation=tf.nn.relu, name="fc1"
            )

        inputs = tf.keras.layers.Input(observation_space.shape)
        last_layer = TF2_GLOBAL_SHARED_LAYER(inputs)
        output = tf.keras.layers.Dense(
            units=num_outputs, activation=None, name="fc_out"
        )(last_layer)
        vf = tf.keras.layers.Dense(units=1, activation=None, name="value_out")(
            last_layer
        )
        self.base_model = tf.keras.models.Model(inputs, [output, vf])
        # self.base_model.summary()
        _ = None

    @override(ModelV2)
    def forward(self, input_dict, state, seq_lens):
        out, self._value_out = self.base_model(input_dict["obs"])
        return out, []

    @override(ModelV2)
    def value_function(self):
        return tf.reshape(self._value_out, [-1])


# TODO: 用户自定义网络模型 SharedWeightsModel1
class SharedWeightsModel1(TFModelV2):
    """两个不同的TensorFlow Model V2（TFModelV2）之间共享权重的例子.

    注意: 这里进在tf1(静态图)模式下工作. 当设置config.framework_str=tf2时, 应使用 TF2SharedWeightsModel.

    这里, 通过显式地使用tf1.AUTO_REUSE进入'shared'变量作用域从而实现变量的共享.
    这样做的结果是在一个名为'shared'的全局作用域中创建了 名为 'fc1', ... 层的变量,
    这个全局作用域位于策略（Policy）的正常变量作用域之外.
    """

    def __init__(
        self,
        observation_space,
        action_space,
        num_outputs,
        model_config,
        name,
        **kwargs
    ):
        super().__init__(
            observation_space, action_space, num_outputs, model_config, name
        )
        # print(f'num_outputs = {num_outputs}')
        inputs = tf.keras.layers.Input(observation_space.shape)
        with tf1.variable_scope(
            tf1.VariableScope(tf1.AUTO_REUSE, "shared"),
            reuse=tf1.AUTO_REUSE,
            auxiliary_name_scope=False
        ):
            hid_layer1 = tf.keras.layers.Dense(
                units=256, activation=tf.nn.relu, name="fc1",
                kernel_initializer='glorot_uniform'
            )(inputs)
            batch_normalization_layer1 = tf.keras.layers.BatchNormalization()(hid_layer1)
            hid_layer2 = tf.keras.layers.Dense(
                units=256, activation=tf.nn.relu, name="fc2",
                kernel_initializer='glorot_uniform'
            )(batch_normalization_layer1)
            batch_normalization_layer2 = tf.keras.layers.BatchNormalization()(hid_layer2)
            hid_layer3 = tf.keras.layers.Dense(
                units=256, activation=tf.nn.tanh, name="fc3",
                kernel_initializer='glorot_uniform'
            )(batch_normalization_layer2)
            batch_normalization_layer3 = tf.keras.layers.BatchNormalization()(hid_layer3)
            last_layer = tf.keras.layers.Dense(
                units=256, activation=tf.nn.tanh, name="fc4",
                kernel_initializer='glorot_uniform'
            )(batch_normalization_layer3)
        output = tf.keras.layers.Dense(
            units=num_outputs, activation=None, name="fc_out",
            kernel_initializer='glorot_uniform'
        )(last_layer)
        vf = tf.keras.layers.Dense(
            units=1, activation=None, name="value_out")(last_layer)
        self.base_model = tf.keras.models.Model(inputs, [output, vf])
        # print(f'已初始化: {action_space}, {num_outputs}')
        # _ = ''
        # self.register_variables(self.base_model.variables)

    @override(ModelV2)
    def forward(self, input_dict, state, seq_lens):
        out, self._value_out = self.base_model(input_dict["obs"])
        return out, []

    @override(ModelV2)
    def value_function(self):
        return tf.reshape(self._value_out, [-1])


# TODO: 用户自定义网络模型 SharedWeightsModel2
class SharedWeightsModel2(TFModelV2):
    """The "other" TFModelV2 using the same shared space as the one above."""

    def __init__(
        self, observation_space, action_space, num_outputs,
        model_config, name
    ):
        super().__init__(observation_space, action_space, num_outputs,
                         model_config, name)

        inputs = tf.keras.layers.Input(observation_space.shape)

        # Weights shared with SharedWeightsModel1.
        with tf1.variable_scope(
            tf1.VariableScope(tf1.AUTO_REUSE, "shared"),
            reuse=tf1.AUTO_REUSE,
            auxiliary_name_scope=False):
            hid_layer1 = tf.keras.layers.Dense(
                units=128, activation=tf.nn.relu, name="fc1")(inputs)
            hid_layer2 = tf.keras.layers.Dense(
                units=128, activation=tf.nn.tanh, name="fc2")(hid_layer1)
            last_layer = tf.keras.layers.Dense(
                units=128, activation=tf.nn.tanh, name="fc3")(hid_layer2)
        output = tf.keras.layers.Dense(
            units=num_outputs, activation=None, name="fc_out")(last_layer)
        vf = tf.keras.layers.Dense(
            units=1, activation=None, name="value_out")(last_layer)
        self.base_model = tf.keras.models.Model(inputs, [output, vf])
        # self.register_variables(self.base_model.variables)

    @override(ModelV2)
    def forward(self, input_dict, state, seq_lens):
        out, self._value_out = self.base_model(input_dict["obs"])
        return out, []

    @override(ModelV2)
    def value_function(self):
        return tf.reshape(self._value_out, [-1])


TORCH_GLOBAL_SHARED_LAYER: SlimFC or None = None
if torch:
    TORCH_GLOBAL_SHARED_LAYER = SlimFC(
        64,
        64,
        activation_fn=nn.ReLU,
        initializer=torch.nn.init.xavier_uniform_,
    )


# TODO: 用户自定义网络模型 SharedWeightsModelTorch
class SharedWeightsModelTorch(TorchModelV2, nn.Module):
    """Example of weight sharing between two different TorchModelV2s.

    The shared (single) layer is simply defined outside of the two Models,
    then used by both Models in their forward pass.
    """

    def __init__(
        self, observation_space, action_space, num_outputs,
        model_config, name
    ):
        TorchModelV2.__init__(self, observation_space, action_space,
                              num_outputs, model_config, name)
        nn.Module.__init__(self)

        # Non-shared initial layer.
        self.first_layer = SlimFC(
            int(np.prod(observation_space.shape)),
            64,
            activation_fn=nn.ReLU,
            initializer=torch.nn.init.xavier_uniform_)

        # Non-shared final layer.
        self.last_layer = SlimFC(
            64,
            self.num_outputs,
            activation_fn=None,
            initializer=torch.nn.init.xavier_uniform_)
        self.vf = SlimFC(
            64,
            1,
            activation_fn=None,
            initializer=torch.nn.init.xavier_uniform_,
        )
        self._output = None

    @override(ModelV2)
    def forward(self, input_dict, state, seq_lens):
        out = self.first_layer(input_dict["obs"])
        self._output = TORCH_GLOBAL_SHARED_LAYER(out)
        model_out = self.last_layer(self._output)
        return model_out, []

    @override(ModelV2)
    def value_function(self):
        assert self._output is not None, "must call forward first!"
        return torch.reshape(self.vf(self._output), [-1])


# TODO: 用户自定义网络模型 RNNModelTF
class RNNModelTF(RecurrentNetwork):
    """Example of using the Keras functional API to define a RNN model."""

    def __init__(
        self,
        obs_space,
        action_space,
        num_outputs,
        model_config,
        name,
        hiddens_size=256,
        cell_size=64
    ):
        super(RNNModelTF, self).__init__(obs_space, action_space, num_outputs,
                                         model_config, name)
        self.cell_size = cell_size

        # Define input layers
        input_layer = tf.keras.layers.Input(
            shape=(None, obs_space.shape[0]), name="inputs")
        state_in_h = tf.keras.layers.Input(shape=(cell_size,), name="h")
        state_in_c = tf.keras.layers.Input(shape=(cell_size,), name="c")
        seq_in = tf.keras.layers.Input(shape=(), name="seq_in", dtype=tf.int32)

        # Preprocess observation with a hidden layer and send to LSTM cell
        dense1 = tf.keras.layers.Dense(
            hiddens_size, activation=tf.nn.relu, name="dense1")(input_layer)
        lstm_out, state_h, state_c = tf.keras.layers.LSTM(
            cell_size, return_sequences=True, return_state=True, name="lstm")(
            inputs=dense1,
            mask=tf.sequence_mask(seq_in),
            initial_state=[state_in_h, state_in_c])

        # Postprocess LSTM output with another hidden layer and compute values
        logits = tf.keras.layers.Dense(
            self.num_outputs,
            activation=tf.keras.activations.linear,
            name="logits")(lstm_out)
        values = tf.keras.layers.Dense(
            1, activation=None, name="values")(lstm_out)

        # Create the RNN model
        self.rnn_model = tf.keras.Model(
            inputs=[input_layer, seq_in, state_in_h, state_in_c],
            outputs=[logits, values, state_h, state_c])
        self.register_variables(self.rnn_model.variables)
        self.rnn_model.summary()

    @override(RecurrentNetwork)
    def forward_rnn(self, inputs, state, seq_lens):
        print("inputs", inputs)
        print("state", state)
        print("seq_lens", seq_lens)
        model_out, self._value_out, h, c = self.rnn_model([inputs, seq_lens] +
                                                          state)
        return model_out, [h, c]

    @override(ModelV2)
    def get_initial_state(self):
        return [
            np.zeros(self.cell_size, np.float32),
            np.zeros(self.cell_size, np.float32),
        ]

    # @override(ModelV2)
    # def value_function(self):
    #     return tf.reshape(self._value_out, [-1])


# TODO: 用户自定义网络模型 RNNModelTorch
class RNNModelTorch(TorchRNN, nn.Module):
    def __init__(
        self,
        obs_space,
        action_space,
        num_outputs,
        model_config,
        name,
        fc_size=64,
        lstm_state_size=256
    ):
        nn.Module.__init__(self)
        super().__init__(obs_space, action_space, num_outputs, model_config,
                         name)

        self.obs_size = get_preprocessor(obs_space)(obs_space).size
        self.fc_size = fc_size
        self.lstm_state_size = lstm_state_size

        # Build the Module from fc + LSTM + 2xfc (action + value outs).
        self.fc1 = nn.Linear(self.obs_size, self.fc_size)
        self.lstm = nn.LSTM(
            self.fc_size, self.lstm_state_size, batch_first=True)
        self.action_branch = nn.Linear(self.lstm_state_size, num_outputs)
        self.value_branch = nn.Linear(self.lstm_state_size, 1)
        # Holds the current "base" output (before logits layer).
        self._features = None

    @override(ModelV2)
    def get_initial_state(self):
        # Place hidden states on same device as model.
        h = [
            self.fc1.weight.new(1, self.lstm_state_size).zero_().squeeze(0),
            self.fc1.weight.new(1, self.lstm_state_size).zero_().squeeze(0)
        ]
        return h

    @override(ModelV2)
    def value_function(self):
        assert self._features is not None, "must call forward() first"
        return torch.reshape(self.value_branch(self._features), [-1])

    @override(TorchRNN)
    def forward_rnn(self, inputs, state, seq_lens):
        """Feeds `inputs` (B x T x ..) through the Gru Unit.

        Returns the resulting outputs as a sequence (B x T x ...).
        Values are stored in self._cur_value in simple (B) shape (where B
        contains both the B and T dims!).

        Returns:
            NN Outputs (B x T x ...) as sequence.
            The state batches as a List of two items (c- and h-states).
        """
        x = nn.functional.relu(self.fc1(inputs))
        self._features, [h, c] = self.lstm(
            x, [torch.unsqueeze(state[0], 0),
                torch.unsqueeze(state[1], 0)])
        action_out = self.action_branch(self._features)
        return action_out, [torch.squeeze(h, 0), torch.squeeze(c, 0)]


class TestTF1Model:

    def __init__(
        self, observation_space, num_outputs
    ):
        inputs = tf.keras.layers.Input(observation_space.shape)
        with tf1.variable_scope(
            tf1.VariableScope(tf1.AUTO_REUSE, "shared"),
            reuse=tf1.AUTO_REUSE,
            auxiliary_name_scope=False,
        ):
            last_layer = tf.keras.layers.Dense(
                units=64, activation=tf.nn.relu, name="fc1"
            )(inputs)
        output = tf.keras.layers.Dense(
            units=num_outputs, activation=None, name="fc_out"
        )(last_layer)
        vf = tf.keras.layers.Dense(units=1, activation=None, name="value_out")(
            last_layer
        )
        self.base_model = tf.keras.models.Model(inputs, [output, vf, last_layer])

    def __call__(self, array):
        out, self._value_out, last_layer = self.base_model(array)
        return out, last_layer


shared_layer = None


class TestTF2Model:

    def __init__(
        self, observation_space, num_outputs
    ):
        global shared_layer
        # The global, shared layer to be used by both models.
        if shared_layer is None:
            shared_layer = tf.keras.layers.Dense(
                units=64, activation=tf.nn.relu, name="fc1"
            )

        inputs = tf.keras.layers.Input(observation_space.shape)
        last_layer = shared_layer(inputs)
        output = tf.keras.layers.Dense(
            units=num_outputs, activation=None, name="fc_out"
        )(last_layer)
        vf = tf.keras.layers.Dense(units=1, activation=None, name="value_out")(
            last_layer
        )
        self.base_model = tf.keras.models.Model(inputs, [output, vf, last_layer])

    def __call__(self, array):
        out, self._value_out, last_layer = self.base_model(array)
        return out, last_layer


if __name__ == '__main__':
    pass
    # import gymnasium as gym
    #
    # observation_space = gym.spaces.Box(0, 10, (37,))
    #
    # # model = TestTF1Model(observation_space, 7)
    # # model1 = TestTF1Model(observation_space, 7)
    # #
    # # x = np.random.random((73, 37))
    # # out, last_layer = model(x)
    # # out1, last_layer1 = model1(x)
    #
    # model = TestTF2Model(observation_space, 7)
    # model1 = TestTF2Model(observation_space, 7)
    #
    # x = np.random.random((73, 37))
    # out, last_layer = model(x)
    # out1, last_layer1 = model1(x)
