/*
 * NANO DRUM ENGINE – BLAST BEAT KICK + SNARE v9.3
 * Piezo ACTIVE-LOW version (trigger when signal goes LOW)
 * Kick pad  → A1
 * Snare pad → A0
 */

#include <SPI.h>
#include <SD.h>
#include <TMRpcm.h>
#include <TM1637Display.h>

// ── Debug & Logging ──────────────────────────────────────────────────────────
#define DEBUG                1
#define DEBUG_ENVELOPE       0

#if DEBUG
  #define LOG_INIT()         Serial.begin(115200)
  #define LOG_THROTTLE(msg, val, interval) do { \
    static unsigned long _last = 0; \
    if (millis() - _last >= interval) { \
      Serial.print(F(msg)); Serial.println(val); \
      _last = millis(); \
    } \
  } while(0)
  #define LOG_HIT(drum, vel, en, lk) do { \
    Serial.print(F("HIT | ")); Serial.print(drum); \
    Serial.print(F(" V:")); Serial.print(vel); \
    Serial.print(F(" E:")); Serial.print(en); \
    Serial.print(F(" L:")); Serial.println(lk); \
  } while(0)
#else
  #define LOG_INIT()
  #define LOG_THROTTLE(msg,val,interval)
  #define LOG_HIT(drum,vel,en,lk)
#endif

// ── Pin Definitions ──────────────────────────────────────────────────────────
#define KICK_PIEZO_PIN     A1
#define SNARE_PIEZO_PIN    A0
#define POT_THRESHOLD      A2
#define POT_LOCK           A3
#define SD_CS_PIN          4

#define AUDIO_PWM_PIN      9

#define SR_CLOCK_PIN       2
#define SR_LATCH_PIN       3
#define SR_DATA_PIN        7

#define TM1637_CLK         5
#define TM1637_DIO         6

TMRpcm tmrpcm;
TM1637Display tm(TM1637_CLK, TM1637_DIO);

// ── Tunable Constants ────────────────────────────────────────────────────────
const int THRESHOLD_MIN     = 15;
const int THRESHOLD_MAX     = 140;

const int LOCK_MIN_MS       = 35;
const int LOCK_MAX_MS       = 80;
const int CANCEL_MIN_MS     = 25;
const int CANCEL_MAX_MS     = 60;

const unsigned long SCAN_US         = 700;
const int       SLOPE_MIN           = 20;
const long      ENERGY_THRESHOLD    = 2600;
const uint8_t   ENVELOPE_DECAY_SHIFT = 3;
const int       MIN_VELOCITY        = 5;

const int       REBOUND_MASK_MS     = 40;
const int       REBOUND_THRESHOLD   = 25;

// ── Velocity Curve LUT ───────────────────────────────────────────────────────
const uint8_t velocityCurve[16] PROGMEM = {
     0,   6,  14,  24,  35,  47,  58,  68,
    77,  84,  89,  93,  96,  98,  99, 100
};

// ── Per-drum state ───────────────────────────────────────────────────────────
struct Drum {
  int     pin;
  int     noiseFloor       = 512;
  int     baseThreshold    = 460;
  int     dynamicThreshold = 460;
  int     lastRaw          = 512;
  int     envelope         = 0;
  int     lastPeakEnv      = 0;
  bool    triggerArmed     = true;
  unsigned long lastHitMillis     = 0;
  unsigned long cancelUntil       = 0;
  unsigned long reboundIgnoreUntil = 0;
  int     lastDip          = 1023;
  int     currentVelocity  = 0;
};

Drum kick;
Drum snare;

int   retriggerLockMs   = 50;
int   cancelWindowMs    = 35;

unsigned long lastPotRead       = 0;
unsigned long lastDisplayUpdate = 0;

// ── Functions ────────────────────────────────────────────────────────────────
void updateLedBar(int velocity) {
  unsigned long now = millis();
  if (now - max(kick.lastHitMillis, snare.lastHitMillis) < 50) {
    int leds = map(velocity, 0, 100, 1, 10);
    displayBar((1UL << leds) - 1);
  } else if (now - max(kick.lastHitMillis, snare.lastHitMillis) < 350) {
    int elapsed = now - max(kick.lastHitMillis, snare.lastHitMillis) - 50;
    int level = map(elapsed, 0, 300, 10, 0);
    int leds = map(level, 0, 10, 0, map(velocity, 0, 100, 1, 10));
    displayBar((1UL << leds) - 1);
  } else {
    displayBar(0);
  }
}

inline void displayBar(uint16_t mask) {
  digitalWrite(SR_LATCH_PIN, LOW);
  shiftOut(SR_DATA_PIN, SR_CLOCK_PIN, MSBFIRST, (mask >> 8) & 0xFF);
  shiftOut(SR_DATA_PIN, SR_CLOCK_PIN, MSBFIRST, mask & 0xFF);
  digitalWrite(SR_LATCH_PIN, HIGH);
}

void updateDisplay() {
  unsigned long now = millis();
  if (now - lastDisplayUpdate < 300) return;
  lastDisplayUpdate = now;

  int latestVel = max(kick.currentVelocity, snare.currentVelocity);
  unsigned long latestHit = max(kick.lastHitMillis, snare.lastHitMillis);

  if (latestVel > 0 && now - latestHit < 1500) {
    tm.showNumberDec(latestVel, false);
  } else {
    static bool showThresh = true;
    if (showThresh) {
      tm.showNumberDec(kick.baseThreshold, true);
      tm.setSegments((const uint8_t[]){0x71, 0x68, 0x00, 0x00}, false, 0); // t h
    } else {
      tm.showNumberDec(retriggerLockMs, true);
      tm.setSegments((const uint8_t[]){0x38, 0x39, 0x00, 0x00}, false, 0); // L o
    }
    showThresh = !showThresh;
  }
}

int applyVelocityCurve(float normalized) {
  if (normalized <= 0.0f) return 0;
  if (normalized >= 1.0f) return 100;

  float indexF = normalized * 15.0f;
  uint8_t idx = (uint8_t)indexF;
  float frac = indexF - (float)idx;

  uint8_t low  = pgm_read_byte(&velocityCurve[idx]);
  uint8_t high = pgm_read_byte(&velocityCurve[idx + 1]);

  return low + (int)((high - low) * frac + 0.5f);
}

void processDrum(Drum &drum, const char* name, const char* soft, const char* med, const char* hard) {
  unsigned long now = millis();

  if (now < drum.cancelUntil || now < drum.reboundIgnoreUntil) {
    analogRead(drum.pin); // dummy
    return;
  }

  int raw = analogRead(drum.pin);

  // Envelope (how much BELOW idle)
  int deviation = drum.noiseFloor - raw;
  if (deviation > drum.envelope) {
    drum.envelope = deviation;
  } else {
    drum.envelope -= drum.envelope >> ENVELOPE_DECAY_SHIFT;
  }

  int slope = drum.lastRaw - raw;
  drum.lastRaw = raw;

  if (drum.triggerArmed && slope > SLOPE_MIN && drum.envelope > drum.dynamicThreshold) {
    if (now - drum.lastHitMillis > (unsigned long)retriggerLockMs) {

      int deepest = raw;
      long energy = 0UL;
      unsigned long start = micros();

      while (micros() - start < SCAN_US) {
        int r = analogRead(drum.pin);
        if (r < deepest) deepest = r;
        long d = (long)drum.noiseFloor - r;
        energy += d * d;
      }

      energy = (energy * 100) / (drum.lastPeakEnv > 0 ? drum.lastPeakEnv : 1);

      if (energy > ENERGY_THRESHOLD) {
        float howMuchBelow = (float)(drum.noiseFloor + drum.dynamicThreshold - deepest);
        float range        = (float)(drum.noiseFloor + drum.dynamicThreshold - 10);
        float normalized   = howMuchBelow / range;
        normalized = constrain(normalized, 0.0f, 1.0f);

        drum.currentVelocity = applyVelocityCurve(normalized);

        if (drum.currentVelocity >= MIN_VELOCITY) {
          const char* sample;
               if (drum.currentVelocity < 35) sample = soft;
          else if (drum.currentVelocity < 75) sample = med;
          else                               sample = hard;

          tmrpcm.play(sample);

          updateLedBar(drum.currentVelocity);
          LOG_HIT(name, drum.currentVelocity, energy, retriggerLockMs);

          drum.lastHitMillis    = now;
          drum.lastPeakEnv      = drum.envelope;
          drum.lastDip          = deepest;
          drum.triggerArmed     = false;

          drum.dynamicThreshold = drum.baseThreshold + (drum.envelope >> 2);

          drum.cancelUntil      = now + cancelWindowMs;
          drum.reboundIgnoreUntil = now + REBOUND_MASK_MS;

          if (!tmrpcm.isPlaying()) digitalWrite(AUDIO_PWM_PIN, LOW);
        }
      }
    }
  }

  // Re-arm
  if (!drum.triggerArmed && drum.envelope < (drum.lastPeakEnv >> (ENVELOPE_DECAY_SHIFT + 1))) {
    drum.triggerArmed = true;
    drum.dynamicThreshold = drum.baseThreshold;
  }
}

// ── Setup ────────────────────────────────────────────────────────────────────
void setup() {
  LOG_INIT();

  // Faster ADC
  ADCSRA |= (1 << ADPS2); ADCSRA &= ~(1 << ADPS1); ADCSRA &= ~(1 << ADPS0);

  pinMode(SR_DATA_PIN, OUTPUT);
  pinMode(SR_LATCH_PIN, OUTPUT);
  pinMode(SR_CLOCK_PIN, OUTPUT);
  pinMode(AUDIO_PWM_PIN, OUTPUT);
  digitalWrite(AUDIO_PWM_PIN, LOW);

  tm.setBrightness(0x0f);
  tm.clear();

  if (!SD.begin(SD_CS_PIN)) {
    LOG_THROTTLE("SD failed", 0, 0);
    displayBar(0b1111111111);
    while (true);
  }

  tmrpcm.speakerPin = AUDIO_PWM_PIN;
  tmrpcm.setVolume(4);
  tmrpcm.quality(0);

  // Calibrate both piezos
  long sumKick = 0, sumSnare = 0;
  for (uint8_t i = 0; i < 120; i++) {
    sumKick += analogRead(KICK_PIEZO_PIN);
    sumSnare += analogRead(SNARE_PIEZO_PIN);
    delay(3);
  }
  kick.noiseFloor  = sumKick  / 120;
  snare.noiseFloor = sumSnare / 120;

  LOG_THROTTLE("Kick idle: ", kick.noiseFloor, 0);
  LOG_THROTTLE("Snare idle:", snare.noiseFloor, 0);

  tm.showNumberDec(kick.noiseFloor, false);
  delay(600);
  tm.showNumberDec(snare.noiseFloor, false);
  delay(600);
  tm.clear();
}

// ── Main Loop ────────────────────────────────────────────────────────────────
void loop() {
  unsigned long now = millis();

  // Potentiometers (shared settings for now)
  if (now - lastPotRead >= 250) {
    int tRaw = analogRead(POT_THRESHOLD);
    int mappedThresh = map(tRaw, 0, 1023, THRESHOLD_MIN, THRESHOLD_MAX);

    int lRaw = analogRead(POT_LOCK);
    retriggerLockMs = map(lRaw, 0, 1023, LOCK_MIN_MS, LOCK_MAX_MS);
    cancelWindowMs  = map(lRaw, 0, 1023, CANCEL_MIN_MS, CANCEL_MAX_MS);

    // Apply to both drums
    kick.baseThreshold  = mappedThresh;
    snare.baseThreshold = mappedThresh;

    if (kick.triggerArmed)  kick.dynamicThreshold  = kick.baseThreshold;
    if (snare.triggerArmed) snare.dynamicThreshold = snare.baseThreshold;

    lastPotRead = now;
  }

  // Process each drum
  processDrum(kick,  "KICK ", "kick_soft.wav",   "kick_medium.wav",   "kick_hard.wav");
  processDrum(snare, "SNARE", "snare_soft.wav",  "snare_medium.wav",  "snare_hard.wav");

  // UI
  updateLedBar(max(kick.currentVelocity, snare.currentVelocity));
  updateDisplay();
}