import torch
from torch import nn
from dsx.util import print_red
from dsx.transformer.layers import SingleHeadAttention, LayerNorm, MLP

from dsx.transformer.layers import create_causal_attn_mask


class EncoderLayer(nn.Module):

    def __init__(self, d_model, feedforward_size):
        super(EncoderLayer, self).__init__()

        self.attn = SingleHeadAttention(d_model)  # 自注意力层
        self.attn_norm = LayerNorm(d_model)  # 自注意力的归一化层
        self.mlp = MLP(d_model, feedforward_size)  # 前馈网络层
        self.mlp_norm = LayerNorm(d_model)  # 前馈网络的归一化层

    def forward(self, hidden_states: torch.Tensor, mask=None):
        # 对输入进行自注意力计算，并对结果进行残差连接和归一化
        residual = hidden_states
        hidden_states, _ = self.attn(
            q_input=hidden_states,
            k_input=hidden_states,
            v_input=hidden_states,
            attn_mask=mask,
        )
        hidden_states = self.attn_norm(hidden_states + residual)

        # 对输入进行前馈网络计算，并对结果进行残差连接和归一化
        residual = hidden_states
        hidden_states = self.mlp(hidden_states)
        hidden_states = self.mlp_norm(hidden_states + residual)

        return hidden_states


from torch import nn

from dsx.transformer.layers import PositionalEncoding, create_cross_attn_pad_mask


class Encoder(nn.Module):
    def __init__(
        self, vocab_size, max_seq_len, d_model, feedforward_size, n_layers, pad_token_id
    ):
        super(Encoder, self).__init__()

        self.pad_token_id = pad_token_id
        self.token_emb = nn.Embedding(vocab_size, d_model)  # 词嵌入层
        self.position_emb = PositionalEncoding(d_model, max_seq_len)  # 位置编码层
        self.layers = nn.ModuleList(
            [EncoderLayer(d_model, feedforward_size) for _ in range(n_layers)]
        )  # 多层encoder

    def forward(self, encoder_input_ids):
        # encoder_inputs: [N, T]
        encoder_outputs = self.token_emb(
            encoder_input_ids
        )  # encoder_outputs: [N, T, D]
        encoder_outputs = self.position_emb(
            encoder_outputs
        )  # encoder_outputs: [N, T, D]

        # # 生成Query和Key的交叉注意力填充mask矩阵，每个Query只关注非填充的Key，只与其计算注意力得分: [N, T, T]
        encoder_attn_pad_mask = create_cross_attn_pad_mask(
            query_ids=encoder_input_ids,
            key_ids=encoder_input_ids,
            pad_token_id=self.pad_token_id,
        )

        # print("encoder_attn_pad_mask: \n", encoder_attn_pad_mask, "\n")

        for layer in self.layers:
            encoder_outputs = layer(encoder_outputs, encoder_attn_pad_mask)

        # 数据的形状为 [批次大小, 序列长度, 模型维度]
        return encoder_outputs


class DecoderLayer(nn.Module):
    def __init__(self, d_model, feedforward_size):
        super(DecoderLayer, self).__init__()
        self.self_attn = SingleHeadAttention(d_model)  # 自注意力层
        self.attn_norm = LayerNorm(d_model)

        self.cross_attn = SingleHeadAttention(d_model)  # 交叉注意力层
        self.cross_attn_norm = LayerNorm(d_model)

        self.mlp = MLP(d_model, feedforward_size)  # 前馈网络层
        self.mlp_norm = LayerNorm(d_model)

    def forward(self, dec_inputs, enc_outputs, dec_self_attn_mask, dec_enc_attn_mask):
        # 残差
        residual = dec_inputs

        # 自注意力计算
        dec_outputs, _ = self.self_attn(
            dec_inputs, dec_inputs, dec_inputs, dec_self_attn_mask
        )
        # 残差连接 + 层归一化
        dec_outputs = self.attn_norm(dec_outputs + residual)

        # 残差
        residual = dec_outputs
        # 交叉注意力计算
        dec_outputs, _ = self.cross_attn(
            dec_outputs, enc_outputs, enc_outputs, dec_enc_attn_mask
        )

        # 残差连接 + 层归一化
        dec_outputs = self.cross_attn_norm(dec_outputs + residual)

        # 残差
        residual = dec_outputs
        # 前馈网络计算
        dec_outputs = self.mlp(dec_outputs)
        # 残差连接 + 层归一化
        dec_outputs = self.mlp_norm(dec_outputs + residual)

        return dec_outputs


class Decoder(nn.Module):
    def __init__(
        self, vocab_size, max_seq_len, d_model, feedforward_size, n_layers, pad_token_id
    ):
        super(Decoder, self).__init__()
        # 解码器的vocab_size可以和编码器是不同的（例如：翻译任务）
        self.pad_token_id = pad_token_id
        self.token_emb = nn.Embedding(vocab_size, d_model)
        self.position_emb = PositionalEncoding(d_model, max_seq_len)
        self.layers = nn.ModuleList(
            [DecoderLayer(d_model, feedforward_size) for _ in range(n_layers)]
        )
        self.output = nn.Linear(d_model, vocab_size)

    def forward(self, dec_inputs, enc_inputs, enc_outputs):
        dec_outputs = self.token_emb(dec_inputs)  # [batch_size, tgt_len, d_model]
        dec_outputs = self.position_emb(dec_outputs)  # [batch_size, tgt_len, d_model]

        # 生成解码器输入的填充掩码矩阵
        dec_self_attn_pad_mask = create_cross_attn_pad_mask(
            dec_inputs, dec_inputs, self.pad_token_id
        ).to(dec_inputs.device)
        # print("dec_self_attn_pad_mask ->\n", dec_self_attn_pad_mask, "\n")

        # 生成解码器输入的causal mask矩阵
        dec_self_attn_subsequence_mask = create_causal_attn_mask(dec_inputs).to(
            dec_inputs.device
        )
        # print("dec_self_attn_subsequence_mask ->\n", dec_self_attn_subsequence_mask, "\n")

        # 通过逻辑与运算，生成解码器输入最终的掩码矩阵
        dec_self_attn_mask = (
            torch.logical_and(dec_self_attn_pad_mask, dec_self_attn_subsequence_mask)
            .byte()
            .to(dec_inputs.device)
        )
        # print("dec_self_attn_mask ->\n", dec_self_attn_mask, "\n")

        # 生成编码器和解码器的交叉注意力矩阵
        dec_enc_attn_mask = create_cross_attn_pad_mask(
            dec_inputs, enc_inputs, self.pad_token_id
        ).to(dec_inputs.device)
        # print("dec_enc_attn_mask ->\n", dec_enc_attn_mask, "\n")

        for layer in self.layers:
            dec_outputs = layer(
                dec_outputs, enc_outputs, dec_self_attn_mask, dec_enc_attn_mask
            )

        dec_outputs = self.output(dec_outputs)

        return dec_outputs


from dsx.transformer.tokenizer import TransformerTokenizer
import lightning as L


class Transformer(L.LightningModule):
    def __init__(
        self,
        tokenizer: TransformerTokenizer,
        d_model,
        feedforward_size,
        encoder_max_length,
        decoder_max_length,
        n_layers=1,
        learning_rate=0.0001,
    ):
        super(Transformer, self).__init__()

        self.encoder = Encoder(
            vocab_size=tokenizer.vocab_size,
            max_seq_len=encoder_max_length,
            d_model=d_model,
            feedforward_size=feedforward_size,
            n_layers=n_layers,
            pad_token_id=tokenizer.pad_token_id,
        )

        self.decoder = Decoder(
            vocab_size=tokenizer.vocab_size,
            max_seq_len=decoder_max_length,
            d_model=d_model,
            feedforward_size=feedforward_size,
            n_layers=n_layers,
            pad_token_id=tokenizer.pad_token_id,
        )

        self.tokenizer = tokenizer
        self.learning_rate = learning_rate
        self.encoder_max_length = encoder_max_length
        self.decoder_max_length = decoder_max_length

        # 示例输入
        self.example_input_array = (
            torch.randint(
                0, tokenizer.vocab_size, (2, encoder_max_length), dtype=torch.long
            ),
            torch.randint(
                0, tokenizer.vocab_size, (2, decoder_max_length), dtype=torch.long
            ),
        )

        # 用于存储训练和验证过程中的指标
        self.train_step_losses = []
        self.train_epoch_losses = []
        self.validation_step_outputs = []
        self.validation_epoch_outputs = []

    def forward(self, encoder_input_ids, decoder_input_ids):
        # 前向传播
        encoder_outputs = self.encoder(encoder_input_ids)
        decoder_outputs = self.decoder(
            decoder_input_ids, encoder_input_ids, encoder_outputs
        )

        return decoder_outputs

    def training_step(self, batch, batch_idx):
        encoder_input_ids, decoder_input_ids = (
            batch["encoder_input_ids"],
            batch["decoder_input_ids"],
        )
        decoder_outputs = self(encoder_input_ids, decoder_input_ids)

        decoder_target_ids = batch["decoder_target_ids"]
        loss = torch.nn.functional.cross_entropy(
            decoder_outputs.view(-1, decoder_outputs.shape[-1]),
            decoder_target_ids.view(-1),
        )

        # 存储损失以便后续使用
        self.train_step_losses.append(loss.detach())

        return loss

    def on_train_epoch_end(self):
        """在每个训练epoch结束时计算整体损失"""

        if self.train_step_losses:  # 确保列表不为空
            # 计算并记录平均训练损失
            avg_train_loss = torch.stack(self.train_step_losses).mean()

            print(
                f"***** 【Epoch {self.current_epoch}】  Train Avg Loss: {avg_train_loss:.4f} *****"
            )

            self.train_epoch_losses.append(
                {"epoch": self.current_epoch, "loss": avg_train_loss.item()}
            )

            # 清空列表为下一个 epoch 做准备
            self.train_step_losses.clear()

    def validation_step(self, batch, batch_idx):
        encoder_input_ids = batch["encoder_input_ids"]
        decoder_input_ids = batch["decoder_input_ids"]
        decoder_target_ids = batch["decoder_target_ids"]

        # 前向传播
        outputs = self(encoder_input_ids, decoder_input_ids)

        # 计算预测值
        preds = torch.argmax(outputs, dim=-1)

        # 创建掩码以忽略填充值
        mask = decoder_target_ids != self.tokenizer.pad_token_id

        # 计算token级别的准确率
        correct_tokens = ((preds == decoder_target_ids) & mask).sum().float()
        total_tokens = mask.sum().float()
        token_acc = (
            correct_tokens / total_tokens if total_tokens > 0 else torch.tensor(0.0)
        )

        # 计算序列级别的准确率
        seq_correct = ((preds == decoder_target_ids) | ~mask).all(dim=1).float().mean()

        # 保存结果供epoch结束时使用
        self.validation_step_outputs.append(
            {
                "preds": preds,
                "target_ids": decoder_target_ids,
                "token_acc": token_acc,
                "seq_acc": seq_correct,
            }
        )

    def on_validation_epoch_end(self):
        """在每个验证epoch结束时计算整体准确率"""
        # 汇总所有预测结果和标签
        all_preds = torch.cat([x["preds"] for x in self.validation_step_outputs])
        all_target_ids = torch.cat(
            [x["target_ids"] for x in self.validation_step_outputs]
        )

        # 1. 计算序列（样本）级别的准确率
        # 创建掩码以标识非填充位置
        mask = all_target_ids != self.tokenizer.pad_token_id
        # 计算每个序列是否完全正确（所有非填充位置都预测正确）
        seq_correct = ((all_preds == all_target_ids) | ~mask).all(dim=1)
        # 序列总数
        total_sequences = all_preds.size(0)
        # 正确预测的序列数
        correct_sequences = seq_correct.sum().item()
        # 序列级别准确率
        seq_acc = correct_sequences / total_sequences if total_sequences > 0 else 0.0

        print_red(
            f"***** Validation: 样本总数 {total_sequences}  正确预测: {correct_sequences}  正确率: {seq_acc:.4f} *****"
        )

        # 2. 计算 Token 级别的准确率
        # 计算所有非填充位置的预测总数
        total_tokens = mask.sum().item()
        # 计算所有非填充位置预测正确的总数
        correct_tokens = ((all_preds == all_target_ids) & mask).sum().item()
        # Token 级别准确率
        token_acc = correct_tokens / total_tokens if total_tokens > 0 else 0.0

        # 3. 将评估结果保存到 validation_epoch_outputs 列表中
        self.validation_epoch_outputs.append(
            {
                "epoch": self.current_epoch,
                "总样本数": total_sequences,
                "正确样本数": correct_sequences,
                "样本准确率": round(seq_acc, 4),
                "总Token数": total_tokens,
                "正确Token数": correct_tokens,
                "Token准确率": round(token_acc, 4),
            }
        )

        # 记录日志
        self.log("total_sequences", total_sequences)
        self.log("correct_sequences", correct_sequences)
        self.log("seq_acc", seq_acc)

        self.log("total_tokens", total_tokens)
        self.log("correct_tokens", correct_tokens)
        self.log("token_acc", token_acc)

        # 清空缓存
        self.validation_step_outputs.clear()

    def clear_cache(self):
        """清除缓存"""
        self.train_step_losses.clear()
        self.train_epoch_losses.clear()
        self.validation_step_outputs.clear()
        self.validation_epoch_outputs.clear()

    def configure_optimizers(self):
        """配置优化器"""
        optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
        return optimizer

    def generate(self, encoder_input_ids):
        encoder_input_ids = encoder_input_ids.unsqueeze(0)
        enc_outputs = self.encoder(encoder_input_ids)
        dec_input = torch.full(
            (1, self.decoder_max_length),
            self.tokenizer.pad_token_id,
            device=encoder_input_ids.device,
            dtype=torch.long,
        )
        next_token = self.tokenizer.bos_token_id

        for i in range(self.decoder_max_length):
            dec_input[0][i] = next_token

            projected = self.decoder(dec_input, encoder_input_ids, enc_outputs)
            prob = projected.squeeze(0).max(dim=-1, keepdim=False)[1]
            next_word = prob.data[i]
            next_token = next_word.item()

        # 解码
        preds = self.tokenizer.decode(dec_input[0].tolist())

        return preds

    def generate_batch(self, encoder_input_ids):
        batch_size = encoder_input_ids.shape[0]

        # 编码器前向传播
        encoder_outputs = self.encoder(encoder_input_ids)

        # 初始化解码器输入，使用开始标记填充
        decoder_input_ids = torch.full(
            (batch_size, self.decoder_max_length),
            self.tokenizer.pad_token_id,
            device=encoder_outputs.device,
            dtype=torch.long,
        )

        # 设置第一个位置为开始标记
        decoder_input_ids[:, 0] = self.tokenizer.bos_token_id

        # 逐步生成每个位置的输出
        for i in range(
            1, self.decoder_max_length
        ):  # 从位置1开始，因为位置0已经设置为BOS
            # 获取当前已生成序列的预测结果
            decoder_outputs = self.decoder(
                dec_inputs=decoder_input_ids,
                enc_inputs=encoder_input_ids,
                enc_outputs=encoder_outputs,
            )

            # 获取当前位置的预测概率分布
            current_logits = decoder_outputs[
                :, i - 1, :
            ]  # [batch_size, vocab_size]，使用前一个位置的输出预测当前位置

            # 选择概率最高的token作为预测结果
            next_tokens = torch.argmax(current_logits, dim=-1)  # [batch_size]

            # 更新解码器输入，将当前位置的预测结果作为已知信息
            decoder_input_ids[:, i] = next_tokens

            # 检查是否所有序列都生成了结束标记，如果是则提前结束
            if torch.all(next_tokens == self.tokenizer.eos_token_id):
                break

        preds = []
        for i in range(batch_size):
            # 解码预测结果，并去除结束标记后的部分
            decoded = self.tokenizer.decode(decoder_input_ids[i].tolist())
            # 如果生成了结束标记，只取到结束标记之前的部分
            if self.tokenizer.eos_token in decoded:
                decoded = decoded.split(self.tokenizer.eos_token)[0]
            preds.append(decoded)

        return preds
