# @Author：青松
import platform
import random

from IPython import display
import numpy as np
import pandas as pd
import torch
import lightning as L
from colorama import Fore, Style
from prettytable import PrettyTable
import matplotlib.pyplot as plt


def set_seed(seed=1024):
    # 设置随机种子
    random.seed(seed)

    # 设置 NumPy 的随机种子
    np.random.seed(seed)

    # 设置 PyTorch 的随机种子
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    # 启用 PyTorch Lightning 的种子管理
    L.seed_everything(seed)

    # 设置 MPS 随机种子（如果可用）
    if torch.backends.mps.is_available():
        torch.mps.manual_seed(seed)

    # 确保结果的可重复性
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


def show_version():
    print_table(
        "本书愿景",
        ["Info", "《动手学大语言模型》"],
        [
            ["作者", "吾辈亦有感"],
            ["哔站", "https://space.bilibili.com/3546632320715420"],
            ["定位", "基于'从零构建'的理念，用实战帮助程序员快速入门大模型。"],
            ["愿景", "若让你的AI学习之路走的更容易一点，我将倍感荣幸！祝好😄"],
        ],
    )

    print_table(
        "环境信息",
        ["Python 版本", "PyTorch 版本", "PyTorch Lightning 版本"],
        [[platform.python_version(), torch.__version__, L.__version__]],
    )

    set_seed(1024)


def show_table(data, headers=None):
    # 设置浮点数显示格式，避免使用科学记数法影响阅读体验
    pd.set_option("display.float_format", "{:.4f}".format)

    # 如果提供了headers，则使用它们作为列名
    if headers is not None:
        if isinstance(data, list):
            # 假设data是列表的列表，第一行为数据而非标题
            df = pd.DataFrame(data, columns=headers)
        else:
            # 如果data已经是字典或其他结构，headers将被忽略
            df = pd.DataFrame(data, columns=headers)
    else:
        df = pd.DataFrame(data)

    return df


def print_table(table_name, field_names: list, data: list, show_headers=True):
    table = PrettyTable()
    if show_headers:
        table.field_names = field_names
    for data_item in data:
        table.add_row(data_item)
    if table_name:
        print_red(f"{table_name}：")
    print(table)


def print_red(text):
    print(Fore.RED + text + Style.RESET_ALL)


def print_regression_results(
    inputs,
    true_values,
    predictions,
    table_name="预测结果",
    field_names=["输入", "真实值", "预测值", "预测偏差", "预测精度"],
):
    """
    打印回归预测的结果
    """

    # 处理 tensor 类型的预测结果
    if isinstance(predictions, torch.Tensor):
        predictions = predictions.tolist()

    # 准备表格数据
    table_data = []

    for i in range(len(inputs)):
        table_data.append(
            [
                inputs[i],
                round(true_values[i], 4),
                round(predictions[i], 4),
                f"{predictions[i] - true_values[i]:.4f}",
                f"{predictions[i] * 100 / true_values[i]:.2f}%",
            ]
        )

    print_table(table_name, field_names, table_data)


def print_classification_results(
    inputs, true_labels, predictions, probabilities, label_map
):
    """
    打印文本分类预测结果，包含预测标签和真实标签的对比

    Args:
        inputs: 待预测的样本列表
        true_labels: 真实标签列表（数字形式）
        predictions: 模型预测结果（数字形式）
        probabilities: 预测概率列表
        label_map: 标签映射字典
    """

    # 处理 tensor 类型的预测结果
    if isinstance(predictions, torch.Tensor):
        predictions = predictions.tolist()
    if isinstance(probabilities, torch.Tensor):
        probabilities = probabilities.tolist()

    # 准备表格数据
    table_data = []
    correct_count = 0

    for i, text in enumerate(inputs):
        # 获取真实标签和预测标签名称
        true_label_name = label_map[true_labels[i]]
        pred_label_name = label_map[predictions[i]]

        # 判断预测是否正确
        is_correct = true_labels[i] == predictions[i]
        if is_correct:
            correct_count += 1
            mark = "\033[92m" + "☑" + "\033[0m"  # 绿色勾号表示正确
        else:
            mark = "\033[91m" + "☒" + "\033[0m"  # 红色叉号表示错误

        # 处理长文本显示
        display_text = text[:30] + "..." if len(text) > 30 else text

        # 获取最高概率
        if isinstance(probabilities[i], (list, tuple)):
            max_prob = max(probabilities[i])
        else:
            max_prob = probabilities[i].max().item()

        # 添加行数据
        table_data.append(
            [display_text, true_label_name, pred_label_name, f"{max_prob:.4f}", mark]
        )

    # 使用 print_table 方法打印结果
    accuracy = correct_count / len(inputs) if len(inputs) > 0 else 0
    print_table(
        f"🎯 分类预测结果 (准确率: {correct_count}/{len(inputs)} = {accuracy * 100:.2f}%)",
        ["输入", "真实标签", "预测标签", "最高概率", "标记"],
        table_data,
    )


def print_aigc_results(questions, real_answers, pred_answers):
    # 准备表格数据
    table_data = []
    correct_count = 0
    for i in range(len(questions)):
        # 获取真实答案和预测答案名称
        real_answer_name = real_answers[i]
        pred_answer_name = pred_answers[i]

        # 判断预测是否正确
        if real_answer_name == pred_answer_name:
            correct_count += 1
            mark = "\033[92m" + "☑" + "\033[0m"  # 绿色勾号表示正确
        else:
            mark = "\033[91m" + "☒" + "\033[0m"  # 红色叉号表示错误

        # 添加行数据
        table_data.append([questions[i], real_answer_name, pred_answer_name, mark])

    accuracy = correct_count / len(questions) if len(questions) > 0 else 0
    print_table(
        f"🎯 生成结果 (准确率: {correct_count}/{len(questions)} = {accuracy * 100:.2f}%)",
        ["输入", "真实值", "预测值", "标记"],
        table_data,
    )


def plot_training_logs(training_logs):
    """
    绘制训练日志图表，包括损失函数和各参数权重的变化趋势

    参数:
    training_logs: 包含训练日志的列表
    """

    # 提取绘图所需数据
    epochs = [log["Epoch"] for log in training_logs]
    losses = [log["Loss"] for log in training_logs]
    mint_leaf_weight = [log["薄荷叶权重"] for log in training_logs]
    cucumber_slice_weight = [log["黄瓜片权重"] for log in training_logs]
    rose_salt_weight = [log["玫瑰盐权重"] for log in training_logs]

    # 设置中文字体支持
    plt.rcParams["font.sans-serif"] = [
        "SimHei",
        "Arial Unicode MS",
        "DejaVu Sans",
    ]  # 支持中文字体
    plt.rcParams["axes.unicode_minus"] = False  # 正常显示负号

    # 创建 2x2 子图布局
    fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 10))

    # 绘制 Loss 图
    ax1.plot(epochs, losses, color="red")
    ax1.set_title("损失函数变化")
    ax1.set_xlabel("训练轮次")
    ax1.set_ylabel("损失值")
    ax1.grid(True)

    # 绘制 薄荷叶权重 图
    ax2.plot(epochs, mint_leaf_weight, color="blue")
    ax2.set_title("薄荷叶权重变化")
    ax2.set_xlabel("训练轮次")
    ax2.set_ylabel("权重值")
    ax2.grid(True)

    # 绘制 黄瓜片权重 图
    ax3.plot(epochs, cucumber_slice_weight, color="green")
    ax3.set_title("黄瓜片权重变化")
    ax3.set_xlabel("训练轮次")
    ax3.set_ylabel("权重值")
    ax3.grid(True)

    # 绘制 玫瑰盐权重 图
    ax4.plot(epochs, rose_salt_weight, color="orange")
    ax4.set_title("玫瑰盐权重变化")
    ax4.set_xlabel("训练轮次")
    ax4.set_ylabel("权重值")
    ax4.grid(True)

    # 调整子图间距
    plt.tight_layout()
    plt.show()


def plot_training_loss(train_losses):
    """
    绘制模型训练过程中的损失变化图

    参数:
    train_losses: 包含训练日志的列表，每个元素应包含'epoch'和'loss'键
    """
    # 提取绘图所需数据
    epochs = [log["epoch"] for log in train_losses]
    losses = [log["loss"] for log in train_losses]

    # 设置中文字体支持
    plt.rcParams["font.sans-serif"] = [
        "SimHei",
        "Arial Unicode MS",
        "DejaVu Sans",
    ]  # 支持中文字体
    plt.rcParams["axes.unicode_minus"] = False  # 正常显示负号

    # 创建图表
    plt.figure(figsize=(10, 6))

    # 绘制 Loss 图
    plt.plot(epochs, losses, color="red", linestyle="-")
    plt.title("损失函数变化")
    plt.xlabel("轮次（Epoch）")
    plt.ylabel("损失（Loss）")
    plt.grid(True)

    # 显示图表
    plt.tight_layout()
    plt.show()


def show_first_n_lines(file_path, n=5):
    """
    显示文本文件的前n行内容

    Args:
        file_path (str): 文件路径
        n (int): 显示的行数，默认为5行
    """
    try:
        data = []
        with open(file_path, "r", encoding="utf-8") as file:
            for i in range(n):
                line = file.readline()
                data.append({"行数": i + 1, "内容": line.strip()})
                if not line:
                    break

        show_table(data)

    except FileNotFoundError:
        print(f"文件 {file_path} 未找到")
    except Exception as e:
        print(f"读取文件时发生错误: {e}")


def display_co_matrix(co_matrix):
    """
    使用 print_table 显示共现矩阵。
    """
    if co_matrix.matrix is None:
        print("错误: 共现矩阵尚未构建。请先调用 build_from_corpus 方法。")
        return

    # 1. 准备表头
    headers = ["Token ID", "Token"]
    for i in range(co_matrix.tokenizer.vocab_size):
        headers.append(co_matrix.tokenizer.ids_to_tokens.get(i, "[UNK]"))

    # 2. 准备数据行
    rows = []
    for i in range(co_matrix.vocab_size):
        row = [i, co_matrix.tokenizer.ids_to_tokens.get(i, "[UNK]")]
        row.extend(co_matrix.matrix[i, :].tolist())
        rows.append(row)

    # 3. 使用 print_table 显示
    print_table(f"共现矩阵 (窗口大小={co_matrix.window_size})", headers, rows)
