/*
  ==============================================================================

    PitchDetector.h
    Created: 24 Apr 2021 2:07:10pm
    Author:  brady

    Pitch detection algorithm using HPCP algorithm from:

    TONAL DESCRIPTION OF MUSIC AUDIO SIGNALS
    Gutierrez 2006

    and

    Melody Extraction Using Chroma-Level Note Tracking and Pitch Mapping
    Weiwei Zhang, Zhe Chen and Fuliang Yin 2018

    http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.76.8192&rep=rep1&type=pdf

    also Melody Detection in Polyphonic Audio
    Rui Pedro Pinto de Carvalho e Paiva 2006

    https://www.researchgate.net/publication/255908118

  ==============================================================================
*/

#pragma once

#include <juce_core/juce_core.h>

#include "DSP/Fft.h"
#include "Utils/Utils.h"
#include "Utils/PitchClass.h"

class PitchDetector : juce::Thread {
 public:
  PitchDetector(double startProgress, double endProgress);
  ~PitchDetector();

  typedef struct Pitch {
    Utils::PitchClass pitchClass;
    float posRatio;  // position in track from 0-1
    float duration;  // note duration from 0-1
    float gain;      // pitch salience
    Pitch() : pitchClass(Utils::PitchClass::NONE), posRatio(0.0), duration(0.0), gain(0.0) {}
    Pitch(Utils::PitchClass pitchClass_, float posRatio_, float duration_, float gain_)
        : pitchClass(pitchClass_), posRatio(posRatio_), duration(duration_), gain(gain_) {}
  } Pitch;

  typedef juce::HashMap<Utils::PitchClass, std::vector<Pitch>> PitchMap;

  std::function<void(Utils::SpecBuffer& hpcp)> onHarmonicProfileReady = nullptr;
  std::function<void(PitchMap& pitchMap, Utils::SpecBuffer& pitchSpec)> onPitchesReady = nullptr;
  std::function<void(double progress)> onProgressUpdated = nullptr;

  void process(const juce::AudioBuffer<float>* audioBuffer, double sampleRate);
  void cancelProcessing();

  void run() override;
  // Clear any data not used after lifetime of run()
  void clear();

 private:
  // FFT
  static constexpr int FFT_SIZE = 4096;
  static constexpr int HOP_SIZE = 512;
  // Spectral Whitening
  static constexpr double MIN_AVG_FRAME_ENERGY = 0.0001;
  // HPCP
  static constexpr int REF_FREQ = 440;
  static constexpr int MAX_SPEC_PEAKS = 60;
  static constexpr int NUM_HPCP_BINS = 120;
  static constexpr float HPCP_WINDOW_LEN = 1.0f;
  static constexpr int NUM_HARMONIC_WEIGHTS = 3;
  static constexpr int MIN_FREQ = 40;
  static constexpr int MAX_FREQ = 5000;
  static constexpr double HARMONIC_PRECISION = 0.00001;
  static constexpr double MAGNITUDE_THRESHOLD = 0.00001;
  static constexpr int PITCH_CLASS_OFFSET = 9;  // Offset from reference freq A to lowest class C
  static constexpr int PITCH_CLASS_OFFSET_BINS = (NUM_HPCP_BINS / Utils::PitchClass::COUNT) * PITCH_CLASS_OFFSET;
  // Pitch segmenting
  static constexpr int NUM_ACTIVE_SEGMENTS = 1;
  static constexpr int MAX_DEVIATION_CENTS = 15;
  static constexpr int INVALID_BIN = -1;
  static constexpr int MAX_DEVIATION_BINS = (NUM_HPCP_BINS / Utils::PitchClass::COUNT) * (MAX_DEVIATION_CENTS / 100.0);
  static constexpr double MAX_IDLE_TIME_MS = 62.5;
  static constexpr int MIN_NOTE_TIME_MS = 125;
  static constexpr int LOOKAHEAD_TIME_MS = 25;

  // Used to show far along the run thread is
  void updateProgress(double progress);
  double mStartProgress;
  double mEndProgress;
  double mDiffProgress;

  typedef struct PitchSegment {
    float binNum;      // Bin number in HPCP
    int startFrame;    // Start frame of segment
    int idleFrame;     // Start frame of when segment began being idle (or -1 when active)
    float salience;    // Confidence level accumulator from gains
    bool isAvailable;  // True when segment isn't being used to track a pitch
    PitchSegment() : binNum(0), startFrame(0), idleFrame(-1), salience(0.0), isAvailable(true) {}
  } PitchSegment;

  typedef struct Peak {
    float binNum;  // Bin number in frame
    float gain;
    Peak() : binNum(-1), gain(0) {}
    Peak(float binNum_, float gain_) : binNum(binNum_), gain(gain_) {}
  } Peak;

  typedef struct HarmonicWeight {
    float semitone;
    float gain;
    HarmonicWeight(float semitone_, float gain_) : semitone(semitone_), gain(gain_) {}
  } HarmonicWeight;

  Fft mFft;
  double mSampleRate;
  // HPCP fields
  std::vector<HarmonicWeight> mHarmonicWeights;
  Utils::SpecBuffer mHPCP;  // harmonic pitch class profile

  // Pitch segments in buffer form
  Utils::SpecBuffer mSegmentedPitches;
  std::array<PitchSegment, NUM_ACTIVE_SEGMENTS> mSegments;

  // Hashmap of detected pitches
  PitchMap mPitchMap;

  bool computeHPCP();
  bool segmentPitches();
  void getSegmentedPitchBuffer();
  bool hasBetterCandidateAhead(int startFrame, float target,
                               float deviation);  // True if a closer target is ahead
  Utils::PitchClass getPitchClass(float binNum);  // Finds the closest pitch class
  Peak interpolatePeak(int frame, int bin);
  void interpolatePeak(const float leftVal, const float middleVal, const float rightVal, int currentBin, float& resultVal,
                       float& resultBin) const;
  std::vector<Peak> getPeaks(int numPeaks, const std::vector<float>& frame);
  void initHarmonicWeights();
};
