/*
  SCW-ROBUST: Shaped-CW Broadcast OOK Modem for Arduino Nano Every
  Author: ㆍArduino Expertㆍ

  Features:
   - Pure OOK keying of a CW transmitter (no FSK/PSK/AFSK).
   - Optional envelope shaping via PWM for click/splatter reduction.
   - Broadcast frames with preamble, Barker-13 sync, CRC-16, rate-1/2 conv. FEC, interleaving.
   - Profiles: 75 / 300 / 1200 bps.

  Pins (Nano Every default):
    KEY_PIN: Digital pin driving a keying transistor/optocoupler (active HIGH = key down).
    PWM_PIN: PWM output -> RC -> RF envelope shaper (optional). If not used, set USE_ENVELOPE_SHAPER = false.

  Serial:
    115200-8N1. Feed bytes; the modem packs into frames and transmits continuously.

  Safety/Compliance:
    - Observe local regulations. Keep symbol rates and rise/fall shaping conservative for spectral cleanliness.
*/

#include <Arduino.h>

//////////////////// USER CONFIG ////////////////////
#define PROFILE_MEDIUM_300   // choose one: PROFILE_ROBUST_75, PROFILE_MEDIUM_300, PROFILE_FAST_1200

const bool USE_ENVELOPE_SHAPER = true;   // true = use PWM envelope ramping on transitions
const bool USE_RETURN_TO_ZERO  = true;   // RZ (25% duty within symbol) improves timing at RX
const bool CONTINUOUS_MODE     = true;   // continuously transmit frames as data arrives

// Framing
const uint16_t FRAME_PAYLOAD_BYTES = 128;   // payload per frame before FEC/Interleave
const uint8_t  PREAMBLE_BITS       = 200;   // ≈ 200 ms if symbolRate=1k, scaled by actual symbol timing

// Pins (Nano Every)
const uint8_t KEY_PIN = 10;   // any digital pin; use transistor to key TX
const uint8_t PWM_PIN = 9;    // Timer PWM-capable pin for envelope (if used)

// If using Trinket M0 (SAMD21), you can #define TARGET_TRINKET_M0 and set pins accordingly.
//#define TARGET_TRINKET_M0
#ifdef TARGET_TRINKET_M0
  #undef KEY_PIN
  #undef PWM_PIN
  #define KEY_PIN  2
  #define PWM_PIN  1  // On M0 you could use DAC on A0 instead for smoother ramps
#endif
/////////////////////////////////////////////////////

// --------- Derived PHY parameters per profile -------
#if defined(PROFILE_ROBUST_75)
  const uint32_t SYMBOL_RATE_BPS = 75;
  const uint16_t EDGE_RAMP_MS    = 8;     // longer ramp = cleaner spectrum
  const uint8_t  RZ_PERCENT      = 25;    // 25% on-time within symbol
#elif defined(PROFILE_MEDIUM_300)
  const uint32_t SYMBOL_RATE_BPS = 300;
  const uint16_t EDGE_RAMP_MS    = 4;
  const uint8_t  RZ_PERCENT      = 25;
#elif defined(PROFILE_FAST_1200)
  const uint32_t SYMBOL_RATE_BPS = 1200;
  const uint16_t EDGE_RAMP_MS    = 2;
  const uint8_t  RZ_PERCENT      = 25;
#else
  #error "Select one profile: PROFILE_ROBUST_75 / PROFILE_MEDIUM_300 / PROFILE_FAST_1200"
#endif

// -------------- Protocol constants -----------------
static const uint16_t CRC16_POLY = 0xA001;  // CRC-16 (IBM) reversed poly
static const uint16_t CRC16_INIT = 0xFFFF;

// Barker-13 sync pattern (MSB first): 1111100110101
static const uint16_t BARKER13 = 0b1111100110101;
static const uint8_t  BARKER13_LEN = 13;

// Convolutional code (rate 1/2, K=7, polynomials in octal 171,133)
static const uint8_t CONV_K = 7;
static const uint8_t G1 = 0b1111001; // 171 octal = 0b1111001
static const uint8_t G2 = 0b1011011; // 133 octal = 0b1011011

// Interleaver size (bits after FEC per frame will be 2 * (header+payload+CRC)*8)
const uint16_t HEADER_BYTES = 4; // [version(1), flags(1), seq(1), length(1)]
const uint16_t MAX_FRAME_BYTES = HEADER_BYTES + FRAME_PAYLOAD_BYTES + 2; // + CRC16
const uint16_t MAX_INFO_BITS = MAX_FRAME_BYTES * 8; // before FEC
const uint16_t MAX_CODED_BITS = MAX_INFO_BITS * 2;  // rate 1/2
// Interleaver: choose square near sqrt(MAX_CODED_BITS)
const uint16_t INTERLEAVER_ROWS = 16;
const uint16_t INTERLEAVER_COLS = ( (MAX_CODED_BITS + INTERLEAVER_ROWS - 1) / INTERLEAVER_ROWS );

// -------------- Buffers ----------------------------
static uint8_t rxBuf[FRAME_PAYLOAD_BYTES * 2];  // incoming serial stash
static uint16_t rxCount = 0;

static uint8_t frameBuf[MAX_FRAME_BYTES];
static uint8_t encodedBytes[(MAX_CODED_BITS + 7)/8 + 4]; // packed coded + margin
static bool    codedBits[MAX_CODED_BITS + 64];           // boolean bitstream
static bool    interBits[INTERLEAVER_ROWS * INTERLEAVER_COLS + 64];

// Timing
static uint32_t symbolPeriodMicros;
static uint16_t onTimeMicros;  // for RZ on-time
static uint16_t rampSteps;     // number of steps in ramp
static uint16_t rampStepDelayMicros;

// State
static uint8_t txSeq = 0;

// ---------- Utility: CRC16-IBM (Modbus variant) ----------
uint16_t crc16_ibm(const uint8_t* data, uint16_t len) {
  uint16_t crc = CRC16_INIT;
  for (uint16_t i = 0; i < len; i++) {
    crc ^= data[i];
    for (uint8_t b = 0; b < 8; b++) {
      if (crc & 1) crc = (crc >> 1) ^ CRC16_POLY;
      else         crc = crc >> 1;
    }
  }
  return crc;
}

// ---------- Bit pack/unpack helpers ----------
inline void bitsWrite(bool* dst, uint32_t idx, bool bit) {
  dst[idx] = bit;
}
inline bool getBit(const uint8_t* src, uint32_t bitIndex) {
  return (src[bitIndex >> 3] >> (7 - (bitIndex & 7))) & 1;
}
inline void setBit(uint8_t* dst, uint32_t bitIndex, bool v) {
  if (v) dst[bitIndex >> 3] |= (1 << (7 - (bitIndex & 7)));
  else   dst[bitIndex >> 3] &= ~(1 << (7 - (bitIndex & 7)));
}

// ---------- Convolutional Encoder (rate 1/2, K=7) ----------
uint32_t convEncode(const uint8_t* in, uint16_t inBits, bool* outBits) {
  uint8_t shift = 0; // K-1 = 6 zeros initial
  uint32_t outIdx = 0;
  for (uint16_t i = 0; i < inBits; i++) {
    bool u = getBit(in, i);
    shift = ((shift << 1) | (u ? 1 : 0)) & 0x7F; // keep 7 bits
    // parity for G1, G2 (LSB of shift is newest)
    uint8_t v1 = __builtin_parity(shift & G1);
    uint8_t v2 = __builtin_parity(shift & G2);
    outBits[outIdx++] = v1;
    outBits[outIdx++] = v2;
  }
  // Tail bits (flush trellis): append K-1 zeros
  for (uint8_t t = 0; t < CONV_K - 1; t++) {
    shift = (shift << 1) & 0x7F;
    uint8_t v1 = __builtin_parity(shift & G1);
    uint8_t v2 = __builtin_parity(shift & G2);
    outBits[outIdx++] = v1;
    outBits[outIdx++] = v2;
  }
  return outIdx;
}

// ---------- Block Interleaver ----------
uint32_t interleave(const bool* inBits, uint32_t inLen, bool* outBits) {
  // write row-wise, read column-wise
  uint16_t rows = INTERLEAVER_ROWS;
  uint16_t cols = ( (inLen + rows - 1) / rows );
  for (uint16_t r = 0; r < rows; r++) {
    for (uint16_t c = 0; c < cols; c++) {
      uint32_t i = r * cols + c;
      uint32_t o = c * rows + r;
      bool v = (i < inLen) ? inBits[i] : 0;
      outBits[o] = v;
    }
  }
  return rows * cols;
}

// ---------- Preamble & Barker Sync emission helpers ----------
void emitPreamble(void (*emitBit)(bool)) {
  // Alternate 1/0 for PREAMBLE_BITS
  for (uint16_t i = 0; i < PREAMBLE_BITS; i++) {
    emitBit((i & 1) != 0);
  }
}
void emitBarker13(void (*emitBit)(bool)) {
  for (int8_t i = BARKER13_LEN - 1; i >= 0; i--) {
    bool b = (BARKER13 >> i) & 1;
    emitBit(b);
  }
}

// ---------- Envelope control (key + optional PWM ramp) ----------
inline void rfKeyDownRaw()   { digitalWrite(KEY_PIN, HIGH); }
inline void rfKeyUpRaw()     { digitalWrite(KEY_PIN, LOW);  }

// Simple PWM ramp on transitions for envelope shaping.
// If not used, we just toggle KEY_PIN with delays.
void envelopeKey(bool keyOn) {
  if (!USE_ENVELOPE_SHAPER) {
    if (keyOn) rfKeyDownRaw();
    else       rfKeyUpRaw();
    return;
  }

  // Using PWM duty ramp to approximate raised-cosine edges; requires external RC & RF gain element.
  static uint8_t last = 0; // last duty
  if (keyOn) {
    // ramp up
    for (uint16_t s = 0; s < rampSteps; s++) {
      // raised-cosine-ish lookup without float
      uint16_t duty = (uint32_t)(s + 1) * 255 / rampSteps;
      analogWrite(PWM_PIN, duty);
      last = duty;
      delayMicroseconds(rampStepDelayMicros);
    }
    rfKeyDownRaw(); // Assert key line once envelope is up
  } else {
    // release key before ramp down helps avoid chirp in some TX chains
    rfKeyUpRaw();
    for (uint16_t s = 0; s < rampSteps; s++) {
      uint16_t duty = (uint32_t)(rampSteps - s - 1) * 255 / rampSteps;
      analogWrite(PWM_PIN, duty);
      last = duty;
      delayMicroseconds(rampStepDelayMicros);
    }
  }
}

// Emit one OOK bit, with optional RZ sub-gating for transition density.
// NRZ: keyOn for '1', keyOff for '0'.
// RZ: If bit=1, turn on for RZ_PERCENT of symbol, then off for remainder.
//     If bit=0, keep off full symbol.
void emitBitTimed(bool bit) {
  uint32_t t0 = micros();

  if (USE_RETURN_TO_ZERO) {
    if (bit) {
      // turn on with shaped edge
      envelopeKey(true);
      delayMicroseconds(onTimeMicros);
      envelopeKey(false);
      // remainder of symbol
      uint32_t elapsed = micros() - t0;
      uint32_t remain = (elapsed >= symbolPeriodMicros) ? 0 : (symbolPeriodMicros - elapsed);
      if (remain > 0) delayMicroseconds(remain);
    } else {
      // keep off full symbol
      envelopeKey(false);
      // busy-wait to symbol boundary
      uint32_t elapsed = micros() - t0;
      if (elapsed < symbolPeriodMicros) delayMicroseconds(symbolPeriodMicros - elapsed);
    }
  } else {
    // NRZ: key is on for '1' entire symbol, off for '0'
    envelopeKey(bit);
    uint32_t elapsed = micros() - t0;
    if (elapsed < symbolPeriodMicros) delayMicroseconds(symbolPeriodMicros - elapsed);
  }
}

// High-level bit emitter function pointer wrapper
void emitBit(bool b) {
  emitBitTimed(b);
}

// ---------- Frame builder and transmitter ----------
void buildAndSendFrame(const uint8_t* payload, uint16_t len) {
  if (len > FRAME_PAYLOAD_BYTES) len = FRAME_PAYLOAD_BYTES;

  // Build header
  uint8_t* p = frameBuf;
  p[0] = 1;           // version
  p[1] = 0;           // flags (reserved)
  p[2] = txSeq++;     // sequence
  p[3] = (uint8_t)len;
  memcpy(p + HEADER_BYTES, payload, len);

  // CRC over header+payload
  uint16_t usedBytes = HEADER_BYTES + len;
  uint16_t crc = crc16_ibm(frameBuf, usedBytes);
  frameBuf[usedBytes] = (uint8_t)(crc & 0xFF);
  frameBuf[usedBytes + 1] = (uint8_t)(crc >> 8);
  usedBytes += 2;

  // FEC encode
  uint16_t inBits = usedBytes * 8;
  uint32_t codedLen = convEncode(frameBuf, inBits, codedBits);

  // Interleave
  uint32_t interLen = interleave(codedBits, codedLen, interBits);

  // Transmit
  emitPreamble(emitBit);
  emitBarker13(emitBit);

  // Emit bits
  for (uint32_t i = 0; i < interLen; i++) {
    emitBit(interBits[i]);
  }
}

// ---------- Serial buffer management ----------
void pumpSerialToFrames() {
  // Fill rxBuf
  while (Serial.available() && rxCount < sizeof(rxBuf)) {
    rxBuf[rxCount++] = (uint8_t)Serial.read();
  }

  // While we have enough bytes for a frame payload, send one
  while (rxCount >= FRAME_PAYLOAD_BYTES) {
    buildAndSendFrame(rxBuf, FRAME_PAYLOAD_BYTES);

    // --- ADDED CONFIRMATION ---
    Serial.println(F("-> Transmitted full frame (128 bytes)."));

    // shift remaining
    uint16_t rem = rxCount - FRAME_PAYLOAD_BYTES;
    memmove(rxBuf, rxBuf + FRAME_PAYLOAD_BYTES, rem);
    rxCount = rem;

    if (!CONTINUOUS_MODE) break;
  }

  // If we have any leftover and in continuous mode, you can choose to flush a short frame.
  // Here we flush if idle for > ~20 ms without new data (simple heuristic).
  static uint32_t lastActivity = 0;
  uint32_t now = millis();
  if (Serial.available()) lastActivity = now;
  if (CONTINUOUS_MODE && rxCount > 0) {
    if (now - lastActivity > 20) {
      uint16_t bytesToSend = rxCount; // Store count before it's cleared
      buildAndSendFrame(rxBuf, rxCount);
      rxCount = 0;

      // --- ADDED CONFIRMATION ---
      Serial.print(F("-> Flushed buffer. Transmitted short frame ("));
      Serial.print(bytesToSend);
      Serial.println(F(" bytes)."));
    }
  }
}


// ---------- Setup & Loop ----------
void setupTimersAndPWM() {
  symbolPeriodMicros = (uint32_t)(1000000UL / SYMBOL_RATE_BPS);
  onTimeMicros = (uint16_t)((uint32_t)symbolPeriodMicros * RZ_PERCENT / 100);
  if (!USE_RETURN_TO_ZERO) onTimeMicros = symbolPeriodMicros;

  // Ramping: allocate ~EDGE_RAMP_MS time to ramps.
  if (USE_ENVELOPE_SHAPER) {
    rampSteps = 32; // 32-step ramp
    uint32_t totalRampMicros = (uint32_t)EDGE_RAMP_MS * 1000UL;
    rampStepDelayMicros = (totalRampMicros / rampSteps);
    // initialize PWM
    analogWrite(PWM_PIN, 0);
  }
}

void setup() {
  pinMode(KEY_PIN, OUTPUT);
  digitalWrite(KEY_PIN, LOW);
  pinMode(PWM_PIN, OUTPUT);

  Serial.begin(115200);
  while (!Serial) { ; }

  setupTimersAndPWM();

  Serial.println(F("SCW-ROBUST OOK Modem Ready."));
  Serial.print(F("SymbolRate=")); Serial.print(SYMBOL_RATE_BPS);
  Serial.print(F(" bps, RZ=")); Serial.print((int)RZ_PERCENT);
  Serial.print(F("%, Ramp=")); Serial.print((int)EDGE_RAMP_MS); Serial.println(F(" ms"));
  Serial.println(F("Send bytes via Serial. Broadcasting frames..."));
}

void loop() {
  pumpSerialToFrames();
}