/*
 SynKlick  -  repair synchronous skips in raw sound recording

 (C) Ichthyostega (Hermann Vosseler), 2014

All code in this package is provided "as is" and WITHOUT ANY WARRANTY.
It may be copied and used under the terms and conditions of the
GNU PUBLIC LICENCE (GPL) Version 2 and above.

*********************************************************************/

#include <sndfile.h>
#include <stdio.h>

#include <string>
#include <vector>
#include <iostream>
#include <cassert>
#include <cstdlib>
#include <cmath>

using std::string;
using std::to_string;
using std::vector;
using std::cout;
using std::cerr;
using std::max;
using std::abs;

const size_t CHUNK_SIZE  = 100000;
const size_t CONTEXT    = 100;
const size_t BUFF_SIZE = CHUNK_SIZE + 2*CONTEXT;

const float EPSILON = 1.0f / (1 << 23);      // one normalised sample step for 24bit samples 
const float SILENCE_GATE = 4096.0f * EPSILON;// signal below this level (-33dB) is too low to detect jumps
const float TRIGGER_FACTOR = 6.0f;           // discontinuity detected by local change compared to average change
const float CONFIRMATION_DEGREE = 0.66f;     // degree of confirmation by multiple tracks for an event to be accepted
const uint  AVERAGE_SAMPLES = 50;            // number of past samples to use for average delta
const uint  POS_LEEWAY = 3;                  // events within this timespan (samples) still count as "synchronous"

const bool VERBOSE = false; // set to true to trace the decision process

/* ==== Helpers ==== */

void
error (string msg)
{
  cerr << "processing FAILED: "
       << msg << "\n\n";
  exit(EXIT_FAILURE);
}

string
name_without_path (string rawName)
{
  auto pos = rawName.rfind('/');
  return (pos == string::npos)? rawName
                              : rawName.substr(pos+1);
}

void
reduceBufferPos (size_t & pos, size_t elementsToRemove)
{
  if (pos < elementsToRemove)
    pos = 0;
  else
    pos -= elementsToRemove;
}




/* ==== Abstraction: single channel to be processed ==== */

/**
 * Helper to decide if a glitch event should be considered,
 * while visiting multiple tracks at the same location.
 */
class EventScore
  {
    float weights_ = 0.0f;
    float score_ = 0.0f;
    uint confs_ = 0;
    bool isHere_ = false;
    static uint minConfs;
    
    uint last_ = 0;
    float fac_ = 0.0f;
    
  public:
    float
    confirmationDegree()
      {
        if (0==weights_) return 0.0;
        float avgScore = score_/weights_;
        return 0.5 * (avgScore + 1.0);
      }
    
    operator bool()
      {
        return CONFIRMATION_DEGREE <= this->confirmationDegree()
            && confs_ >= minConfs
            && isHere_;
      }
    
    void
    record (bool unclear, bool confirmed, bool noOffset, float factor)
      {
        fac_  = factor;
        last_ = unclear? 0 : confirmed? 2 : 1; // record for diagnostics
        if (unclear) return;
        if (confirmed) ++confs_;
        if (noOffset) isHere_ = true;
        
        float weight = confirmed? (factor/TRIGGER_FACTOR)
                                : (0.0<factor? TRIGGER_FACTOR/factor : 10);
        weights_ += weight;
        score_   += weight * (confirmed? 1.0 : -1.0);
      }
    
    void
    recordNeutral()
      {
        record (true, false, true, 0.0);
      }
    
    bool
    possible()
      {
        return 0.0 < weights_ && 0.1 < confirmationDegree();
      }
    
    string
    showLastUpdate()
      {
        return 0==last_? "neutral"
                       : (2==last_? "conf\t"
                                  : "rej \t"
                         ) + to_string(fac_) +'\t'+ to_string(confirmationDegree());
      }
    
    /**
     * to avoid triggering by noise peaks, we require a confirmation
     * by multiple channels; but the minimum fraction of confirming
     * channels is lowered when processing a large number of channels.
     */
    static void
    adjustRequiredMinimalConfirmations (uint channelCount)
      {
        minConfs = 1 + uint(std::log2(channelCount));
      }
  };




/**
 * single sound channel with the ability to investigate itself.
 * Internally holds a work buffer to be filled through libsndfile.
 * The #visit function performs the decision logic to detect a defect
 * at the "current" position, including the next neighbouring samples.
 */
class Channel
  {
    SNDFILE *in_ = nullptr,
            *out_ = nullptr;
    SF_INFO info_in_,
            info_out_;
    
    float buff_[BUFF_SIZE];
    size_t numSamples_ = 0;
    size_t workPos_    = 0;
    float ema_ = 1.0f / TRIGGER_FACTOR;
    
    
  public:
    string name_;
    static bool hitEnd;
    
    void
    open (string name, string target)
      {
        this->name_ = name_without_path(name);
        target += name_;
        
        info_in_.format = 0;
        in_ = sf_open (name.c_str(), SFM_READ, &info_in_);
        if (!in_)
          error ("could not open "+name+" -- " + sf_strerror (NULL));
        if (1 != info_in_.channels)
          error ("only Mono streams supported -- but '"+name_
                +"' was found to have "+to_string(info_in_.channels)+" channels.");
        
        info_out_.format = info_in_.format;
        info_out_.channels = info_in_.channels;
        info_out_.samplerate = info_in_.samplerate;
        out_ = sf_open(target.c_str(), SFM_WRITE, &info_out_);
        if (!out_)
          error ("unable to create target "+target+" -- " + sf_strerror (NULL));
      }
    
    
   ~Channel()
      {
        if (in_) sf_close (in_);
        if (out_)
          if (0 != sf_close (out_))
            cerr << "WARNING: problems while closing target '" << name_
                 << "' -- " << sf_strerror(out_) << "\n";
        
        in_ = out_ = nullptr;
      }
    
    
    void
    fillBuffer()
      {
        assert(numSamples_ < BUFF_SIZE);
        float* startPos = buff_ + numSamples_;
        size_t samplesToRead = BUFF_SIZE - numSamples_;
        numSamples_ = sf_read_float(in_, startPos, samplesToRead);
        if (numSamples_ < samplesToRead)
          hitEnd = true; // causes the processing to end with this chunk.
        
        numSamples_ += (startPos - buff_);
        assert(workPos_ <= numSamples_);
      }
    
    void
    writeBuffer()
      {
        size_t samplesToWrite = numSamples_ - (hitEnd? 0 : 2*CONTEXT);
        size_t actuallyWritten = sf_write_float(out_, buff_, samplesToWrite);
        if (actuallyWritten < samplesToWrite)
          cerr << "Alarm: only "<<actuallyWritten<<" instead of "<<samplesToWrite
               <<" samples written for " << name_ <<"\n";
        
        if (!hitEnd)
          // move processing context samples to the beginning of the buffer;
          // they will be written together with the next chunk of data
          {
            for (uint i=0; i < 2*CONTEXT; ++i)
              buff_[i] = buff_[i + actuallyWritten];
            reduceBufferPos(numSamples_, actuallyWritten);
            reduceBufferPos(workPos_,    actuallyWritten);
          }
        assert(workPos_ <= numSamples_);
      }
    
    /**
     * investigate a single sample position in the current buffer content.
     * The delta to the preceding sample(s) is used to detect a signal glitch.
     * @param i position of current sample
     * @param event state of other tracks investigated thus far
     * @return the event detection state including the results of investigating this sample.
     *         If the sample is positively an glitch event, and the state is undecided thus far,
     *         the event state is confirmed; if a glitch can not be confirmed, the state is
     *         cleared. Otherwise, the undecided state is passed on.
     */
    void
    visit (uint i, EventScore& event)
      {
        if (i + CONTEXT < numSamples_
            && CONTEXT  <= i)
          {
            bool noOffset = fineTuneWorkPos (i);
            float averageDelta = updateMovingAverage();
            float changeFactor = delta(workPos_) / (averageDelta + EPSILON);
            bool unclear = (sampleLevel() < SILENCE_GATE );
            bool trigger = (changeFactor > TRIGGER_FACTOR);
            event.record (unclear, trigger, noOffset, changeFactor);
          }
        else
          event.recordNeutral();
      }
    
    /** if it ain't broken, don't fix it */
    bool
    reallyNeedsFix()
      {
        float changeFactor = delta(workPos_) / (ema_ + EPSILON);
        return changeFactor > 0.25 * TRIGGER_FACTOR;
      }
    
    
  private:
    float
    sampleLevel()
      {
        auto curr = abs(buff_[workPos_]);
        auto prev = abs(buff_[workPos_ - 1]);
        
        return 0 < workPos_? max(prev, curr)
                           : curr;
      }
    
    float
    delta(size_t pos)
      {
        if (0 < pos)
          return abs(buff_[pos] - buff_[pos-1]);
        else
          return 0;
      }
    
    float
    updateMovingAverage()
      {
        float old_ema = ema_;  // avoid poisoning EMA used for detection
        
        ema_ = (delta(workPos_) + ema_*(AVERAGE_SAMPLES - 1)) / AVERAGE_SAMPLES;
        return old_ema;
      }
    
    /**
     * sometimes the actual jump appears slightly offset on neighbouring tracks.
     * We allow for some leeway by moving the work position to maximise delta.
     * @return if the triggering event is at the current position
     */
    bool
    fineTuneWorkPos (size_t currPos)
      {
        size_t maxPos = 0;
        float largestDelta = 0.0f;
        for (size_t j=currPos; j < currPos+POS_LEEWAY && j < numSamples_; ++j)
          {
            float thisDelta = delta(j);
            if (thisDelta > largestDelta)
              {
                largestDelta = thisDelta;
                maxPos = j;
              }
          }
        
        workPos_ = maxPos;
        return (workPos_ == currPos);
      }
    
    float*
    accessCurrentContext()
      {
        assert(CONTEXT <= workPos_);
        assert(CONTEXT + workPos_ <= numSamples_);
        return buff_ + (workPos_ - CONTEXT);
      }
    friend class Patch;
    
  };



/**
 * Fix a glitch detected in the corresponding channel.
 * When the Patch is applied, internal work buffers are used
 * to create two (partially synthetic) copies of the channel's signal
 * which are then cross faded. The result is pushed back into the
 * output of the underlying channel.
 */
class Patch
  {
    static const size_t LEN = 2 * CONTEXT;
    
    float workBuff[LEN];
    float* rawBuff;
    
    
    void
    reconstructMirroredSignal()
      {
        for (size_t i=0; i < LEN; ++i)
          workBuff[i] = rawBuff[LEN-1 - i];
      }
    
    void
    maybe_fix_Gap()
      {
        const size_t CENTRE = CONTEXT;
        size_t i = CENTRE-1;
        while (i > 0 && 0.0f == workBuff[i]) --i;
        if (i < CENTRE-1)
          fix_Gap_by_Interpolation(i, CENTRE);
        
        i = CENTRE;
        while (i < LEN-1 && 0.0f == workBuff[i]) ++i;
        if (i > CENTRE)
          fix_Gap_by_Interpolation(CENTRE-1, i);
      }
    
    void
    fix_Gap_by_Interpolation(size_t s, size_t e)
      {
        assert (s+1 < e);
        assert (0 <= s && s+1 < LEN);
        assert (0 <= e && e < LEN);
        
        float before = workBuff[s];
        float after  = workBuff[e];
        for (size_t i = s+1; i<e; ++i)
          {
            float fade = float(i-s) / float(e-s);
            workBuff[i] = rawBuff[LEN-1 - i]
                        = (1.0-fade)*before + fade*after;
          }
      }
    
    void
    mix_faded_with_mirroredSignal()
      {
        for (uint i=0; i < LEN; ++i)
          {
            float fade = (i<CONTEXT) ? CONTEXT - (i+1)
                                     : i - CONTEXT;
            fade = 1.0 - fade/CONTEXT;
            // i=-1        -> 0
            // i=CONTEXT-1 -> 1
            // i=CONTEXT   -> 1
            // i=2*CONTEXT -> 0
            
            float p = fade * 0.5f,
                  q = 1.0 - p;
            
            rawBuff[i] = q*rawBuff[i] + p*workBuff[i];
          }
      }
    
    
  public:
    Patch (Channel& c)
      : rawBuff{c.accessCurrentContext()}
      {
        reconstructMirroredSignal();
      }
    
    void
    apply()
      {
        maybe_fix_Gap();
        mix_faded_with_mirroredSignal();
      }
  };





/* ==== high-level operation ==== */

int num_sources = 0;
vector<Channel> streams(0);
bool Channel::hitEnd = false;
uint EventScore::minConfs = 1;


void
openStreams (char* name[], string targetDir)
{
  assert (num_sources);
  assert (0 == streams.size());
  streams.resize(num_sources);
  EventScore::adjustRequiredMinimalConfirmations (num_sources);
  
  assert (targetDir.length());
  if (targetDir.back() != '/')
    targetDir += '/';
  
  for (uint i=0; i < num_sources; ++i)
    streams[i].open (name[i], targetDir);
}

void
closeStreams()
{
  streams.clear();
}


void
eventConfirmationDiagnostic(EventScore& event, string name, size_t absPos)
{
    static size_t trail = 0;
    if (trail && trail != absPos) trail = 0;
    if (!trail && event.possible())
      {
        trail = absPos;
        cerr << "...\n";
      }
    
    if (trail)
      cerr <<absPos<<"\t jump? "
           <<name<<"\t :: "
           <<event.showLastUpdate()
           <<"\n";
}

void
process()
{
  size_t absStartPos = 0; // for diagnostic messages
  
  while (!Channel::hitEnd)
    {
      for (Channel& channel : streams)
        channel.fillBuffer();
      
      for (uint i=0; i < BUFF_SIZE; ++i)
        {
          EventScore event;
          for (Channel& channel : streams)
            {
              channel.visit(i, event);
              if (VERBOSE)
                eventConfirmationDiagnostic (event, channel.name_, absStartPos+i);
            }
          if (event)
            {
              cout << "PATCH at: "<< absStartPos+i <<"\n";
              for (Channel& channel : streams)
                if (channel.reallyNeedsFix())
                  Patch(channel).apply();
            }
        }
      
      for (Channel& channel : streams)
        channel.writeBuffer();
      
      absStartPos += CHUNK_SIZE;
    }
}





int
main (int argc, char* argv[]) {
    if (argc < 3 || string("-h") == argv[1])
      {
        cerr << "SynKlick <INPUTFILE> [<INPUTFILE2> ...] <TARGETDIR>\n"
             << "\nRepair a very special kind of defects in sound recordings:\n"
             << "synchronous skips in a set of input sound files. Possible events\n"
             << "will be detected by scanning the deltas between neighbouring samples\n"
             << "and confirmed by jumps at the same position in all files.\n"
             << "The jump will be amended by a short cross-blend of about 4.5ms.\n"
             << "\nA fixed copy of each file will be written to <TARGETDIR>\n\n";
        
        return EXIT_FAILURE;
      }
    num_sources = argc - 2;
    string targetDir = argv[argc-1];
    
    openStreams (&argv[1], targetDir);
    process();
    closeStreams();
    cout << "\nSUCCESS\n"
         << num_sources << " files read and processed --> "
         << targetDir << "\n\n";
    return EXIT_SUCCESS;
}