"""
RNN 文本生成模型构建组件。

这个文件放诗歌 RNN 任务使用的训练模型和推理模型装配逻辑，由文本生成 Pipeline 调用。
"""

from dataclasses import dataclass
from functools import partial
from pathlib import Path

import keras
import tensorflow as tf
from keras import layers

from deep_learning.pipeline.impl.text_generation.generation import generate_with_stateful_model, generate_with_training_model
from deep_learning.models.spec import TextGenerationModel, TextGenerationModelBuilder


@dataclass
class RNNModelBuilder(TextGenerationModelBuilder):
    num_layers: int = 2
    embedding_dim: int = 100
    hidden_dim: int = 1024

    def build_training_artifact(
        self,
        vocab_size: int,
        sequence_length: int
    ) -> TextGenerationModel:
        inputs = keras.Input(shape=(None,), dtype="int32", name="inputs")
        x = layers.Embedding(
            input_dim=vocab_size,
            output_dim=self.embedding_dim,
            mask_zero=True,
            name="embedding"
        )(inputs)

        for i in range(self.num_layers):
            x = layers.LSTM(
                self.hidden_dim,
                return_sequences=True,
                recurrent_dropout=0.1,
                name=f"lstm_{i}"
            )(x)
            x = layers.Dropout(0.1, name=f"dropout_{i}")(x)

        outputs = layers.Dense(vocab_size, name="logits")(x)
        model = keras.Model(inputs=inputs, outputs=outputs, name="rnn_training")
        return TextGenerationModel(
            model=model,
            generate=partial(generate_with_training_model, model)
        )

    def compile_training_model(self, model: keras.Model) -> None:
        from deep_learning.pipeline.impl.text_generation.pipeline import WarmupSchedule

        model.compile(
            optimizer=keras.optimizers.Adam(WarmupSchedule()),
            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
            metrics=["accuracy"]
        )

    def convert_to_inference_artifact(
        self,
        training_artifact: TextGenerationModel
    ) -> TextGenerationModel:
        inference_model = self._build_inference_model_from_training_model(
            training_artifact.model
        )
        return TextGenerationModel(
            model=inference_model,
            generate=partial(
                generate_with_stateful_model,
                inference_model,
                initial_states=self._initial_states(batch_size=1)
            )
        )

    def load_inference_artifact(self, model_path: Path) -> TextGenerationModel:
        model = keras.models.load_model(str(model_path))
        initial_states = [
            tf.zeros((1, state_input.shape[-1]))
            for state_input in model.inputs[1:]
        ]
        return TextGenerationModel(
            model=model,
            generate=partial(
                generate_with_stateful_model,
                model,
                initial_states=initial_states
            )
        )

    def _build_inference_model_from_training_model(
        self,
        training_model: keras.Model
    ) -> keras.Model:
        token_input = keras.Input(shape=(None,), dtype="int32", name="token_input")
        state_inputs = []
        for i in range(self.num_layers):
            h_input = keras.Input(shape=(self.hidden_dim,), name=f"h_{i}_input")
            c_input = keras.Input(shape=(self.hidden_dim,), name=f"c_{i}_input")
            state_inputs.extend([h_input, c_input])

        embedding = training_model.get_layer("embedding")
        logits_layer = training_model.get_layer("logits")
        x = embedding(token_input)

        new_states = []
        inference_lstm_layers = []
        for i in range(self.num_layers):
            inference_lstm = layers.LSTM(
                self.hidden_dim,
                return_sequences=i < self.num_layers - 1,
                return_state=True,
                recurrent_dropout=0.1,
                name=f"lstm_{i}"
            )
            h_input = state_inputs[i * 2]
            c_input = state_inputs[i * 2 + 1]
            x, new_h, new_c = inference_lstm(x, initial_state=[h_input, c_input])
            new_states.extend([new_h, new_c])
            dropout = training_model.get_layer(f"dropout_{i}")
            x = dropout(x)
            inference_lstm_layers.append(inference_lstm)

        logits = logits_layer(x)
        inference_model = keras.Model(
            [token_input] + state_inputs,
            [logits] + new_states,
            name="rnn_inference"
        )

        for i, inference_lstm in enumerate(inference_lstm_layers):
            training_lstm = training_model.get_layer(f"lstm_{i}")
            inference_lstm.set_weights(training_lstm.get_weights())

        return inference_model

    def _initial_states(self, batch_size: int) -> list:
        states = []
        for _ in range(self.num_layers):
            states.append(tf.zeros((batch_size, self.hidden_dim)))
            states.append(tf.zeros((batch_size, self.hidden_dim)))
        return states