#include <Arduino.h>

// ==================== 全局变量和配置 ====================
const uint8_t pwmPin = 9;  // PWM输出引脚
const uint16_t baseFreq = 1000;  // 基础频率1kHz
const int MAX_INPUT_LENGTH = 10;  // 最多10个字符
const int MAX_UNIQUE_CHARS = 10;  // 最多10个唯一字符

// 霍夫曼编码相关结构
struct CharFreq {
  char character;  // ASCII字符
  uint8_t frequency;
  uint8_t codeLength;
  uint8_t huffmanCode;  // 最多8位编码
};

// 动态字符频率表
CharFreq charTable[MAX_UNIQUE_CHARS];
uint8_t uniqueCharCount = 0;

// 汉明码(15,11)生成矩阵 - 存储在PROGMEM中
const uint8_t PROGMEM G_matrix_15_11[15][11] = {
  {1, 0, 0, 0, 1, 1, 0, 1, 1, 1, 0},  // P1
  {0, 1, 0, 0, 1, 0, 1, 1, 1, 0, 1},  // P2  
  {1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0},  // D1
  {0, 0, 1, 0, 0, 1, 1, 0, 1, 1, 1},  // P3
  {0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0},  // D2
  {0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0},  // D3
  {0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0},  // D4
  {0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0},  // D5
  {1, 1, 1, 0, 0, 0, 0, 1, 0, 0, 1},  // P4
  {0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0},  // D6
  {0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0},  // D7
  {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1},  // D8
  {0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0},  // D9
  {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},  // D10
  {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}   // D11
};

// ==================== 字符处理函数 ====================
void countCharFrequencies(const String& message) {
  uniqueCharCount = 0;
  
  // 清空字符表
  for (int i = 0; i < MAX_UNIQUE_CHARS; i++) {
    charTable[i].character = 0;
    charTable[i].frequency = 0;
    charTable[i].huffmanCode = 0;
    charTable[i].codeLength = 0;
  }
  
  // 只处理ASCII字符
  for (int pos = 0; pos < message.length() && pos < MAX_INPUT_LENGTH; pos++) {
    char currentChar = message[pos];
    
    // 查找字符是否已存在
    bool found = false;
    for (int i = 0; i < uniqueCharCount; i++) {
      if (charTable[i].character == currentChar) {
        charTable[i].frequency++;
        found = true;
        break;
      }
    }
    
    // 如果是新字符，添加到表中
    if (!found && uniqueCharCount < MAX_UNIQUE_CHARS) {
      charTable[uniqueCharCount].character = currentChar;
      charTable[uniqueCharCount].frequency = 1;
      uniqueCharCount++;
    }
  }
  
  Serial.println(F("字符频率统计:"));
  for (int i = 0; i < uniqueCharCount; i++) {
    Serial.print(F("'"));
    Serial.print(charTable[i].character);
    Serial.print(F("': "));
    Serial.println(charTable[i].frequency);
  }
}

// 简化的霍夫曼编码生成
void generateSimpleHuffmanCodes() {
  if (uniqueCharCount == 1) {
    charTable[0].huffmanCode = 0;
    charTable[0].codeLength = 1;
    return;
  }
  
  // 按频率排序（冒泡排序）
  for (int i = 0; i < uniqueCharCount - 1; i++) {
    for (int j = 0; j < uniqueCharCount - i - 1; j++) {
      if (charTable[j].frequency < charTable[j + 1].frequency) {
        CharFreq temp = charTable[j];
        charTable[j] = charTable[j + 1];
        charTable[j + 1] = temp;
      }
    }
  }
  
  // 简化的霍夫曼编码分配
  uint8_t codeLength = 1;
  uint8_t maxCodesForLength = 2;
  uint8_t assignedCodes = 0;
  
  for (int i = 0; i < uniqueCharCount; i++) {
    if (assignedCodes >= maxCodesForLength) {
      codeLength++;
      maxCodesForLength *= 2;
      assignedCodes = 0;
    }
    
    charTable[i].huffmanCode = assignedCodes;
    charTable[i].codeLength = codeLength;
    assignedCodes++;
  }
  
  Serial.println(F("生成的霍夫曼编码:"));
  for (int i = 0; i < uniqueCharCount; i++) {
    Serial.print(F("'"));
    Serial.print(charTable[i].character);
    Serial.print(F("': "));
    // 打印二进制编码
    for (int bit = charTable[i].codeLength - 1; bit >= 0; bit--) {
      Serial.print((charTable[i].huffmanCode >> bit) & 1);
    }
    Serial.println();
  }
}

// ==================== 霍夫曼编码并返回位数 ====================
int huffmanEncode(const String& message, String& huffmanBits) {
  huffmanBits = "";
  
  for (int pos = 0; pos < message.length() && pos < MAX_INPUT_LENGTH; pos++) {
    char currentChar = message[pos];
    
    // 查找字符的霍夫曼编码
    for (int i = 0; i < uniqueCharCount; i++) {
      if (charTable[i].character == currentChar) {
        uint8_t code = charTable[i].huffmanCode;
        uint8_t codeLen = charTable[i].codeLength;
        
        // 将编码位添加到字符串
        for (int bit = codeLen - 1; bit >= 0; bit--) {
          huffmanBits += ((code >> bit) & 1) ? "1" : "0";
        }
        break;
      }
    }
  }
  
  Serial.print(F("霍夫曼编码结果: "));
  Serial.println(huffmanBits);
  Serial.print(F("总位数: "));
  Serial.println(huffmanBits.length());
  
  return huffmanBits.length();
}

// ==================== 汉明码(15,11)编码函数 ====================
void hamming_encode_11bits(uint8_t data[11], uint8_t encoded[15]) {
  for (int i = 0; i < 15; i++) {
    encoded[i] = 0;
    for (int j = 0; j < 11; j++) {
      encoded[i] ^= (pgm_read_byte(&G_matrix_15_11[i][j]) & data[j]);
    }
  }
}

String applyHammingEncoding(const String& binaryStr) {
  String result = "";
  int totalBits = binaryStr.length();
  
  // 计算需要多少次汉明编码
  int hammingBlocks = (totalBits + 10) / 11;  // 向上取整
  int paddingBits = hammingBlocks * 11 - totalBits;
  
  Serial.print(F("原始霍夫曼编码位数: "));
  Serial.println(totalBits);
  Serial.print(F("汉明编码块数: "));
  Serial.println(hammingBlocks);
  Serial.print(F("填充位数: "));
  Serial.println(paddingBits);
  
  String paddedBinary = binaryStr;
  // 用0填充到11的倍数
  for (int i = 0; i < paddingBits; i++) {
    paddedBinary += "0";
  }
  
  Serial.print(F("填充后的霍夫曼编码: "));
  Serial.println(paddedBinary);
  
  // 每11位进行一次汉明编码
  for (int block = 0; block < hammingBlocks; block++) {
    uint8_t data[11];
    uint8_t encoded[15];
    
    // 提取11位数据
    String blockData = "";
    for (int i = 0; i < 11; i++) {
      int bitPos = block * 11 + i;
      data[i] = (bitPos < paddedBinary.length() && paddedBinary[bitPos] == '1') ? 1 : 0;
      blockData += (data[i] == 1) ? "1" : "0";
    }
    
    // 汉明编码
    hamming_encode_11bits(data, encoded);
    
    // 显示每个块的编码过程
    Serial.print(F("块 "));
    Serial.print(block + 1);
    Serial.print(F(" 输入数据(11位): "));
    Serial.println(blockData);
    
    String encodedBlock = "";
    // 将编码结果添加到输出
    for (int i = 0; i < 15; i++) {
      encodedBlock += (encoded[i] == 1) ? "1" : "0";
      result += (encoded[i] == 1) ? "1" : "0";
    }
    
    Serial.print(F("块 "));
    Serial.print(block + 1);
    Serial.print(F(" 汉明编码(15位): "));
    Serial.println(encodedBlock);
  }
  
  // 添加编码块数和填充位数的11位二进制表示
  String blockCountBinary = "";
  String paddingCountBinary = "";
  
  // 转换为11位二进制
  for (int i = 10; i >= 0; i--) {
    blockCountBinary += ((hammingBlocks >> i) & 1) ? "1" : "0";
    paddingCountBinary += ((paddingBits >> i) & 1) ? "1" : "0";
  }
  
  result += blockCountBinary + paddingCountBinary;
  
  Serial.println(F(""));
  Serial.print(F("编码块数信息 (11位): "));
  Serial.println(blockCountBinary);
  Serial.print(F("填充位数信息 (11位): "));
  Serial.println(paddingCountBinary);
  
  Serial.println(F(""));
  Serial.println(F("==================== 最终完整编码 ===================="));
  Serial.print(F("完整编码结果: "));
  Serial.println(result);
  Serial.print(F("最终总位数: "));
  Serial.println(result.length());
  Serial.println(F("编码组成: [汉明编码数据块] + [块数信息11位] + [填充位数信息11位]"));
  Serial.println(F("========================================================"));
  
  return result;
}

// ==================== 自适应调制位数计算 ====================
int calculateOptimalModulationBits(int totalBits) {
  Serial.print(F("计算最优调制位数，总位数: "));
  Serial.println(totalBits);
  
  int bestN = 2;  // 默认值
  int minRemainder = totalBits % 2;
  
  Serial.println(F("测试不同调制位数:"));
  
  // 首先找出所有能够整除的n值，按照从大到小的顺序测试
  int divisibleN = 0;
  for (int n = 8; n >= 2; n--) {  // 从8到2，优先选择大的n值
    int remainder = totalBits % n;
    int groups = (totalBits + n - 1) / n;  // 向上取整
    
    Serial.print(F("n="));
    Serial.print(n);
    Serial.print(F(": 组数="));
    Serial.print(groups);
    Serial.print(F(", 余数="));
    Serial.print(remainder);
    
    if (remainder == 0) {
      Serial.print(F(" -> 完全整除"));
      if (divisibleN == 0) {  // 第一个找到的完全整除的n值（也是最大的）
        divisibleN = n;
      }
    }
    Serial.println();
    
    // 同时记录余数最小的情况，作为备选
    if (remainder < minRemainder) {
      minRemainder = remainder;
      bestN = n;
    }
  }
  
  // 如果找到了能完全整除的n值，优先选择最大的那个
  if (divisibleN > 0) {
    Serial.print(F("找到能完全整除的n值，选择最大的: n="));
    Serial.println(divisibleN);
    return divisibleN;
  }
  
  // 如果没有完全整除的，选择余数最小的
  Serial.print(F("没有完全整除的n值，选择余数最小的: n="));
  Serial.print(bestN);
  Serial.print(F(" (余数="));
  Serial.print(minRemainder);
  Serial.println(F(")"));
  return bestN;
}

// ==================== PWM调制函数 ====================
void setupPWM() {
  pinMode(pwmPin, OUTPUT);
  digitalWrite(pwmPin, LOW);  // 初始化为低电平
  
  // 暂时不启动PWM定时器，等到需要传输时再启动
  Serial.println(F("PWM引脚初始化完成，输出为低电平（待机状态）"));
}

void startPWMTimer() {
  // 设置Timer1为Fast PWM模式
  TCCR1A = _BV(WGM11) | _BV(COM1A1);
  TCCR1B = _BV(WGM13) | _BV(WGM12) | _BV(CS11);
  
  ICR1 = 1999;  // 1kHz频率
  OCR1A = 0;    // 初始占空比为0
  
  Serial.println(F("PWM定时器已启动"));
}

void stopPWMTimer() {
  // 停止PWM定时器
  TCCR1A = 0;
  TCCR1B = 0;
  
  // 确保引脚输出低电平
  digitalWrite(pwmPin, LOW);
  
  Serial.println(F("PWM定时器已停止，引脚输出低电平"));
}

void setPWMDutyCycle(uint8_t dutyCycle) {
  uint16_t ocrValue = (uint32_t)ICR1 * dutyCycle / 255;
  OCR1A = ocrValue;
}

void modulateData(const String& bits) {
  int totalBits = bits.length();
  if (totalBits == 0) return;
  
  int n = calculateOptimalModulationBits(totalBits);
  int maxValue = (1 << n) - 1;  // 2^n - 1
  int totalGroups = (totalBits + n - 1) / n;  // 向上取整
  
  Serial.print(F("调制参数: n="));
  Serial.print(n);
  Serial.print(F(", 最大值="));
  Serial.print(maxValue);
  Serial.print(F(", 总组数="));
  Serial.println(totalGroups);
  
  Serial.println(F("准备开始传输..."));
  delay(1000);  // 给接收端1秒准备时间
  
  // 启动PWM定时器
  startPWMTimer();
  
  Serial.println(F("开始单次PWM调制传输:"));
  Serial.print(F("将输出 "));
  Serial.print(totalGroups);
  Serial.println(F(" 个PWM波形，然后停止"));
  
  // 每n位进行一次调制
  for (int group = 0; group < totalGroups; group++) {
    String nBits = "";
    int value = 0;
    
    // 提取n位数据
    for (int bit = 0; bit < n; bit++) {
      int bitIndex = group * n + bit;
      if (bitIndex < totalBits) {
        nBits += bits[bitIndex];
        if (bits[bitIndex] == '1') {
          value |= (1 << (n - 1 - bit));
        }
      } else {
        nBits += "0";  // 填充0
      }
    }
    
    // 使用你指定的映射函数计算占空比
    uint8_t dutyCycle = map(value + 1, 0, (1 << n) + 1, 0, 255);
    
    Serial.print(F("波形"));
    Serial.print(group + 1);
    Serial.print(F("/"));
    Serial.print(totalGroups);
    Serial.print(F(": "));
    Serial.print(nBits);
    Serial.print(F(" (值="));
    Serial.print(value);
    Serial.print(F(") -> 占空比="));
    Serial.print(dutyCycle);
    Serial.print(F(" ("));
    Serial.print((float)dutyCycle * 100 / 255, 1);
    Serial.print(F("%) - 发送中..."));
    
    setPWMDutyCycle(dutyCycle);
    delay(1000);  // 每个波形持续1秒
    
    Serial.println(F(" 完成"));
  }
  
  // 数据传输完成，停止PWM定时器并设置引脚为低电平
  stopPWMTimer();
  
  Serial.println(F(""));
  Serial.println(F("==================== 单次传输完成 ===================="));
  Serial.print(F("已成功发送 "));
  Serial.print(totalGroups);
  Serial.println(F(" 个PWM波形"));
  Serial.println(F("PWM定时器已停止，引脚恢复低电平待机状态"));
  Serial.println(F("========================================================"));
}

// ==================== 主要处理函数 ====================
void processMessage(const String& message) {
  if (message.length() == 0) {
    Serial.println(F("错误: 输入消息为空"));
    return;
  }
  
  if (message.length() > MAX_INPUT_LENGTH) {
    Serial.println(F("错误: 输入消息超过10个字符限制"));
    return;
  }
  
  Serial.println(F("\n==================== 开始处理 ===================="));
  Serial.print(F("原始消息: "));
  Serial.println(message);
  Serial.print(F("消息长度: "));
  Serial.println(message.length());
  
  // 步骤1: 统计字符频率
  Serial.println(F("\n步骤1: 统计字符频率"));
  countCharFrequencies(message);
  
  // 步骤2: 生成霍夫曼编码
  Serial.println(F("\n步骤2: 生成霍夫曼编码"));
  generateSimpleHuffmanCodes();
  
  // 步骤3: 霍夫曼编码
  Serial.println(F("\n步骤3: 应用霍夫曼编码"));
  String huffmanBits;
  int huffmanLength = huffmanEncode(message, huffmanBits);
  
  // 步骤4: 汉明码编码（包含块数和填充位信息）
  Serial.println(F("\n步骤4: 汉明码(15,11)编码"));
  String hammingEncoded = applyHammingEncoding(huffmanBits);
  
  // 步骤5: 自适应PWM调制
  Serial.println(F("\n步骤5: 自适应PWM调制"));
  modulateData(hammingEncoded);
  
  Serial.println(F("==================== 处理完成 ====================\n"));
}

// ==================== Arduino主函数 ====================
void setup() {
  Serial.begin(115200);
  setupPWM();
  
  Serial.println(F("=== 改进的霍夫曼+汉明码(15,11)+自适应PWM调制系统 ==="));
  Serial.println(F("系统特性:"));
  Serial.println(F("- 支持ASCII字符输入（最多10个字符）"));
  Serial.println(F("- 动态生成霍夫曼编码"));
  Serial.println(F("- 汉明码(15,11)包含编码信息"));
  Serial.println(F("- 自适应调制位数选择"));
  Serial.println(F("- 显示完整的最终编码结果"));
  Serial.println(F("- 使用指定的占空比映射函数"));
  Serial.println(F("========================================"));
  Serial.println(F("请输入要处理的消息（最多10个字符）:"));
}

void loop() {
  // 检查串口输入
  if (Serial.available()) {
    String input = Serial.readStringUntil('\n');
    input.trim();
    
    if (input.length() > 0) {
      processMessage(input);
      Serial.println(F("\n请输入下一个消息:"));
    }
  }
  
  delay(100);
}