let loader = new dl.CheckpointLoader(
  'https://teropa.info/drums_rnn_checkpoint'
);

const DRUM_CLASSES = [
  'Kick',
  'Snare',
  'Hi-hat closed',
  'Hi-hat open',
  'Tom low',
  'Tom mid',
  'Tom high',
  'Clap',
  'Rim'
];
const TIME_HUMANIZATION = 0.01;

let sampleBaseUrl = 'https://s3-us-west-2.amazonaws.com/s.cdpn.io/969699';

let reverb = new Tone.Convolver(
  `${sampleBaseUrl}/small-drum-room.wav`
).toMaster();
reverb.wet.value = 0.35;

let snarePanner = new Tone.Panner().connect(reverb);
new Tone.LFO(0.13, -0.25, 0.25).connect(snarePanner.pan).start();

let drumKit = [
  new Tone.Players({
    high: `${sampleBaseUrl}/808-kick-vh.mp3`,
    med: `${sampleBaseUrl}/808-kick-vm.mp3`,
    low: `${sampleBaseUrl}/808-kick-vl.mp3`
  }).toMaster(),
  new Tone.Players({
    high: `${sampleBaseUrl}/flares-snare-vh.mp3`,
    med: `${sampleBaseUrl}/flares-snare-vm.mp3`,
    low: `${sampleBaseUrl}/flares-snare-vl.mp3`
  }).connect(snarePanner),
  new Tone.Players({
    high: `${sampleBaseUrl}/808-hihat-vh.mp3`,
    med: `${sampleBaseUrl}/808-hihat-vm.mp3`,
    low: `${sampleBaseUrl}/808-hihat-vl.mp3`
  }).connect(new Tone.Panner(-0.5).connect(reverb)),
  new Tone.Players({
    high: `${sampleBaseUrl}/808-hihat-open-vh.mp3`,
    med: `${sampleBaseUrl}/808-hihat-open-vm.mp3`,
    low: `${sampleBaseUrl}/808-hihat-open-vl.mp3`
  }).connect(new Tone.Panner(-0.5).connect(reverb)),
  new Tone.Players({
    high: `${sampleBaseUrl}/slamdam-tom-low-vh.mp3`,
    med: `${sampleBaseUrl}/slamdam-tom-low-vm.mp3`,
    low: `${sampleBaseUrl}/slamdam-tom-low-vl.mp3`
  }).connect(new Tone.Panner(-0.4).connect(reverb)),
  new Tone.Players({
    high: `${sampleBaseUrl}/slamdam-tom-mid-vh.mp3`,
    med: `${sampleBaseUrl}/slamdam-tom-mid-vm.mp3`,
    low: `${sampleBaseUrl}/slamdam-tom-mid-vl.mp3`
  }).connect(reverb),
  new Tone.Players({
    high: `${sampleBaseUrl}/slamdam-tom-high-vh.mp3`,
    med: `${sampleBaseUrl}/slamdam-tom-high-vm.mp3`,
    low: `${sampleBaseUrl}/slamdam-tom-high-vl.mp3`
  }).connect(new Tone.Panner(0.4).connect(reverb)),
  new Tone.Players({
    high: `${sampleBaseUrl}/909-clap-vh.mp3`,
    med: `${sampleBaseUrl}/909-clap-vm.mp3`,
    low: `${sampleBaseUrl}/909-clap-vl.mp3`
  }).connect(new Tone.Panner(0.5).connect(reverb)),
  new Tone.Players({
    high: `${sampleBaseUrl}/909-rim-vh.wav`,
    med: `${sampleBaseUrl}/909-rim-vm.wav`,
    low: `${sampleBaseUrl}/909-rim-vl.wav`
  }).connect(new Tone.Panner(0.5).connect(reverb))
];

let attentionLength = 32;
let attentionSize = 256;
let dropoutKeepProb = dl.Scalar.new(1.0);
let temperature = dl.variable(dl.Scalar.new(1.2));

/*
 * Magenta's model encodes each timestep as a one-hot vector of length 512.
 * There are 9 drum classes, and any combination of those 9 classes can play simultaneously.
 * Each of the 2^9 = 512 possible combinations has its own element in the vector.
 * 
 * The one-hot vector is then augmented with six meter-tracking bits, implemented as binary counters.
 */

function encodeDrums(drumClasses, nextPosition) {
  let hotIndex = 0;
  for (let drumClass of drumClasses) {
    hotIndex += 2 ** drumClass;
  }
  let vector = _.times(2 ** DRUM_CLASSES.length, () => 0);
  vector[hotIndex] = 1;

  for (let i = 0; i < 6; i++) {
    vector.push(Math.floor(nextPosition / 2 ** i) % 2 ? 1.0 : -1.0);
  }

  return vector;
}

function decodeDrums(index) {
  let bin = (index >>> 0).toString(2);
  let drums = [];
  for (let i = bin.length - 1; i >= 0; i--) {
    if (bin[i] === '1') {
      drums.push(bin.length - i - 1);
    }
  }
  return drums;
}

Promise.all([
  loader.getAllVariables(),
  new Promise(res => Tone.Buffer.on('load', res))
]).then(([vars]) => {
  let prefix = 'RNN/AttentionCellWrapper/';
  let lstm = {
    attnInputMatrix: vars[`${prefix}Linear/Matrix`],
    attnInputBias: vars[`${prefix}Linear/Bias`].as2D(1, -1),
    attnW: vars[`${prefix}Attention/AttnW`],
    attnV: vars[`${prefix}Attention/AttnV`],
    attnMatrix: vars[`${prefix}Attention/Linear/Matrix`],
    attnBias: vars[`${prefix}Attention/Linear/Bias`],
    attnOutputMatrix: vars[`${prefix}AttnOutputProjection/Linear/Matrix`],
    attnOutputBias: vars[`${prefix}AttnOutputProjection/Linear/Bias`],
    kernel1: vars[`${prefix}MultiRNNCell/Cell0/BasicLSTMCell/Linear/Matrix`],
    bias1: vars[`${prefix}MultiRNNCell/Cell0/BasicLSTMCell/Linear/Bias`],
    kernel2: vars[`${prefix}MultiRNNCell/Cell1/BasicLSTMCell/Linear/Matrix`],
    bias2: vars[`${prefix}MultiRNNCell/Cell1/BasicLSTMCell/Linear/Bias`],
    kernel3: vars[`${prefix}MultiRNNCell/Cell2/BasicLSTMCell/Linear/Matrix`],
    bias3: vars[`${prefix}MultiRNNCell/Cell2/BasicLSTMCell/Linear/Bias`],
    fullyConnectedBiases: vars['fully_connected/biases'],
    fullyConnectedWeights: vars['fully_connected/weights']
  };

  function initNnState() {
    return {
      lstmStates: [
        dl.Array2D.zeros([1, lstm.bias1.shape[0] / 4]),
        dl.Array2D.zeros([1, lstm.bias2.shape[0] / 4]),
        dl.Array2D.zeros([1, lstm.bias3.shape[0] / 4])
      ],
      lstmOutputs: [
        dl.Array2D.zeros([1, lstm.bias1.shape[0] / 4]),
        dl.Array2D.zeros([1, lstm.bias2.shape[0] / 4]),
        dl.Array2D.zeros([1, lstm.bias3.shape[0] / 4])
      ],
      attention: dl.Array1D.zeros([attentionSize]),
      attentionState: dl.Array2D.zeros([1, attentionSize * attentionLength])
    };
  }

  function disposeNnState(state) {
    state.lstmStates.forEach(s => s.dispose());
    state.lstmOutputs.forEach(s => s.dispose());
    state.attention.dispose();
    state.attentionState.dispose();
  }

  function generateStep(
    drumClasses,
    nextPosition,
    { lstmStates, lstmOutputs, attention, attentionState }
  ) {
    return dl.tidy(async () => {
      let lstm1 = dl.basicLSTMCell.bind(
        dl,
        dropoutKeepProb,
        lstm.kernel1,
        lstm.bias1
      );
      let lstm2 = dl.basicLSTMCell.bind(
        dl,
        dropoutKeepProb,
        lstm.kernel2,
        lstm.bias2
      );
      let lstm3 = dl.basicLSTMCell.bind(
        dl,
        dropoutKeepProb,
        lstm.kernel3,
        lstm.bias3
      );

      let input = dl.Array1D.new(encodeDrums(drumClasses, nextPosition));
      let attnStates = attentionState.reshape([
        -1,
        attentionLength,
        attentionSize
      ]);

      // Apply attention to input

      let rnnInput = dl.add(
        dl.matMul(
          dl.concat([input, attention]).as2D(1, -1),
          lstm.attnInputMatrix
        ),
        lstm.attnInputBias
      );

      // LSTM

      let [newLstmStates, newLstmOutputs] = dl.multiRNNCell(
        [lstm1, lstm2, lstm3],
        rnnInput,
        lstmStates,
        lstmOutputs
      );

      // Attention

      let attnHidden = dl.reshape(attentionState, [
        -1,
        attentionLength,
        1,
        attentionSize
      ]);

      let attnHiddenFeatures = dl.conv2d(
        attnHidden,
        lstm.attnW,
        [1, 1],
        'same'
      );

      let attnQuery = dl.concat(
        [
          newLstmStates[0],
          newLstmOutputs[0],
          newLstmStates[1],
          newLstmOutputs[1],
          newLstmStates[2],
          newLstmOutputs[2]
        ],
        1
      );

      let attnY = dl
        .matMul(attnQuery, lstm.attnMatrix)
        .reshape([-1, 1, 1, attentionSize]);
      let attnS = dl.sum(
        dl.mul(lstm.attnV, dl.tanh(dl.add(attnHiddenFeatures, attnY))),
        [2, 3]
      );
      let attnA = dl.softmax(attnS);
      let attnD = dl.sum(
        dl.mul(dl.reshape(attnA, [-1, attentionLength, 1, 1]), attnHidden),
        [1, 2]
      );
      let newAttns = attnD.reshape([-1, attentionSize]);

      // Remove first (oldest) attention state
      let newAttnStates = dl.slice(
        attnStates,
        [0, 1, 0],
        [attnStates.shape[0], attnStates.shape[1] - 1, attnStates.shape[2]]
      );

      // Output projection

      let newOutput = dl.add(
        dl.matMul(
          dl.concat([newLstmOutputs[2], newAttns], 1),
          lstm.attnOutputMatrix
        ),
        lstm.attnOutputBias
      );

      // Add latest output to attention state
      newAttnStates = dl
        .concat(
          [
            newAttnStates,
            newOutput.as3D(newOutput.shape[0], 1, newOutput.shape[1])
          ],
          1
        )
        .reshape([-1, attentionLength * attentionSize]);

      lstmStates.forEach(s => s.dispose());
      lstmOutputs.forEach(o => o.dispose());
      attention.dispose();
      attentionState.dispose();
      lstmStates = newLstmStates;
      lstmOutputs = newLstmOutputs;
      attention = newAttns.flatten();
      attentionState = newAttnStates;
      lstmStates.forEach(s => dl.keep(s));
      lstmOutputs.forEach(o => dl.keep(o));
      dl.keep(attention);
      dl.keep(attentionState);

      let weightedResult = dl.matMul(newOutput, lstm.fullyConnectedWeights);
      let logits = dl.add(weightedResult, lstm.fullyConnectedBiases);
      let softmax = dl.softmax(dl.div(logits.as1D(), temperature));
      let sampledOutput = dl.ENV.math.multinomial(softmax, 1).asScalar();
      return {
        drums: decodeDrums(await sampledOutput.val()),
        nnState: { lstmStates, lstmOutputs, attention, attentionState }
      };
    });
  }

  let patternLength = 16,
    seedLength = 4,
    swing = 0.55,
    stepEls = [],
    pattern = [[0], [], [2]].concat(_.times(patternLength - 3, i => [])),
    hasBeenStarted = false,
    sequence;

  function generatePattern(
    result,
    fromStep = 0,
    currentStep = 0,
    nnState = initNnState()
  ) {
    let nextStep = currentStep + 1;
    return generateStep(result[currentStep], nextStep, nnState).then(
      ({ drums, nnState }) => {
        if (nextStep >= fromStep) {
          result[nextStep] = drums;
        }
        if (nextStep < patternLength - 1) {
          return generatePattern(result, fromStep, nextStep, nnState);
        } else {
          disposeNnState(nnState);
          return result;
        }
      }
    );
  }

  function getStepVelocity(step) {
    if (step % 4 === 0) {
      return 'high';
    } else if (step % 2 === 0) {
      return 'med';
    } else {
      return 'low';
    }
  }

  function humanizeTime(time) {
    return time - TIME_HUMANIZATION / 2 + Math.random() * TIME_HUMANIZATION;
  }

  function playPattern() {
    sequence = new Tone.Sequence(
      (time, { drums, stepIdx }) => {
        let isSwung = stepIdx % 2 !== 0;
        if (isSwung) {
          time += (swing - 0.5) * Tone.Time('8n').toSeconds();
        }
        let velocity = getStepVelocity(stepIdx);
        drums.forEach(d => {
          let humanizedTime = stepIdx === 0 ? time : humanizeTime(time);
          drumKit[d].get(velocity).start(humanizedTime);
          visualizePlay(humanizedTime, stepIdx, d);
        });
      },
      pattern.map((drums, stepIdx) => ({ drums, stepIdx })),
      '16n'
    ).start();
  }

  function visualizePlay(time, stepIdx, drumIdx) {
    Tone.Draw.schedule(() => {
      if (!stepEls[stepIdx]) return;
      let animTime = Tone.Time('2n').toMilliseconds();
      let cellEl = stepEls[stepIdx].cellEls[drumIdx];
      if (cellEl.classList.contains('on')) {
        let baseColor = stepIdx < seedLength ? '#e91e63' : '#64b5f6';
        cellEl.animate(
          [
            {
              transform: 'translateZ(-100px)',
              backgroundColor: '#fad1df'
            },
            {
              transform: 'translateZ(50px)',
              offset: 0.7
            },
            { transform: 'translateZ(0)', backgroundColor: baseColor }
          ],
          { duration: animTime, easing: 'cubic-bezier(0.23, 1, 0.32, 1)' }
        );
      }
    }, time);
  }

  function renderPattern(regenerating = false) {
    let seqEl = document.querySelector('.sequencer .steps');
    while (stepEls.length > pattern.length) {
      let { stepEl, gutterEl } = stepEls.pop();
      stepEl.remove();
      if (gutterEl) gutterEl.remove();
    }
    for (let stepIdx = 0; stepIdx < pattern.length; stepIdx++) {
      let step = pattern[stepIdx];
      let stepEl, gutterEl, cellEls;
      if (stepEls[stepIdx]) {
        stepEl = stepEls[stepIdx].stepEl;
        gutterEl = stepEls[stepIdx].gutterEl;
        cellEls = stepEls[stepIdx].cellEls;
      } else {
        stepEl = document.createElement('div');
        stepEl.classList.add('step');
        stepEl.dataset.stepIdx = stepIdx;
        seqEl.appendChild(stepEl);
        cellEls = [];
      }

      stepEl.style.flex = stepIdx % 2 === 0 ? swing : 1 - swing;

      if (!gutterEl && stepIdx < pattern.length - 1) {
        gutterEl = document.createElement('div');
        gutterEl.classList.add('gutter');
        seqEl.insertBefore(gutterEl, stepEl.nextSibling);
      } else if (gutterEl && stepIdx >= pattern.length) {
        gutterEl.remove();
        gutterEl = null;
      }

      if (gutterEl && stepIdx === seedLength - 1) {
        gutterEl.classList.add('seed-marker');
      } else if (gutterEl) {
        gutterEl.classList.remove('seed-marker');
      }

      for (let cellIdx = 0; cellIdx < DRUM_CLASSES.length; cellIdx++) {
        let cellEl;
        if (cellEls[cellIdx]) {
          cellEl = cellEls[cellIdx];
        } else {
          cellEl = document.createElement('div');
          cellEl.classList.add('cell');
          cellEl.classList.add(_.kebabCase(DRUM_CLASSES[cellIdx]));
          cellEl.dataset.stepIdx = stepIdx;
          cellEl.dataset.cellIdx = cellIdx;
          stepEl.appendChild(cellEl);
          cellEls[cellIdx] = cellEl;
        }
        if (step.indexOf(cellIdx) >= 0) {
          cellEl.classList.add('on');
        } else {
          cellEl.classList.remove('on');
        }
      }
      stepEls[stepIdx] = { stepEl, gutterEl, cellEls };

      let stagger = stepIdx * (300 / (patternLength - seedLength));
      setTimeout(() => {
        if (stepIdx < seedLength) {
          stepEl.classList.add('seed');
        } else {
          stepEl.classList.remove('seed');
          if (regenerating) {
            stepEl.classList.add('regenerating');
          } else {
            stepEl.classList.remove('regenerating');
          }
        }
      }, stagger);
    }

    setTimeout(repositionRegenerateButton, 0);
  }

  function repositionRegenerateButton() {
    let regenButton = document.querySelector('.regenerate');
    let sequencerEl = document.querySelector('.sequencer');
    let seedMarkerEl = document.querySelector('.gutter.seed-marker');
    let regenLeft =
      sequencerEl.offsetLeft +
      seedMarkerEl.offsetLeft +
      seedMarkerEl.offsetWidth / 2 -
      regenButton.offsetWidth / 2;
    let regenTop =
      sequencerEl.offsetTop +
      seedMarkerEl.offsetTop +
      seedMarkerEl.offsetHeight / 2 -
      regenButton.offsetHeight / 2;
    regenButton.style.left = `${regenLeft}px`;
    regenButton.style.top = `${regenTop}px`;
    regenButton.style.visibility = 'visible';
  }

  function regenerate() {
    let result = _.take(pattern, seedLength).concat(
      _.times(length - seedLength, () => [])
    );
    renderPattern(true);
    return generatePattern(result, seedLength).then(result => {
      pattern = result;
      onPatternUpdated();
    });
  }

  function onPatternUpdated() {
    if (sequence) {
      sequence.dispose();
      sequence = null;
    }
    renderPattern();
  }

  function toggleStep(cellEl) {
    if (pattern && cellEl.classList.contains('cell')) {
      let stepIdx = +cellEl.dataset.stepIdx;
      let cellIdx = +cellEl.dataset.cellIdx;
      let isOn = cellEl.classList.contains('on');
      if (isOn) {
        _.pull(pattern[stepIdx], cellIdx);
        cellEl.classList.remove('on');
      } else {
        pattern[stepIdx].push(cellIdx);
        cellEl.classList.add('on');
      }
      sequence.at(stepIdx, { stepIdx, drums: pattern[stepIdx] });
    }
  }

  function setSwing(newSwing) {
    swing = newSwing;
    renderPattern();
  }

  function setPatternLength(newPatternLength) {
    if (newPatternLength < patternLength) {
      pattern.length = newPatternLength;
    } else {
      for (let i = pattern.length; i < newPatternLength; i++) {
        pattern.push([]);
      }
    }
    let lengthRatio = newPatternLength / patternLength;
    seedLength = Math.max(
      1,
      Math.min(newPatternLength - 1, Math.round(seedLength * lengthRatio))
    );
    patternLength = newPatternLength;
    onPatternUpdated();
    if (Tone.Transport.state === 'started') {
      playPattern();
    }
  }

  function updatePlayPauseIcons() {
    if (Tone.Transport.state === 'started') {
      document.querySelector('.playpause .pause-icon').style.display = null;
      document.querySelector('.playpause .play-icon').style.display = 'none';
    } else {
      document.querySelector('.playpause .play-icon').style.display = null;
      document.querySelector('.playpause .pause-icon').style.display = 'none';
    }
  }

  document.querySelector('.app').addEventListener('click', event => {
    if (event.target.classList.contains('cell')) {
      toggleStep(event.target);
    }
  });
  document.querySelector('.regenerate').addEventListener('click', event => {
    event.currentTarget.classList.remove('pulse');
    regenerate().then(() => {
      if (!hasBeenStarted) {
        Tone.context.resume();
        Tone.Transport.start();
        updatePlayPauseIcons();
        hasBeenStarted = true;
      }
      if (Tone.Transport.state === 'started') {
        playPattern();
      }
    });
  });
  document.querySelector('.playpause').addEventListener('click', event => {
    if (Tone.Transport.state !== 'started') {
      Tone.context.resume();
      Tone.Transport.start();
      playPattern();
      updatePlayPauseIcons();
      hasBeenStarted = true;
    } else {
      if (sequence) {
        sequence.dispose();
        sequence = null;
      }
      Tone.Transport.pause();
      updatePlayPauseIcons();
    }
  });

  let draggingSeedMarker = false;
  document.querySelector('.app').addEventListener('mousedown', evt => {
    let el = evt.target;
    if (
      el.classList.contains('gutter') &&
      el.classList.contains('seed-marker')
    ) {
      draggingSeedMarker = true;
      evt.preventDefault();
    }
  });
  document.querySelector('.app').addEventListener('mouseup', () => {
    draggingSeedMarker = false;
  });
  document.querySelector('.app').addEventListener('mouseover', evt => {
    if (draggingSeedMarker) {
      let el = evt.target;
      while (el) {
        if (el.classList.contains('step')) {
          let stepIdx = +el.dataset.stepIdx;
          if (stepIdx > 0) {
            seedLength = stepIdx;
            renderPattern();
          }
          break;
        }
        el = el.parentElement;
      }
    }
  });
  document
    .querySelector('#swing')
    .addEventListener('input', evt => setSwing(+evt.target.value));
  document
    .querySelector('#temperature')
    .addEventListener('input', evt =>
      temperature.assign(dl.scalar(+evt.target.value))
    );
  document
    .querySelector('#pattern-length')
    .addEventListener('change', evt => setPatternLength(+evt.target.value));
  document
    .querySelector('#tempo')
    .addEventListener(
      'input',
      evt => (Tone.Transport.bpm.value = +evt.target.value)
    );

  $('select')
    .on('change', evt => setPatternLength(+evt.target.value))
    .material_select();

  window.addEventListener('resize', repositionRegenerateButton);

  renderPattern();

  document.querySelector('.progress').remove();
  document.querySelector('.app').style.display = null;
});