#!/usr/bin/env python3
"""
电力波形故障识别 — 1D-CNN 训练脚本 (完整版)

使用方法:
  python train_1dcnn.py --data_dir ./data --epochs 100

数据准备:
  data/
    train_data.npy   (N_train, 128)  float32  训练样本
    train_labels.npy (N_train,)      int64    训练标签
    val_data.npy     (N_val, 128)
    val_labels.npy   (N_val,)
    test_data.npy    (N_test, 128)
    test_labels.npy  (N_test,)

输出:
  checkpoints/best_model.pth     最佳权重
  checkpoints/fault_1dcnn.onnx   ONNX 模型 (用于 RKNN 转换)
"""

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from sklearn.metrics import classification_report
import os
import argparse
from collections import Counter


# ============================================================
# 故障类型名称
# ============================================================
FAULT_NAMES = [
    'Normal',       # 0  正常
    'A-G',          # 1  A相单相接地
    'BC',           # 2  BC相间短路
    'AB-G',         # 3  AB两相接地
    'ABC',          # 4  三相短路
]

NUM_CLASSES = len(FAULT_NAMES)
SEQ_LEN     = 128
INPUT_CH    = 1


# ============================================================
# 模型定义
# ============================================================
class Fault1DCNN(nn.Module):
    """
    轻量级 1D-CNN，面向 RK3568 NPU 部署
    参数量 ~25K，模型 FP32 大小 ~100KB
    """

    def __init__(self, num_classes=NUM_CLASSES,
                 input_channels=INPUT_CH, seq_len=SEQ_LEN):
        super().__init__()
        _ = seq_len  # unused, kept for API compatibility

        self.block1 = nn.Sequential(
            nn.Conv1d(input_channels, 16, kernel_size=7,
                      stride=1, padding='same', bias=False),
            nn.BatchNorm1d(16),
            nn.ReLU(inplace=True),
            nn.MaxPool1d(kernel_size=2),       # 128 → 64
        )

        self.block2 = nn.Sequential(
            nn.Conv1d(16, 32, kernel_size=5,
                      stride=1, padding='same', bias=False),
            nn.BatchNorm1d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool1d(kernel_size=2),       # 64 → 32
        )

        self.block3 = nn.Sequential(
            nn.Conv1d(32, 64, kernel_size=3,
                      stride=1, padding='same', bias=False),
            nn.BatchNorm1d(64),
            nn.ReLU(inplace=True),
        )

        self.gap = nn.AdaptiveAvgPool1d(1)      # 32 → 1
        self.classifier = nn.Linear(64, num_classes)
        self._init_weights()

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv1d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out')
            elif isinstance(m, nn.BatchNorm1d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)

    def forward(self, x):
        x = self.block1(x)       # (B, 16, 64)
        x = self.block2(x)       # (B, 32, 32)
        x = self.block3(x)       # (B, 64, 32)
        x = self.gap(x)          # (B, 64, 1)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)   # (B, num_classes)
        return x


# ============================================================
# 数据集 (含数据增强)
# ============================================================
class WaveformDataset(Dataset):
    def __init__(self, data_path, label_path, augment=False):
        self.data   = np.load(data_path).astype(np.float32)
        self.labels = np.load(label_path).astype(np.int64)
        self.augment = augment

        if self.data.ndim == 2:
            self.data = self.data[:, np.newaxis, :]

        # 统计归一化参数
        self.mean = self.data.mean()
        self.std  = self.data.std() + 1e-8

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

    def __getitem__(self, idx):
        x = self.data[idx].copy()
        y = self.labels[idx]

        # Z-score 归一化
        x = (x - self.mean) / self.std

        if self.augment:
            x = self._augment(x)

        return torch.from_numpy(x), torch.tensor(y)

    @staticmethod
    def _augment(x):
        """时间域数据增强 (numpy, shape=(1, L))"""
        shift = np.random.randint(-8, 9)
        if shift != 0:
            x = np.roll(x, shift, axis=-1)
        x *= (1.0 + np.random.uniform(-0.03, 0.03))
        x += np.random.normal(0, 1e-3, x.shape).astype(np.float32)
        return x


# ============================================================
# 训练 / 验证
# ============================================================
def train_one_epoch(model, loader, criterion, optimizer, device):
    model.train()
    total_loss, correct, total = 0.0, 0, 0
    for x, y in loader:
        x, y = x.to(device), y.to(device)
        optimizer.zero_grad()
        loss = criterion(model(x), y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * x.size(0)
        correct += (model(x).argmax(1) == y).sum().item()  # re-run for stats
        total += x.size(0)

    # Recalculate accuracy properly
    model.eval()
    correct, total = 0, 0
    with torch.no_grad():
        for x, y in loader:
            x, y = x.to(device), y.to(device)
            logits = model(x)
            correct += (logits.argmax(1) == y).sum().item()
            total += x.size(0)

    return total_loss / total, correct / total


@torch.no_grad()
def validate(model, loader, criterion, device):
    model.eval()
    total_loss, correct, total = 0.0, 0, 0
    all_preds, all_labels = [], []
    for x, y in loader:
        x, y = x.to(device), y.to(device)
        logits = model(x)
        loss = criterion(logits, y)
        total_loss += loss.item() * x.size(0)
        preds = logits.argmax(1)
        correct += (preds == y).sum().item()
        total += x.size(0)
        all_preds.extend(preds.cpu().numpy())
        all_labels.extend(y.cpu().numpy())
    return total_loss / total, correct / total, all_preds, all_labels


# ============================================================
# ONNX 导出
# ============================================================
def export_onnx(model, checkpoint_path, output_path, device):
    state = torch.load(checkpoint_path, map_location=device)
    model.load_state_dict(state)
    model.eval()

    dummy = torch.randn(1, INPUT_CH, SEQ_LEN).to(device)

    torch.onnx.export(
        model, dummy, output_path,
        input_names=['waveform'],
        output_names=['logits'],
        dynamic_axes=None,
        opset_version=12,
        do_constant_folding=True,
        export_params=True,
        verbose=False,
    )

    # 验证 ONNX
    try:
        import onnx
        onnx_model = onnx.load(output_path)
        onnx.checker.check_model(onnx_model)
        op_types = [n.op_type for n in onnx_model.graph.node]
        print(f"\nONNX 算子统计: {dict(Counter(op_types))}")
        print("ONNX 模型验证通过 ✓")
    except ImportError:
        print("WARNING: 未安装 onnx 库，跳过 ONNX 验证")

    print(f"ONNX 模型已导出: {output_path}")


# ============================================================
# Main
# ============================================================
def main():
    parser = argparse.ArgumentParser(description='电力波形故障识别 1D-CNN')
    parser.add_argument('--data_dir', default='data', help='数据目录')
    parser.add_argument('--batch_size', type=int, default=64)
    parser.add_argument('--epochs', type=int, default=100)
    parser.add_argument('--lr', type=float, default=0.001)
    parser.add_argument('--device', default='auto')
    args = parser.parse_args()

    # 设备
    if args.device == 'auto':
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    else:
        device = torch.device(args.device)
    print(f"设备: {device}")

    # 数据
    train_ds = WaveformDataset(
        f'{args.data_dir}/train_data.npy',
        f'{args.data_dir}/train_labels.npy', augment=True)
    val_ds   = WaveformDataset(
        f'{args.data_dir}/val_data.npy',
        f'{args.data_dir}/val_labels.npy', augment=False)

    train_loader = DataLoader(train_ds, batch_size=args.batch_size,
                              shuffle=True, num_workers=2)
    val_loader   = DataLoader(val_ds, batch_size=args.batch_size,
                              shuffle=False, num_workers=2)

    print(f"训练集: {len(train_ds):,}  验证集: {len(val_ds):,}")
    print(f"归一化: mean={train_ds.mean:.3f}, std={train_ds.std:.3f}")

    # 数据不平衡检测
    label_counts = Counter(train_ds.labels)
    print(f"训练集类别分布: {dict(label_counts)}")

    # 模型
    model = Fault1DCNN().to(device)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"模型参数量: {n_params:,}  (~{n_params*4/1024:.0f} KB FP32)")

    # 训练
    optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-4)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode='max', factor=0.5, patience=10, verbose=True)
    criterion = nn.CrossEntropyLoss()

    os.makedirs('checkpoints', exist_ok=True)
    best_acc = 0.0

    for epoch in range(1, args.epochs + 1):
        train_loss, train_acc = train_one_epoch(
            model, train_loader, criterion, optimizer, device)
        val_loss, val_acc, preds, labels = validate(
            model, val_loader, criterion, device)
        scheduler.step(val_acc)

        if epoch % 10 == 0 or epoch == 1:
            print(f"Epoch {epoch:3d} | "
                  f"Train: loss={train_loss:.4f} acc={train_acc:.2%} | "
                  f"Val: loss={val_loss:.4f} acc={val_acc:.2%}")

        if val_acc > best_acc:
            best_acc = val_acc
            torch.save(model.state_dict(), 'checkpoints/best_model.pth')

    print(f"\n最佳验证准确率: {best_acc:.2%}")

    # 最终分类报告
    print("\n" + "=" * 55)
    print(classification_report(labels, preds, target_names=FAULT_NAMES))
    print("=" * 55)

    # 导出 ONNX
    export_onnx(model, 'checkpoints/best_model.pth',
                'checkpoints/fault_1dcnn.onnx', device)


if __name__ == '__main__':
    main()