#include "AESLib.h"
#include <arduinoFFT.h>
#include <Arduino_KNN.h>

extern char __bss_end; 
extern char *__brkval; 

#define SAMPLES 16        // 必须是2的幂，FFT的采样点数
#define SAMPLING_FREQUENCY 10000 // 采样频率，单位 Hz

const int micPin = A0;     // 麦克风连接的引脚

float vReal[SAMPLES];     // 实部数组
float vImag[SAMPLES];     // 虚部数组
ArduinoFFT<float> FFT(vReal, vImag, SAMPLES, SAMPLING_FREQUENCY);

KNNClassifier knn(3);
AESLib aesLib;

// 预设的明文
char cleartext[64];
char ciphertext[128]; // 存储加密后的密文



// AES 加密密钥
byte aes_key[] = { 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30 };

// 初始化向量（IV），在实际生产环境中应随机生成并存储
byte aes_iv[N_BLOCK] = { 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 };

String encrypt_impl(char *msg, byte iv[]) {
  int msgLen = strlen(msg);
  char *encrypted = new char[2 * msgLen]();
  aesLib.encrypt64((const byte*)msg, msgLen, encrypted, aes_key, sizeof(aes_key), iv);
  return String(encrypted);
}

String decrypt_impl(char *msg, byte iv[]) {
  int msgLen = strlen(msg);
  char *decrypted = new char[msgLen]();
  aesLib.decrypt64(msg, msgLen, (byte*)decrypted, aes_key, sizeof(aes_key), iv);
  return String(decrypted);
}

// 初始化 AES 加密
void aes_init() {
  Serial.println("\nGenerating IV...");
  aesLib.gen_iv(aes_iv);
}
unsigned long loopcount = 0;
byte enc_iv[N_BLOCK] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };
byte dec_iv[N_BLOCK] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };

void Encrypt_setup(){
  aes_init();
  aesLib.set_paddingmode(paddingMode::CMS);
  char b64in[16] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};

  char b64out[base64_enc_len(sizeof(aes_iv))];
  base64_encode(b64out, b64in, 16);

  char b64enc[base64_enc_len(10)];
  base64_encode(b64enc, (char*) "0123456789", 10);

  char b64dec[ base64_dec_len(b64enc, sizeof(b64enc))];
  base64_decode(b64dec, b64enc, sizeof(b64enc));
  Serial.println("Starting AES encryption and decryption example...");
  
  // Encrypt
  String encrypted = encrypt_impl(cleartext, enc_iv);
  sprintf(ciphertext, "%s", encrypted.c_str());
  Serial.print("Ciphertext: ");
  Serial.println(encrypted);
  delay(1000);

  // Decrypt
  delay(1000);
  String decrypted = decrypt_impl( ciphertext, dec_iv);
  Serial.print("Cleartext: ");
  //decrypted.trim();
  Serial.println(decrypted);

  String cleartextstr = String(cleartext);
  if (decrypted.substring(0, cleartextstr.length()).equals(cleartextstr)) {
    Serial.println("SUCCES\n");
  } else {
    Serial.println("FAILURE\n");
  }
  for (int i = 0; i < 16; i++) {
    enc_iv[i] = 0;
    dec_iv[i] = 0;
  }
}

void setup() {
  Serial.begin(9600);
  Serial.println("FFT and kNN Classifier Start...");

  // 添加训练数据
  float sample1[] = {8500, 8600, 8700};
  float sample2[] = {8600, 8600, 8600};
  float sample3[] = {9000, 9000, 9000};


  knn.addExample(sample1, 1);
  knn.addExample(sample2, 2);
  knn.addExample(sample3, 3);
}

// 采集指定频率的 FFT 特征
void collectFFTFeatures(float* features) {
  const float targetFrequencies[3] = {100, 200, 300};
  for (int f = 0; f < 3; f++) {
    for (int i = 0; i < SAMPLES; i++) {
      vReal[i] = analogRead(micPin);  // 从麦克风读取信号
      vImag[i] = 0;                  // 初始化虚部为 0
      delayMicroseconds(1000000 / SAMPLING_FREQUENCY); // 控制采样频率
    }

    // 执行 FFT 计算
    FFT.windowing(vReal, SAMPLES, FFT_WIN_TYP_HAMMING, FFT_FORWARD); // 加窗处理
    FFT.compute(vReal, vImag, SAMPLES, FFT_FORWARD);                // 计算 FFT
    FFT.complexToMagnitude(vReal, vImag, SAMPLES);                  // 转换为幅值

    // 找到目标频率对应的幅值
    int index = (targetFrequencies[f] * SAMPLES) / SAMPLING_FREQUENCY; // 频率对应的索引
    features[f] = vReal[index]; // 保存目标频率的幅值
  }
}

int freeMemory() {
  char top; // 栈的当前地址
  if (__brkval == 0) {
    return (&top - &__bss_end); 
  } else {
    return (&top - __brkval); 
  }
}
// 初始化 kNN 分类器，k=3
void loop() {
  // 测试采样时间
  unsigned long startTime = micros();
  for (int i = 0; i < SAMPLES; i++) {
      vReal[i] = analogRead(A0);
      delayMicroseconds(1000000 / SAMPLING_FREQUENCY);
  }
  unsigned long samplingTime = micros() - startTime;
  Serial.print("\n=== Sampling Time: ");
  Serial.print(samplingTime);
  Serial.println(" us");

  // 实时采集 FFT 特征
  float feature[3]; // 修改为 3 个特征值数组，适应训练样本维度
  collectFFTFeatures(feature); // 从特定频率采集特征


  unsigned long fftTime = micros() - startTime;
  Serial.print("FFT Computation Time: ");
  Serial.print(fftTime);
  Serial.println(" us");

  int classification = knn.classify(feature); // 使用 kNN 进行分类

  Serial.print("\nFeature: ");
  for (int i = 0; i < 3; i++) {
    Serial.print(feature[i]);
    Serial.print(" ");
  }
  String result = String(feature[0], 2) + " " + String(feature[1], 2) + " " + String(feature[2], 2);
  
  Serial.print("=> Class: ");
  Serial.println(classification);

  unsigned long classificationTime = micros() - startTime;
  Serial.print("Classification Time: ");
  Serial.print(classificationTime);
  Serial.println(" us");

  strcpy(cleartext, result.c_str());

  Encrypt_setup();
  
  unsigned long encryptionTime = micros() - startTime;
  Serial.print("Encryption Time: ");
  Serial.print(encryptionTime);
  Serial.println(" us");
  Serial.print("Remaining stack for ");
  Serial.println(freeMemory());
  delay(1000); // 每秒更新一次
}