-- @description Singtall Drum Align
-- @version 1.0
-- @author Singtall Studio
-- @website
--   Singtall Drum Align https://singtallstudio.com/drum-align.html
--   Singtall Studio https://singtallstudio.com
-- @changelog
--   First release.
-- @about
--   Time- and polarity-aligns drum tracks to a master track (usually the snare)
--   by physically moving the items. No plugin, no render, the audio is untouched.
--
--   1. Select all the drum tracks (click their names in the track panel).
--   2. Run the script. The master is the selected track named like a snare
--      (snare top preferred), otherwise the topmost selected track.
--   3. Check the report and press OK. One undo step puts everything back.
--
--   The script picks one clean snare hit itself (the one where the most mics
--   agree with the rest of the song), lines every track up on it, flips
--   polarity where needed, and drops a "Drum Align" marker on the crest of that
--   hit so you can zoom in and check. Your own markers are never touched.
--
--   Use it at the start of a mix, before any plugins. Then run Singtall Drum
--   Clean to cut the bleed between kick and tom hits.

local SCRIPT_NAME = "Singtall Drum Align"
local VERSION     = "1.0"
local WEBSITE     = "singtallstudio.com"

local CFG = {
  maxLagMs    = 25,     -- furthest a mic can be from the master (25 ms is about 8.5 m)
  maxHits     = 40,     -- strongest master hits used
  minHitGapMs = 150,    -- two hits closer than this count as one
  hitRiseDb   = 12,     -- a hit is a jump of this much over the last ~20 ms
  hitRangeDb  = 10,     -- keep hits within this of the strong ones (drops kick/hat bleed)
  onsetFrac   = 0.25,   -- a mic's snare starts with the first wave above this share of its peak
  agreeMs     = 0.2,    -- a hit agrees with the song when it is this close to the track's median
  minAgree    = 0.3,    -- tracks with fewer agreeing hits than this share get "check by ear"
  minHitsAgree = 5,     -- and the agreeing hits must be at least this many
  minMoveMs   = 0.05,   -- smaller offsets than this are left alone (about 2 samples)
  chunkBlocks = 700,    -- audio read in chunks of this many 64-sample blocks
}

local TEST = SINGTALL_DRUMALIGN_TEST   -- set by the test rig only

-- Each run overwrites a short log next to the script, so a run that seems to
-- do nothing can be diagnosed afterwards.
local LOG_PATH = (({ reaper.get_action_context() })[2] or ""):match("^(.*[\\/])") or (reaper.GetResourcePath() .. "/Scripts/")
LOG_PATH = LOG_PATH .. "Singtall Drum Align log.txt"
local logFile = io.open(LOG_PATH, "w")
local function log(s)
  if logFile then logFile:write(os.date("%H:%M:%S  ") .. s .. "\n"); logFile:flush() end
end

---------------------------------------------------------------------------
-- Audio reading
---------------------------------------------------------------------------
local NCH = 2
local readBuf, readBufSize

-- Mono (L+R average) samples of a track from project time t0, n samples long.
local function readMono(acc, srate, t0, n)
  if not readBuf or readBufSize < n * NCH then
    readBufSize = n * NCH
    readBuf = reaper.new_array(readBufSize)
  end
  readBuf.clear()
  reaper.GetAudioAccessorSamples(acc, srate, NCH, t0, n, readBuf)
  local t = readBuf.table(1, n * NCH)
  local out = {}
  for i = 1, n do out[i] = 0.5 * (t[2 * i - 1] + t[2 * i]) end
  return out
end

---------------------------------------------------------------------------
-- Hit detection on the master track
---------------------------------------------------------------------------
local function findHits(acc, srate, rs, re)
  local B = 64
  local env, nBlocks = {}, 0
  local chunk = B * CFG.chunkBlocks
  local t = rs
  while t < re do
    local s = readMono(acc, srate, t, chunk)
    for b = 0, CFG.chunkBlocks - 1 do
      local m = 0
      for i = b * B + 1, b * B + B do
        local v = s[i]; if v < 0 then v = -v end
        if v > m then m = v end
      end
      nBlocks = nBlocks + 1
      env[nBlocks] = m
    end
    t = t + chunk / srate
  end
  local maxBlocks = math.floor((re - rs) * srate / B)
  if nBlocks > maxBlocks then nBlocks = maxBlocks end

  local gmax = 0
  for i = 1, nBlocks do if env[i] > gmax then gmax = env[i] end end
  if gmax <= 1e-6 then return {} end

  local rise  = 10 ^ (CFG.hitRiseDb / 20)
  local floor = gmax * 10 ^ (-24 / 20)
  local gapB  = math.floor(CFG.minHitGapMs / 1000 * srate / B)
  local look  = 16
  local cands, last = {}, -1e9
  for i = look + 1, nBlocks - 8 do
    if i - last >= gapB and env[i] >= floor then
      local slow = 0
      for j = i - look, i - 1 do slow = slow + env[j] end
      slow = slow / look
      if env[i] >= slow * rise then
        local pk = 0
        for j = i, i + 8 do if env[j] > pk then pk = env[j] end end
        cands[#cands + 1] = { time = rs + (i - 1) * B / srate, peak = pk }
        last = i
      end
    end
  end
  if #cands == 0 then return {} end

  -- Keep only hits near the strong ones (the 5th strongest is the reference,
  -- so one freak peak can't set it), then the strongest maxHits of those.
  table.sort(cands, function(a, b) return a.peak > b.peak end)
  local ref = cands[math.min(5, #cands)].peak
  local keepAbove = ref * 10 ^ (-CFG.hitRangeDb / 20)
  local hits = {}
  for _, c in ipairs(cands) do
    if c.peak >= keepAbove and #hits < CFG.maxHits then hits[#hits + 1] = c end
  end
  table.sort(hits, function(a, b) return a.time < b.time end)
  return hits
end

---------------------------------------------------------------------------
-- Onset analysis: where each mic's snare starts, the way it is done by eye
---------------------------------------------------------------------------
local function maxAbs(t)
  local m = 0
  for i = 1, #t do local v = math.abs(t[i]); if v > m then m = v end end
  return m
end

-- Start of the first wave that rises clearly out of the background: the first
-- sample above both onsetFrac of the window's peak and twice the noise before
-- it, walked back to the zero crossing where that wave began. Returns the start
-- (fractional samples from the window start, 0-based) and the wave's direction.
local function firstWave(w, noise)
  local peak = maxAbs(w)
  if peak <= 0 then return nil end
  local thr = math.max(CFG.onsetFrac * peak, 2 * noise)
  local i = 1
  while i <= #w and math.abs(w[i]) <= thr do i = i + 1 end
  if i > #w then return nil end
  local up = w[i] > 0
  local j = i
  while j > 1 and w[j - 1] ~= 0 and (w[j - 1] > 0) == up do j = j - 1 end
  local z = j - 1
  if j > 1 and w[j] ~= w[j - 1] then z = (j - 2) + w[j - 1] / (w[j - 1] - w[j]) end
  return z, up and 1 or -1
end

-- Onset (project time) and direction of a track's first wave near time t. The
-- search starts `before` seconds ahead of t and runs `len` seconds; background
-- noise is measured over the 10 ms ending 1 ms before the search starts.
local function onsetNear(acc, srate, t, before, len)
  local noise = maxAbs(readMono(acc, srate, t - before - 0.011, math.floor(0.010 * srate)))
  local z, s = firstWave(readMono(acc, srate, t - before, math.floor(len * srate)), noise)
  if not z then return nil end
  return t - before + z / srate, s
end

-- One follower against the master's hits: its onset delay and polarity at each
-- hit, plus the median over all hits, which tells a clean hit from a messy one.
local function analyseTrack(track, master, srate)
  local acc = reaper.CreateTrackAudioAccessor(track)
  local len = (5 + CFG.maxLagMs + 5) / 1000
  local byHit, list = {}, {}
  for h, hit in ipairs(master.hits) do
    local on, s = onsetNear(acc, srate, hit.on, 0.005, len)
    if on then
      local p = { lag = (on - hit.on) * srate, sign = s * hit.sign }
      byHit[h] = p
      list[#list + 1] = p
    end
  end
  reaper.DestroyAudioAccessor(acc)
  if #list == 0 then return { silent = true } end

  local lags, up = {}, 0
  for _, p in ipairs(list) do
    lags[#lags + 1] = p.lag
    if p.sign > 0 then up = up + 1 end
  end
  table.sort(lags)
  local n = #lags
  local med = (n % 2 == 1) and lags[(n + 1) // 2] or 0.5 * (lags[n // 2] + lags[n // 2 + 1])
  local sign = (up * 2 >= n) and 1 or -1
  local tol = CFG.agreeMs / 1000 * srate
  local agree = 0
  for _, p in ipairs(list) do
    if p.sign == sign and math.abs(p.lag - med) <= tol then agree = agree + 1 end
  end
  return { lag = med, sign = sign, agree = agree, used = n, byHit = byHit, tol = tol,
           confident = agree >= math.max(CFG.minHitsAgree, CFG.minAgree * n) }
end

-- The sample with the biggest swing just after a hit's onset: the crest.
local function crestTime(acc, srate, onset)
  local n = math.floor(0.03 * srate)
  local s = readMono(acc, srate, onset, n)
  local best, at = -1, 1
  for i = 1, n do
    local v = math.abs(s[i]); if v > best then best, at = v, i end
  end
  return onset + (at - 1) / srate
end

---------------------------------------------------------------------------
-- Track helpers
---------------------------------------------------------------------------
local function trackName(tr)
  local _, name = reaper.GetTrackName(tr)
  return name
end

local function trackHasItems(tr) return reaper.CountTrackMediaItems(tr) > 0 end

local function itemsExtent(tr)
  local s, e = math.huge, -math.huge
  for i = 0, reaper.CountTrackMediaItems(tr) - 1 do
    local it = reaper.GetTrackMediaItem(tr, i)
    local p = reaper.GetMediaItemInfo_Value(it, "D_POSITION")
    local l = reaper.GetMediaItemInfo_Value(it, "D_LENGTH")
    if p < s then s = p end
    if p + l > e then e = p + l end
  end
  return s, e
end

local function trackPolarity(tr)
  return reaper.GetMediaTrackInfo_Value(tr, "B_PHASE") == 1 and -1 or 1
end

-- Move every item on the track earlier by `offset` seconds (later if negative).
-- An item pushed before project start is trimmed at its start instead.
local function moveTrackItems(tr, offset)
  for i = 0, reaper.CountTrackMediaItems(tr) - 1 do
    local it = reaper.GetTrackMediaItem(tr, i)
    local pos = reaper.GetMediaItemInfo_Value(it, "D_POSITION") - offset
    if pos < 0 then
      local cut = -pos
      for t = 0, reaper.CountTakes(it) - 1 do
        local tk = reaper.GetTake(it, t)
        if tk then
          local rate = reaper.GetMediaItemTakeInfo_Value(tk, "D_PLAYRATE")
          local so = reaper.GetMediaItemTakeInfo_Value(tk, "D_STARTOFFS")
          reaper.SetMediaItemTakeInfo_Value(tk, "D_STARTOFFS", so + cut * rate)
        end
      end
      local len = reaper.GetMediaItemInfo_Value(it, "D_LENGTH")
      reaper.SetMediaItemInfo_Value(it, "D_LENGTH", math.max(0.001, len - cut))
      pos = 0
    end
    reaper.SetMediaItemInfo_Value(it, "D_POSITION", pos)
  end
end

-- The script's own marker. Its ID is kept in the project, so a rerun only ever
-- replaces the marker this script made, never one of the user's.
local MARKER_NAME = "Drum Align"
local EXT_SECTION, EXT_KEY = "SingtallDrumAlign", "markerID"

local function placeMarker(t)
  local _, old = reaper.GetProjExtState(0, EXT_SECTION, EXT_KEY)
  old = tonumber(old)
  if old then
    local i = 0
    while true do
      local ok, isrgn, _, _, name, id = reaper.EnumProjectMarkers(i)
      if ok == 0 then break end
      if not isrgn and id == old and name == MARKER_NAME then
        reaper.DeleteProjectMarker(0, id, false)
        break
      end
      i = i + 1
    end
  end
  local id = reaper.AddProjectMarker2(0, false, t, 0, MARKER_NAME, -1, 0)
  reaper.SetProjExtState(0, EXT_SECTION, EXT_KEY, tostring(id))
end

local function flipTrackItems(tr)
  for i = 0, reaper.CountTrackMediaItems(tr) - 1 do
    local it = reaper.GetTrackMediaItem(tr, i)
    for t = 0, reaper.CountTakes(it) - 1 do
      local tk = reaper.GetTake(it, t)
      if tk then
        local v = reaper.GetMediaItemTakeInfo_Value(tk, "D_VOL")
        reaper.SetMediaItemTakeInfo_Value(tk, "D_VOL", -v)
      end
    end
  end
end

---------------------------------------------------------------------------
-- UI helpers
---------------------------------------------------------------------------
local function msg(text, kind)
  log("message box (" .. (kind == 1 and "OK/Cancel" or "OK") .. "): " .. text:gsub("\n", " | "):sub(1, 300))
  if TEST then TEST.log = (TEST.log or "") .. text .. "\n"; return 1 end
  local t0 = reaper.time_precise()
  local r = reaper.ShowMessageBox(text, SCRIPT_NAME, kind or 0)
  log(string.format("  box returned %s after %.1f s (1 = OK, 2 = Cancel)", tostring(r), reaper.time_precise() - t0))
  return r
end

local function chooseMaster(tracks)
  if TEST and TEST.master then
    for _, tr in ipairs(tracks) do if trackName(tr) == TEST.master then return tr end end
    return nil
  end
  -- The selected track named like a snare (top preferred over bottom), else the
  -- topmost selected track. The report names it, so Cancel backs out of a wrong pick.
  local pick
  for _, tr in ipairs(tracks) do
    local n = " " .. trackName(tr):lower() .. " "
    if n:find("snare") or n:find("snr") or n:find("%Wsn%W") then
      if not (n:find("bot") or n:find("under") or n:find("btm")) then return tr end
      pick = pick or tr
    end
  end
  return pick or tracks[1]
end

local function fmtMs(s) return string.format("%.2f ms", math.abs(s) * 1000) end

---------------------------------------------------------------------------
-- Main
---------------------------------------------------------------------------
local function main()
  local tracks = {}
  for i = 0, reaper.CountSelectedTracks(0) - 1 do tracks[#tracks + 1] = reaper.GetSelectedTrack(0, i) end
  if TEST and TEST.tracks then
    tracks = {}
    for i = 0, reaper.CountTracks(0) - 1 do
      local tr = reaper.GetTrack(0, i)
      if TEST.tracks == true or TEST.tracks[trackName(tr)] then tracks[#tracks + 1] = tr end
    end
  end
  if #tracks < 2 then
    msg("Select all the drum tracks first (Ctrl+click their names in the track panel, snare included), then run "
        .. SCRIPT_NAME .. " again.\n\n" .. reaper.CountSelectedTracks(0) .. " track(s) selected right now.")
    return
  end

  local names = {}
  for _, tr in ipairs(tracks) do names[#names + 1] = trackName(tr) end
  log(#tracks .. " tracks selected: " .. table.concat(names, ", "))
  local masterTr = chooseMaster(tracks)
  if not masterTr then log("no master track") return end
  log("master: " .. trackName(masterTr))
  if not trackHasItems(masterTr) then msg("The master track has no items.") return end

  local projSr = reaper.GetSetProjectInfo(0, "PROJECT_SRATE", 0, false)
  local srate = (projSr >= 44100 and projSr <= 96000) and projSr or 48000

  -- Always the whole master track: a short or forgotten time selection gives
  -- far fewer usable hits (and cut-up tracks may have audio in none of them).
  local rs, re = itemsExtent(masterTr)

  local macc = reaper.CreateTrackAudioAccessor(masterTr)
  local hits = findHits(macc, srate, rs, re)
  if #hits < 4 then
    reaper.DestroyAudioAccessor(macc)
    msg("Only found " .. #hits .. " clear hits on '" .. trackName(masterTr) .. "'.\nPick the snare as the master.")
    return
  end

  local master = { hits = {} }
  for _, h in ipairs(hits) do
    local on, s = onsetNear(macc, srate, h.time, 0.005, 0.035)
    if on then master.hits[#master.hits + 1] = { on = on, sign = s, time = h.time, peak = h.peak } end
  end

  local results = {}
  for _, tr in ipairs(tracks) do
    if tr ~= masterTr then
      local r = { track = tr, name = trackName(tr) }
      if not trackHasItems(tr) then r.skip = "no items"
      else
        r.a = analyseTrack(tr, master, srate)
        if r.a.silent then r.skip = "silent where the snare hits" end
      end
      results[#results + 1] = r
    end
  end

  -- Pick the one hit to align on: the hit where the most tracks give the same
  -- answer as they do across the rest of the song, i.e. the cleanest hit, with
  -- no hi-hat or kick landing on top of it. Louder wins a tie.
  local pick, pickScore = 1, -1
  for h, hit in ipairs(master.hits) do
    local score = 0
    for _, r in ipairs(results) do
      local a = r.a
      if a and not a.silent and a.confident then
        local p = a.byHit[h]
        if p and p.sign == a.sign and math.abs(p.lag - a.lag) <= a.tol then score = score + 1 end
      end
    end
    if score > pickScore or (score == pickScore and hit.peak > master.hits[pick].peak) then
      pick, pickScore = h, score
    end
  end
  local hit = master.hits[pick]
  log(string.format("%d hits found, using hit %d at %.3f s", #master.hits, pick, hit.time))
  local crest = crestTime(macc, srate, hit.on)
  reaper.DestroyAudioAccessor(macc)

  -- Every track is aligned on that hit. A track with no audio there (cut into
  -- pieces) uses its result from the rest of the song, if that is clear.
  local masterPol = trackPolarity(masterTr)
  for _, r in ipairs(results) do
    local a = r.a
    if a and not a.silent then
      local p = a.byHit[pick]
      if p then
        r.lag, r.sign = p.lag, p.sign
        if not (a.confident and p.sign == a.sign and math.abs(p.lag - a.lag) <= a.tol) then
          r.note = "check by ear, other hits disagree"
        end
      else
        -- No clear snare start in this mic at the marker hit (cut into pieces,
        -- or buried in its own ringing): use its typical value across the song.
        r.lag, r.sign = a.lag, a.sign
        r.note = a.confident and "no clear snare at the marker, used the rest of the song"
                              or "no clear snare at the marker, used the rest of the song - check by ear"
      end
      if r.lag then
        r.offset = r.lag / srate
        r.flip = (r.sign * masterPol * trackPolarity(r.track)) < 0
      end
    end
  end

  -- If most tracks are opposite polarity to the master, the master is the odd
  -- one out: flip it once instead of flipping all of them.
  local count, flips = 0, 0
  for _, r in ipairs(results) do
    if r.lag then
      count = count + 1
      if r.flip then flips = flips + 1 end
    end
  end
  local flipMaster = count >= 2 and flips * 2 > count
  for _, r in ipairs(results) do
    if r.lag then
      if flipMaster then r.flip = not r.flip end
      if math.abs(r.offset) * 1000 < CFG.minMoveMs and not r.flip then r.skip = "already aligned" end
    end
  end

  -- Report
  local crestStr = reaper.format_timestr_pos(crest, "", 0)
  local lines = {
    string.format("Master: %s, aligned on its hit at %s", trackName(masterTr), crestStr) }
  local toApply, trimmed = 0, false
  if flipMaster then
    lines[#lines + 1] = string.format("  flip polarity - it is opposite to %d of the %d tracks below", flips, count)
    toApply = toApply + 1
  end
  lines[#lines + 1] = ""
  for _, r in ipairs(results) do
    local line
    if r.skip == "already aligned" then
      line = string.format("%s:  already aligned", r.name)
    elseif r.skip then
      line = string.format("%s:  skipped (%s)", r.name, r.skip)
    else
      local dir = r.offset >= 0 and "earlier" or "later"
      line = string.format("%s:  move %s %s", r.name, fmtMs(r.offset), dir)
      if r.flip then line = line .. ", flip polarity" end
      if r.offset > 0 and itemsExtent(r.track) < r.offset then trimmed = true end
      if r.note then line = line .. "   <- " .. r.note end
      toApply = toApply + 1
    end
    lines[#lines + 1] = line
  end
  lines[#lines + 1] = ""
  if TEST then TEST.results = results; TEST.flipMaster = flipMaster; TEST.crest = crest end
  if toApply == 0 then
    lines[#lines + 1] = "Nothing to move."
    lines[#lines + 1] = "\nFree from Singtall Studio - " .. WEBSITE
    msg(table.concat(lines, "\n"))
    return
  end
  if trimmed then
    lines[#lines + 1] = "Items that start at 0:00 can't move earlier, so their starts are trimmed instead."
  end
  lines[#lines + 1] = "Press OK to make these changes and mark the hit. One undo puts everything back."
  lines[#lines + 1] = "\nFree from Singtall Studio - " .. WEBSITE

  local ok = msg(table.concat(lines, "\n"), 1)
  if ok ~= 1 or (TEST and not TEST.apply) then return end

  reaper.Undo_BeginBlock()
  reaper.PreventUIRefresh(1)
  if flipMaster then flipTrackItems(masterTr) end
  for _, r in ipairs(results) do
    if not r.skip then
      if math.abs(r.offset) * 1000 >= CFG.minMoveMs then moveTrackItems(r.track, r.offset) end
      if r.flip then flipTrackItems(r.track) end
    end
  end
  placeMarker(crest)
  reaper.PreventUIRefresh(-1)
  reaper.UpdateArrange()
  reaper.Undo_EndBlock(SCRIPT_NAME .. ": align to " .. trackName(masterTr), -1)
  log("changes applied")

  local changed = flipMaster and 1 or 0
  for _, r in ipairs(results) do if not r.skip then changed = changed + 1 end end
  msg(string.format("Done: changed %d track%s.\n\nZoom in on the \"%s\" marker (%s): every track now lines up on that hit.\nCtrl+Z puts everything back.",
      changed, changed == 1 and "" or "s", MARKER_NAME, crestStr))
end

log(SCRIPT_NAME .. " " .. VERSION .. " started")
local ok, err = xpcall(main, debug.traceback)
if not ok then
  log("ERROR: " .. tostring(err))
  if TEST then error(err) end
  reaper.ShowMessageBox("Something went wrong:\n\n" .. tostring(err) .. "\n\nDetails were saved to:\n" .. LOG_PATH,
                        SCRIPT_NAME, 0)
end
log("finished")
if logFile then logFile:close() end