#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
DirectML 兼容文本生成器训练与生成脚本（工业协同版）
=====================================================
核心创新：
  - 手工实现 BasicLSTM，仅使用 Linear、Sigmoid、Tanh 等基础算子，
    完全规避 DirectML 对原生 nn.LSTM 的混合精度缺失问题。
  - 采用 GPU（DirectML）主力计算 + CPU 兜底采样的协同架构，
    在老破小集成显卡上实现稳定的端到端训练。
  - 内置 5 个专业报告模板，模拟八维稳态分析输出，自动构建数据集。

适用环境：
  - Windows + DirectML（任意支持 DX12 的显卡）
  - 纯 CPU 也可运行（自动降级）

作者：李文龙
理论支撑：全域共生复杂动力系统
发布日期：2026-05-06
"""

import random
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader

# ===================== 设备检测与分工 =====================
def setup_devices():
    """检测 DirectML 可用性，返回 GPU 设备及 CPU 设备"""
    try:
        import torch_directml
        gpu = torch_directml.device()
        cpu = torch.device("cpu")
        print("[设备] 协同模式：GPU (DirectML) 主力 + CPU 兜底")
        return gpu, cpu
    except ImportError:
        gpu = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        cpu = torch.device("cpu")
        print(f"[设备] DirectML 未安装，使用：{gpu}")
        return gpu, cpu

GPU_DEVICE, CPU_DEVICE = setup_devices()

# ===================== 手工 LSTM 算子（兼容 DirectML） =====================
class BasicLSTM(nn.Module):
    """
    纯基础算子实现的 LSTM。
    仅使用 nn.Linear、Sigmoid、Tanh 等 DirectML 原生支持的操作，
    无任何 AMP 不兼容问题。
    """
    def __init__(self, input_size: int, hidden_size: int, batch_first: bool = True):
        super().__init__()
        self.hidden_size = hidden_size
        self.batch_first = batch_first
        # 输入门、遗忘门、细胞候选、输出门 合并为一个线性变换
        self.W = nn.Linear(input_size + hidden_size, 4 * hidden_size)

    def forward(self, x: torch.Tensor, hidden: tuple = None):
        if self.batch_first:
            x = x.transpose(0, 1)                # (seq, batch, feat)
        seq_len, batch, _ = x.size()

        if hidden is None:
            h = torch.zeros(1, batch, self.hidden_size, device=x.device)
            c = torch.zeros(1, batch, self.hidden_size, device=x.device)
        else:
            h, c = hidden

        outputs = []
        for t in range(seq_len):
            combined = torch.cat([x[t], h[0]], dim=1)   # (batch, input_size+hidden)
            gates = self.W(combined)                    # (batch, 4*hidden)
            i, f, g, o = gates.chunk(4, 1)
            i = torch.sigmoid(i)
            f = torch.sigmoid(f)
            g = torch.tanh(g)
            o = torch.sigmoid(o)

            c = f * c + i * g
            h = o * torch.tanh(c)
            outputs.append(h.unsqueeze(1) if self.batch_first else h)

        if self.batch_first:
            outputs = torch.cat(outputs, dim=1)
        else:
            outputs = torch.stack(outputs)

        return outputs, (h, c)

# ===================== 文本生成模型 =====================
class LightTextGenerator(nn.Module):
    """
    轻量级八维报告生成器。
    输入八维状态向量，输出稳态评估文本。
    """
    def __init__(self, vocab_size: int, embed_dim: int, hidden_dim: int, max_seq_len: int = 256):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.pos_encoding = nn.Parameter(torch.zeros(1, max_seq_len, embed_dim))
        self.dim_proj = nn.Linear(8, embed_dim)           # 将8维状态映射到 embedding 空间
        self.lstm = BasicLSTM(embed_dim, hidden_dim, batch_first=True)
        self.output_proj = nn.Linear(hidden_dim, vocab_size)

    def forward(self, dims: torch.Tensor, tgt_ids: torch.Tensor):
        seq_len = tgt_ids.shape[1]
        emb = self.embedding(tgt_ids) + self.pos_encoding[:, :seq_len, :]
        mem = self.dim_proj(dims).unsqueeze(1)            # (batch, 1, embed_dim)
        out, _ = self.lstm(torch.cat([mem, emb], dim=1))
        return self.output_proj(out[:, 1:, :])            # 预测下一个 token

    @torch.no_grad()
    def generate(self, dims: torch.Tensor, vocab: dict, rev_vocab: dict,
                 max_len: int = 100, temperature: float = 0.7) -> str:
        """
        自回归生成报告文本，GPU 做前向，CPU 负责采样（避免 DirectML 采样算子缺失）。
        """
        self.eval()
        dims = dims.to(GPU_DEVICE)
        input_ids = torch.tensor([[vocab['<BOS>']]], dtype=torch.long, device=GPU_DEVICE)
        hidden = None
        generated_ids = []

        for _ in range(max_len):
            seq_len = input_ids.shape[1]
            emb = self.embedding(input_ids) + self.pos_encoding[:, :seq_len, :]
            mem = self.dim_proj(dims).unsqueeze(1)
            out, hidden = self.lstm(torch.cat([mem, emb], dim=1), hidden)
            logits = self.output_proj(out[:, -1, :])      # (batch, vocab_size)

            # ---- 采样回退到 CPU（DirectML 对 multinomial 支持有限）----
            logits_cpu = logits.to(CPU_DEVICE).squeeze()
            probs = F.softmax(logits_cpu / temperature, dim=-1)
            next_id_cpu = torch.multinomial(probs, num_samples=1)
            next_id = next_id_cpu.to(GPU_DEVICE).view(1, 1)

            if next_id.item() == vocab['<EOS>']:
                break
            generated_ids.append(next_id.item())
            input_ids = torch.cat([input_ids, next_id], dim=1)

        return "".join([rev_vocab.get(i, "") for i in generated_ids])

# ===================== 数据集构建 =====================
# 5 套专业报告模板（覆盖诊断、剖析、协同、收敛、进化等维度）
_REPORT_TEMPLATES = [
    "【全域诊断】系统当前处于稳态收敛期。X1={X1:.2f}表明认知耦合度良好，X2={X2:.2f}反映因果推理链完整。短板维度为X4={X4:.2f}（存储可靠性），建议启动增量备份任务。X6={X6:.2f}显示自进化活跃度适中。综合评分：M={m_val:.3f}，系统健康。",
    "【八维剖析】认知耦合X1={X1:.2f}，因果强度X2={X2:.2f}，知识熵X3={X3:.2f}，存储可靠性X4={X4:.2f}，推理深度X5={X5:.2f}，自进化活跃度X6={X6:.2f}，资源效率X7={X7:.2f}，时间稳定性X8={X8:.2f}。当前主要矛盾集中在X4维度，建议优先加固存储层。",
    "【协同评估】X1={X1:.2f}与X2={X2:.2f}形成正向耦合，但X5={X5:.2f}相对滞后。X8={X8:.2f}显示系统时间稳定性良好。建议增强推理模块递归深度。",
    "【收敛分析】X1-X2簇得分{X1:.2f}-{X2:.2f}，与X5-X6簇得分{X5:.2f}-{X6:.2f}之间存在{coupling_gap:.2f}的梯度。X4={X4:.2f}处于临界状态。当前M={m_val:.3f}，偏差{m_dev:.3f}。",
    "【进化轨迹】X1={X1:.2f}较上一周期变化{delta_x1:+.2f}，X6={X6:.2f}维持高位。需关注X3={X3:.2f}上升趋势。X7={X7:.2f}表现优异。综合稳态M={m_val:.3f}。",
]

def build_dataset(num_samples: int = 2000):
    """用模板随机生成训练数据，每条样本附带随机八维得分"""
    samples = []
    for _ in range(num_samples):
        tpl = random.choice(_REPORT_TEMPLATES)
        dims = {f"X{i}": random.uniform(0.5, 0.98) for i in range(1, 9)}
        m_val = sum(dims.values()) / 8
        coupling_gap = abs(dims["X1"] - dims["X6"])
        m_dev = abs(m_val - 0.928)
        delta_x1 = random.uniform(-0.08, 0.08)

        report = tpl.format(
            X1=dims["X1"], X2=dims["X2"], X3=dims["X3"], X4=dims["X4"],
            X5=dims["X5"], X6=dims["X6"], X7=dims["X7"], X8=dims["X8"],
            m_val=m_val, coupling_gap=coupling_gap, m_dev=m_dev, delta_x1=delta_x1
        )
        samples.append({"dims": dims, "report": report})
    return samples

def build_vocab(samples: list) -> tuple[dict, dict]:
    """构建字符级词表"""
    text = "".join(s["report"] for s in samples)
    vocab = {'<PAD>': 0, '<UNK>': 1, '<BOS>': 2, '<EOS>': 3}
    for ch in sorted(set(text)):
        if ch not in vocab:
            vocab[ch] = len(vocab)
    rev_vocab = {v: k for k, v in vocab.items()}
    return vocab, rev_vocab

def text_to_ids(text: str, vocab: dict, max_len: int) -> torch.Tensor:
    """文本 → 固定长度 token id 序列"""
    ids = [vocab['<BOS>']] + [vocab.get(c, vocab['<UNK>']) for c in text] + [vocab['<EOS>']]
    if len(ids) < max_len:
        ids += [vocab['<PAD>']] * (max_len - len(ids))
    return torch.tensor(ids[:max_len], dtype=torch.long)

class ReportDataset(Dataset):
    """八维报告数据集"""
    def __init__(self, samples, vocab, max_len):
        self.samples = samples
        self.vocab = vocab
        self.max_len = max_len

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        s = self.samples[idx]
        dims = torch.tensor([s["dims"][f"X{i}"] for i in range(1, 9)], dtype=torch.float32)
        tgt = text_to_ids(s["report"], self.vocab, self.max_len)
        return dims, tgt

# ===================== 训练与生成主流程 =====================
def main():
    # 超参数
    BATCH_SIZE = 16
    EPOCHS = 5
    MAX_SEQ_LEN = 128
    EMBED_DIM = 64
    HIDDEN_DIM = 128

    print("=" * 60)
    print("DirectML 兼容文本生成器 - 工业协同版")
    print("=" * 60)

    # 1. 准备数据
    print("[1/4] 生成训练数据...")
    data = build_dataset(2000)
    vocab, rev_vocab = build_vocab(data)
    print(f"      词表大小: {len(vocab)}")

    # 2. 构建 DataLoader
    loader = DataLoader(
        ReportDataset(data, vocab, MAX_SEQ_LEN),
        batch_size=BATCH_SIZE,
        shuffle=True
    )

    # 3. 初始化模型
    print("[2/4] 初始化模型...")
    model = LightTextGenerator(len(vocab), EMBED_DIM, HIDDEN_DIM).to(GPU_DEVICE)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    loss_fn = nn.CrossEntropyLoss(ignore_index=vocab['<PAD>'])

    # 4. 训练
    print("[3/4] 开始训练 (FP32, 无 AMP)...")
    for epoch in range(EPOCHS):
        model.train()
        total_loss = 0.0
        for dims, tgt in loader:
            dims, tgt = dims.to(GPU_DEVICE), tgt.to(GPU_DEVICE)
            pred = model(dims, tgt[:, :-1])
            loss = loss_fn(pred.reshape(-1, len(vocab)), tgt[:, 1:].reshape(-1))
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        print(f"  Epoch {epoch+1}/{EPOCHS} | Loss: {total_loss:.4f}")

    # 5. 测试生成
    print("[4/4] 生成示例报告...")
    test_dims = torch.tensor([[0.85, 0.86, 0.50, 0.88, 0.90, 0.92, 0.93, 0.94]], dtype=torch.float32)
    report = model.generate(test_dims, vocab, rev_vocab, temperature=0.7)
    print(f"生成结果：{report}")
    print("=" * 60)
    print("训练完成。该模型可用于老破小集成显卡，无 AMP 兼容性问题。")

if __name__ == "__main__":
    main()

---

## DirectML 兼容文本生成器（工业协同版）— 老破小专用

### 1. 硬件环境

* **操作系统**：Windows 10 / 11（64 位）
* **显卡**：任何支持 DirectX 12 的显卡（集成显卡、独立显卡均可），包括 Intel UHD、AMD Radeon、NVIDIA 等
* **内存**：建议 8 GB 以上（4 GB 也可运行，但需减少数据量）
* **纯 CPU 模式**：若完全没有独显，脚本会自动退化为 CPU 训练，无需额外配置

### 2. 软件依赖

* **Python**：3.12 版本（推荐 [Python 3.12.9](https://www.python.org/downloads/)）
* **PyTorch**：2.5+（CPU 版或 CUDA 版均可）
* **DirectML 插件**：`torch-directml`

> 注意：`torch-directml` 目前最高支持 Python 3.12，请勿使用 3.13 或更高版本。

### 3. 安装步骤

```powershell
# 1. 创建虚拟环境（可选但推荐）
python -m venv dml_env
dml_env\Scripts\activate

# 2. 安装 PyTorch（CPU 版本即可，DirectML 会自动接管）
pip install torch torchvision torchaudio

# 3. 安装 DirectML 后端
pip install torch-directml
```

*如果你的 Python 版本过高导致 `torch-directml` 安装失败，请降级到 Python 3.12。*

### 4. 运行脚本

直接将以下代码保存为 `train_text_gen.py`，然后在项目目录下执行：

```powershell
python train_text_gen.py
```

脚本会自动：
1. 检测硬件（DirectML 可用则使用 GPU，否则回退 CPU）
2. 生成 2000 条八维诊断报告作为训练数据
3. 训练一个手工实现的 LSTM 文本生成器（纯 FP32 精度，无 AMP 兼容问题）
4. 最后输出一条由模型生成的示例报告

### 5. 输出示例

```
[设备] 协同模式：GPU (DirectML) 主力 + CPU 兜底
[1/4] 生成训练数据...
      词表大小: 112
[2/4] 初始化模型...
[3/4] 开始训练 (FP32, 无 AMP)...
  Epoch 1/5 | Loss: 560.2542
  Epoch 2/5 | Loss: 524.7070
  ...
[4/4] 生成示例报告...
生成结果：域。统。测统建。议统板。建缺统建。测建统。测统...
```

### 6. 工作原理简述

* **手工 LSTM**：完全使用 `Linear`、`Sigmoid`、`Tanh` 等 DirectML 原生支持的算子重写 LSTM，不调用任何 `nn.LSTM`，从而彻底规避 DirectML 对混合精度训练的支持缺失问题。
* **FP32 全精度训练**：关闭 `autocast`，强制在 DirectML 上运行全精度计算，即便集成显卡也能稳定运行。
* **GPU + CPU 协同生成**：推理时 GPU 负责前向传播，CPU 负责采样和后续处理，进一步避免 DirectML 对 `multinomial` 采样的不完整支持。

### 7. 注意事项

* 训练时的 `aten::lerp.Scalar_out` 警告是 Adam 优化器的一个步骤，DirectML 会自动交给 CPU 执行，**不影响训练结果和稳定性**。
* 若生成文本质量不理想，可增大 `EPOCHS` 和 `num_samples` 参数，或增加更丰富的训练模板。