-- @description Singtall Drum Clean
-- @version 1.0
-- @author Singtall Studio
-- @website
--   Singtall Drum Clean https://singtallstudio.com/drum-clean.html
--   Singtall Studio https://singtallstudio.com
-- @changelog
--   First release.
-- @about
--   Removes the bleed between hits on kick and tom tracks by cutting the items,
--   the way it is done by hand, then glues each track back into one item.
--
--   1. Select all the drum tracks (kick, snare, toms, overheads, rooms...).
--      Only kick and tom tracks are cleaned. Every other selected mic is used to
--      recognise bleed: when the snare mic is far louder than a tom mic, what the
--      tom mic hears is snare, not tom.
--   2. Run the script, check the report and press OK. One undo puts everything back.
--
--   Hits it isn't sure about are kept. Before cutting, plain copies of the tracks
--   as they were (no FX) go into a muted "Drum Clean originals" folder at the
--   bottom of the project, on the same timeline, in case anything is missing.
--
--   Use it at the start of a mix, before any plugins: run Singtall Drum Align,
--   then Drum Clean. Aligned, clean drums need far fewer plugins (no gates).

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

local CFG = {
  envRate     = 16000,  -- envelopes are read at this sample rate (level only)
  blockMs     = 2,      -- envelope resolution
  riseDb      = 10,     -- a transient is a jump of this much over the last ~16 ms
  aboveMedDb  = 10,     -- ...and this far above the track's median level
  minGapMs    = 80,     -- transients closer than this count as one
  leakHits    = 30,     -- leak from a mic is learned at its strongest hits
  realDb      = 10,     -- this far above the expected bleed = a real hit
  unsureDb    = 5,      -- between unsureDb and realDb = kept and marked for checking
  preMs       = 4,      -- each kept piece starts this long before the attack
  fadeInMs    = 2,
  fadeOutMs   = 30,
  minHoldMs   = 60,     -- a hit is kept at least this long
  maxHoldMs   = 3000,   -- ...and at most this long
  quietMs     = 20,     -- it has rung out once it stays near the background this long
  floorPct    = 0.6,    -- background level = this percentile of the track's level
  floorDb     = 3,      -- "near the background" = within this of it
  mergeMs     = 50,     -- kept pieces closer than this are joined
}

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

local LOG_PATH = (({ reaper.get_action_context() })[2] or ""):match("^(.*[\\/])") or (reaper.GetResourcePath() .. "/Scripts/")
LOG_PATH = LOG_PATH .. "Singtall Drum Clean 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

---------------------------------------------------------------------------
-- Track names: which drum a mic belongs to
---------------------------------------------------------------------------
local function trackName(tr)
  local _, name = reaper.GetTrackName(tr)
  return name
end

local function has(n, ...)
  for _, w in ipairs({ ... }) do if n:find(w) then return true end end
  return false
end

-- Mics of the same drum (kick in / kick out, tom 1 top / bottom) share a group,
-- so they are never treated as bleed sources for each other.
local function drumGroup(name)
  local n = " " .. name:lower():gsub("[%p_]", " ") .. " "
  if has(n, "kick", "kik", "%sbd%s", "bass drum") then return "kick", true end
  if has(n, "snare", "snr", "%ssn%s") then return "snare", false end
  local isTom = has(n:gsub("bottom", " "), "tom", "floor", "%sft%s", "rack", "%stm%s")
  for _, w in ipairs({ "top", "bottom", "bot", "btm", "in", "out", "sub", "inside", "outside",
                       "reso", "batter", "mic", "close" }) do
    n = n:gsub("%s" .. w .. "%s", " ")
  end
  return n:gsub("%s+", " "), isTom
end

---------------------------------------------------------------------------
-- Envelopes
---------------------------------------------------------------------------
local readBuf, readBufSize
local function readMono(acc, srate, t0, n)
  if not readBuf or readBufSize < n * 2 then
    readBufSize = n * 2
    readBuf = reaper.new_array(readBufSize)
  end
  readBuf.clear()
  reaper.GetAudioAccessorSamples(acc, srate, 2, t0, n, readBuf)
  return readBuf.table(1, n * 2)
end

-- Peak level per block (dB) of a track's items before its FX, over rs..re.
local function readEnv(track, rs, nBlocks)
  local srate = CFG.envRate
  local B = math.floor(CFG.blockMs / 1000 * srate + 0.5)
  local perChunk = 1000
  local acc = reaper.CreateTrackAudioAccessor(track)
  local env, b = {}, 0
  while b < nBlocks do
    local blocks = math.min(perChunk, nBlocks - b)
    local s = readMono(acc, srate, rs + b * B / srate, blocks * B)
    for k = 0, blocks - 1 do
      local m = 0
      for i = k * B * 2 + 1, (k + 1) * B * 2 do
        local v = s[i]; if v < 0 then v = -v end
        if v > m then m = v end
      end
      b = b + 1
      env[b] = m > 1e-9 and 20 * math.log(m, 10) or -180
    end
  end
  reaper.DestroyAudioAccessor(acc)
  return env
end

local function percentile(t, p)
  local s = {}
  for i = 1, #t do s[i] = t[i] end
  table.sort(s)
  return s[math.max(1, math.min(#s, math.floor(p * #s + 0.5)))]
end

-- Transients: a jump of riseDb over the previous ~16 ms, well above the median.
local function findEvents(env)
  local med = percentile(env, 0.5)
  local blockS = CFG.blockMs / 1000
  local gap = math.floor(CFG.minGapMs / CFG.blockMs + 0.5)
  local ev, last = {}, -1e9
  for i = 11, #env - 15 do
    if i - last >= gap and env[i] > med + CFG.aboveMedDb then
      local prev = 0
      for j = i - 10, i - 3 do prev = prev + env[j] end
      prev = prev / 8
      if env[i] - prev > CFG.riseDb then
        local pk = i
        for j = i, i + 14 do if env[j] > env[pk] then pk = j end end
        ev[#ev + 1] = { on = i, peak = pk }
        last = i
      end
    end
  end
  return ev
end

local function levelAt(env, i)
  local m = -180
  for j = math.max(1, i - 2), math.min(#env, i + 2) do if env[j] > m then m = env[j] end end
  return m
end

---------------------------------------------------------------------------
-- Items
---------------------------------------------------------------------------
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

-- Cut the track's items down to the kept pieces, with fades, then glue them.
local function cleanTrack(tr, pieces)
  local items = {}
  for i = 0, reaper.CountTrackMediaItems(tr) - 1 do items[#items + 1] = reaper.GetTrackMediaItem(tr, i) end
  local kept = {}
  for _, item in ipairs(items) do
    local ip = reaper.GetMediaItemInfo_Value(item, "D_POSITION")
    local ie = ip + reaper.GetMediaItemInfo_Value(item, "D_LENGTH")
    local cuts = {}
    for _, p in ipairs(pieces) do
      if p.s > ip and p.s < ie then cuts[#cuts + 1] = p.s end
      if p.e > ip and p.e < ie then cuts[#cuts + 1] = p.e end
    end
    table.sort(cuts)
    local parts, cur = { item }, item
    for _, t in ipairs(cuts) do
      local right = reaper.SplitMediaItem(cur, t)
      if right then parts[#parts + 1] = right; cur = right end
    end
    for _, part in ipairs(parts) do
      local ps = reaper.GetMediaItemInfo_Value(part, "D_POSITION")
      local pe = ps + reaper.GetMediaItemInfo_Value(part, "D_LENGTH")
      local mid, keep = 0.5 * (ps + pe), false
      for _, p in ipairs(pieces) do if mid > p.s and mid < p.e then keep = true; break end end
      if keep then
        reaper.SetMediaItemInfo_Value(part, "D_FADEINLEN", CFG.fadeInMs / 1000)
        reaper.SetMediaItemInfo_Value(part, "D_FADEOUTLEN", math.min(CFG.fadeOutMs / 1000, 0.5 * (pe - ps)))
        kept[#kept + 1] = part
      else
        reaper.DeleteTrackMediaItem(tr, part)
      end
    end
  end
  if #kept > 0 then
    reaper.Main_OnCommand(40289, 0)                -- Item: Unselect all items
    for _, it in ipairs(kept) do reaper.SetMediaItemSelected(it, true) end
    reaper.Main_OnCommand(40362, 0)                -- Item: Glue items
    reaper.Main_OnCommand(40289, 0)
  end
  return #kept
end

---------------------------------------------------------------------------
-- Version 1.0 marked unsure hits; those markers are removed (only the ones this
-- script made, by the IDs it kept in the project).
---------------------------------------------------------------------------
local MARKER_NAME = "Drum Clean: check"
local EXT_SECTION, EXT_KEY = "SingtallDrumClean", "markerIDs"

local function removeOldMarkers()
  local _, old = reaper.GetProjExtState(0, EXT_SECTION, EXT_KEY)
  local mine = {}
  for id in (old or ""):gmatch("%d+") do mine[tonumber(id)] = true end
  local gone, i = {}, 0
  while true do
    local ok, isrgn, _, _, name, id = reaper.EnumProjectMarkers(i)
    if ok == 0 then break end
    if not isrgn and mine[id] and name:sub(1, #MARKER_NAME) == MARKER_NAME then gone[#gone + 1] = id end
    i = i + 1
  end
  for _, id in ipairs(gone) do reaper.DeleteProjectMarker(0, id, false) end
  reaper.SetProjExtState(0, EXT_SECTION, EXT_KEY, "")
  return #gone
end

---------------------------------------------------------------------------
-- Backup: copies of the tracks as they were, muted, in a folder at the bottom
---------------------------------------------------------------------------
local FOLDER_NAME = "Drum Clean originals"

local function backupTracks(list)
  -- Remember the user's track selection, duplicate only the tracks being cleaned.
  local had = {}
  for i = 0, reaper.CountSelectedTracks(0) - 1 do had[#had + 1] = reaper.GetSelectedTrack(0, i) end
  reaper.Main_OnCommand(40297, 0)                      -- Track: Unselect all tracks
  for _, tr in ipairs(list) do reaper.SetTrackSelected(tr, true) end
  reaper.Main_OnCommand(40062, 0)                      -- Track: Duplicate tracks (copies end up selected)
  local copies = {}
  for i = 0, reaper.CountSelectedTracks(0) - 1 do copies[#copies + 1] = reaper.GetSelectedTrack(0, i) end

  -- A new folder track at the very end, named so a second run doesn't mix with the first.
  local name, n = FOLDER_NAME, 1
  local taken = {}
  for i = 0, reaper.CountTracks(0) - 1 do taken[trackName(reaper.GetTrack(0, i))] = true end
  while taken[name] do n = n + 1; name = FOLDER_NAME .. " " .. n end
  local last = reaper.CountTracks(0)
  -- If the old last track closes a folder, the new track must not end up inside it.
  local prev = last > 0 and reaper.GetTrack(0, last - 1)
  reaper.InsertTrackAtIndex(last, false)
  local folder = reaper.GetTrack(0, last)
  reaper.GetSetMediaTrackInfo_String(folder, "P_NAME", name, true)
  if prev and reaper.GetMediaTrackInfo_Value(prev, "I_FOLDERDEPTH") > 0 then
    reaper.SetMediaTrackInfo_Value(prev, "I_FOLDERDEPTH", 0)
  end

  -- Move the copies into it, then mute everything and fold it up.
  reaper.Main_OnCommand(40297, 0)
  for _, c in ipairs(copies) do reaper.SetTrackSelected(c, true) end
  reaper.ReorderSelectedTracks(reaper.CountTracks(0), 1)   -- children of the track before (the folder)
  for _, c in ipairs(copies) do
    reaper.GetSetMediaTrackInfo_String(c, "P_NAME", trackName(c) .. " (before Drum Clean)", true)
    reaper.SetMediaTrackInfo_Value(c, "B_MUTE", 1)
    -- Plain audio backups: Drum Clean belongs before any processing, and muted
    -- copies shouldn't cost CPU.
    for fx = reaper.TrackFX_GetCount(c) - 1, 0, -1 do reaper.TrackFX_Delete(c, fx) end
  end
  reaper.SetMediaTrackInfo_Value(folder, "B_MUTE", 1)
  reaper.SetMediaTrackInfo_Value(folder, "I_FOLDERCOMPACT", 2)

  reaper.Main_OnCommand(40297, 0)
  for _, tr in ipairs(had) do reaper.SetTrackSelected(tr, true) end
  return name, #copies
end

---------------------------------------------------------------------------
-- UI
---------------------------------------------------------------------------
local function msg(text, kind)
  log("message box (" .. (kind == 1 and "OK/Cancel" or "OK") .. "): " .. text:gsub("\n", " | "):sub(1, 400))
  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

---------------------------------------------------------------------------
-- Main
---------------------------------------------------------------------------
local function main()
  local tracks = {}
  if TEST and TEST.tracks then
    for i = 0, reaper.CountTracks(0) - 1 do
      local tr = reaper.GetTrack(0, i)
      if TEST.tracks[trackName(tr)] then tracks[#tracks + 1] = tr end
    end
  else
    for i = 0, reaper.CountSelectedTracks(0) - 1 do tracks[#tracks + 1] = reaper.GetSelectedTrack(0, i) end
  end

  local mics = {}
  for _, tr in ipairs(tracks) do
    if reaper.CountTrackMediaItems(tr) > 0 then
      local group, target = drumGroup(trackName(tr))
      mics[#mics + 1] = { track = tr, name = trackName(tr), group = group, target = target }
    end
  end
  local names = {}
  for _, m in ipairs(mics) do names[#names + 1] = m.name .. (m.target and " (clean)" or "") end
  log(#mics .. " tracks with items: " .. table.concat(names, ", "))

  local targets = 0
  for _, m in ipairs(mics) do if m.target then targets = targets + 1 end end
  if targets == 0 or #mics < 2 then
    msg("Select all the drum tracks first (Ctrl+click their names in the track panel): the kick and tom "
        .. "tracks to clean, plus the snare and the other mics so bleed can be recognised.\n\n"
        .. "Kick and tom tracks are found by name (kick, tom, floor...). Found " .. targets
        .. " to clean among " .. #mics .. " selected track(s) with items.")
    return
  end

  -- Envelopes of every selected mic over the same stretch of time.
  local rs, re = math.huge, -math.huge
  for _, m in ipairs(mics) do
    local s, e = itemsExtent(m.track)
    if s < rs then rs = s end
    if e > re then re = e end
  end
  local blockS = CFG.blockMs / 1000
  local nBlocks = math.floor((re - rs) / blockS)
  for _, m in ipairs(mics) do
    m.env = readEnv(m.track, rs, nBlocks)
    m.events = findEvents(m.env)
  end
  log(string.format("read %d tracks, %.0f s each", #mics, re - rs))

  -- How much of mic o leaks into mic k, learned at o's strongest hits.
  for _, o in ipairs(mics) do
    local top = {}
    for _, e in ipairs(o.events) do top[#top + 1] = e end
    table.sort(top, function(a, b) return o.env[a.peak] > o.env[b.peak] end)
    o.leak = {}
    for _, k in ipairs(mics) do
      if k ~= o then
        local d = {}
        for i = 1, math.min(CFG.leakHits, #top) do
          d[#d + 1] = levelAt(k.env, top[i].peak) - levelAt(o.env, top[i].peak)
        end
        o.leak[k] = #d > 0 and percentile(d, 0.5) or -180
      end
    end
  end

  -- Score every transient on the kick and tom mics.
  local results, checks = {}, {}
  local q = math.floor(CFG.quietMs / CFG.blockMs + 0.5)
  for _, k in ipairs(mics) do
    if k.target then
      local floor = percentile(k.env, CFG.floorPct) + CFG.floorDb
      local function expected(i)
        local x = -180
        for _, o in ipairs(mics) do
          if o ~= k and o.group ~= k.group then
            local v = levelAt(o.env, i) + o.leak[k]
            if v > x then x = v end
          end
        end
        return x
      end
      local r = { mic = k, real = 0, unsure = 0, bleed = 0, pieces = {} }
      for _, e in ipairs(k.events) do
        local over = levelAt(k.env, e.peak) - expected(e.peak)
        local kind = over >= CFG.realDb and "real" or (over >= CFG.unsureDb and "unsure" or "bleed")
        r[kind] = r[kind] + 1
        if kind ~= "bleed" then
          -- Keep from just before the attack until the ring has stayed near the
          -- background for quietMs (between minHold and maxHold).
          local i = e.peak + math.floor(CFG.minHoldMs / CFG.blockMs)
          local last = e.on + math.floor(CFG.maxHoldMs / CFG.blockMs)
          local run = 0
          while i < math.min(last, #k.env) do
            if k.env[i] < floor then run = run + 1 else run = 0 end
            if run >= q then break end
            i = i + 1
          end
          local s = rs + (e.on - 1) * blockS - CFG.preMs / 1000
          local t = rs + (i - run) * blockS + CFG.fadeOutMs / 1000
          r.pieces[#r.pieces + 1] = { s = s, e = t }
          if kind == "unsure" then checks[#checks + 1] = { t = rs + (e.on - 1) * blockS, name = k.name } end
        end
      end
      -- Join overlapping or nearly touching pieces (fills, flams).
      table.sort(r.pieces, function(a, b) return a.s < b.s end)
      local merged = {}
      for _, p in ipairs(r.pieces) do
        local m = merged[#merged]
        if m and p.s - m.e < CFG.mergeMs / 1000 then
          if p.e > m.e then m.e = p.e end
        else
          merged[#merged + 1] = { s = p.s, e = p.e }
        end
      end
      r.pieces = merged
      local keptS = 0
      for _, p in ipairs(merged) do keptS = keptS + (p.e - p.s) end
      r.keptPct = 100 * keptS / math.max(1e-9, re - rs)
      results[#results + 1] = r
      log(string.format("%s: %d real, %d unsure, %d bleed, %d pieces, %.0f%% kept",
          k.name, r.real, r.unsure, r.bleed, #merged, r.keptPct))
    end
  end

  -- Report
  local lines = { "Cleaning " .. #results .. " track(s); bleed is judged against "
                  .. (#mics - 1) .. " other mic(s).", "" }
  local any = false
  for _, r in ipairs(results) do
    lines[#lines + 1] = string.format("%s:  keep %d hit%s%s, remove %d bleed  (%.0f%% of the track stays)",
      r.mic.name, r.real + r.unsure, (r.real + r.unsure) == 1 and "" or "s",
      r.unsure > 0 and string.format(" (%d of them unsure)", r.unsure) or "", r.bleed, r.keptPct)
    if #r.pieces > 0 then any = true end
  end
  lines[#lines + 1] = ""
  if TEST then TEST.results = results; TEST.checks = checks end
  if not any then
    lines[#lines + 1] = "No hits found to keep, so nothing was changed."
    msg(table.concat(lines, "\n"))
    return
  end
  lines[#lines + 1] = "Press OK to cut away the bleed and glue each track back into one item."
  lines[#lines + 1] = "Unsure hits are kept. Copies of these tracks as they are now go into a muted"
  lines[#lines + 1] = "\"" .. FOLDER_NAME .. "\" folder at the bottom, in case anything is missed."
  lines[#lines + 1] = "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)
  local cleaned = {}
  for _, r in ipairs(results) do if #r.pieces > 0 then cleaned[#cleaned + 1] = r.mic.track end end
  local folder = backupTracks(cleaned)
  for _, r in ipairs(results) do
    if #r.pieces > 0 then cleanTrack(r.mic.track, r.pieces) end
  end
  local oldMarkers = removeOldMarkers()
  reaper.PreventUIRefresh(-1)
  reaper.TrackList_AdjustWindows(false)
  reaper.UpdateArrange()
  reaper.Undo_EndBlock(SCRIPT_NAME, -1)
  log(string.format("changes applied; backups in '%s'; %d old markers removed", folder, oldMarkers))
  msg(string.format("Done: cleaned %d track(s).\n\nThe tracks as they were are in the muted \"%s\" folder at the bottom: unmute a copy, or drag a piece back, if anything turns out to be missing.%s\nCtrl+Z puts everything back.",
      #cleaned, folder, oldMarkers > 0 and string.format("\n\nRemoved %d old \"%s\" marker(s).", oldMarkers, MARKER_NAME) or ""))
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