musiclight2/main.c

/*
 * vim: sw=2 ts=2 expandtab
 *
 * "THE PIZZA-WARE LICENSE" (derived from "THE BEER-WARE LICENCE"):
 * Thomas Kolb <cfr34k@tkolb.de> wrote this file. As long as you retain this
 * notice you can do whatever you want with this stuff. If we meet some day,
 * and you think this stuff is worth it, you can buy me a pizza in return.
 * - Thomas Kolb
 */

#include <pthread.h>
#include <semaphore.h>

#include <stdio.h>
#include <math.h>
#include <string.h>
#include <stdint.h>
#include <malloc.h>

#include <lua5.3/lua.h>
#include <lua5.3/lualib.h>
#include <lua5.3/lauxlib.h>

#include "lua_utils.h"
#include "lua_wrappers.h"

#include "fft.h"
#include "utils.h"
#include "sk6812.h"

#include "config.h"

// Frames per second
#define FPS ((value_type)BUFFER_PARTS * SAMPLE_RATE / BLOCK_LEN)

// Number of new samples put into the buffer each frame
#define READ_SAMPLES (BLOCK_LEN / BUFFER_PARTS)

value_type fft[BLOCK_LEN];
value_type rms;
value_type lastUpdateTime = 0;

sample signal[BLOCK_LEN];

sem_t  fftSemaphore;

struct sk6812_ctx sk6812;

int num_modules;
int num_strips;

int running = 1;

void* fft_thread(void *param) {
	sample       buffer[BLOCK_LEN];
	sample       block[BLOCK_LEN];
	value_type   fftOutReal[BLOCK_LEN], fftOutImag[BLOCK_LEN];
	value_type   tmpFFT[BLOCK_LEN];
  value_type   tmpRMS;

	double       nextFrame = get_hires_time() + 0.05;
	double       curTime;

  int i;

  init_fft();

  while(running) {
    // shift the buffer left
    memmove(buffer, buffer + READ_SAMPLES,
        (BLOCK_LEN - READ_SAMPLES) * sizeof(sample));

    size_t ret = fread(buffer + (BLOCK_LEN - READ_SAMPLES),
        sizeof(sample), READ_SAMPLES, stdin);
    if(ret != READ_SAMPLES) {
      break;
    }

    memcpy(block, buffer, BLOCK_LEN * sizeof(sample));

    tmpRMS = 0;
    for(i = 0; i < BLOCK_LEN; i++) {
      tmpRMS += buffer[i]*buffer[i];
    }
    tmpRMS = sqrt(tmpRMS/BLOCK_LEN);

    if(tmpRMS == 0) {
      continue;
    }

    apply_hanning(block);
    fft_transform(block, fftOutReal, fftOutImag);
    complex_to_absolute(fftOutReal, fftOutImag, tmpFFT);

    // --- SAFE SECTION ---
    sem_wait(&fftSemaphore);

    memcpy(fft, tmpFFT, sizeof(fft));
    memcpy(signal, buffer, sizeof(signal));
    rms = tmpRMS;

    curTime = get_hires_time();
    lastUpdateTime = curTime;

    sem_post(&fftSemaphore);
    // --- END SAFE SECTION ---

    if(curTime > nextFrame + 0.05) {
      printf("Frame too late! Skipping.\n");
      nextFrame = -1;
    }

    if(curTime < nextFrame - 0.05) {
      printf("Frame too early.\n");
      nextFrame = -1;
    }

    if(nextFrame < 0) {
      printf("Frame time reset.\n");
      nextFrame = curTime;
    }

    nextFrame += 1.000/FPS;
    //sleep_until(nextFrame);
  }

  return NULL;
}

value_type gamma_correct(value_type d, value_type gamma) {
  return pow(d, gamma);
}

int main(int argc, char **argv) {
  double nextFrame = get_hires_time() + LED_INTERVAL;

  int i;
  pthread_t fftThread;

  int active = 1;
  //int frame = 0;

  double *red;
  double *green;
  double *blue;
  double *white;

  int useFading, fadeStep;

  if(argc < 2) {
    fprintf(stderr, "LUA script file must be given as command line argument!\n");
    return 1;
  }

  // initialize lua
  lua_State *L = luaL_newstate();

  // load the lua libraries
  luaL_openlibs(L);

  // register local functions
  lua_register_funcs(L);

  // load the configuration from "config.lua"
  if(luaL_dofile(L, "config.lua")) {
    lua_showerror(L, "luaL_dofile(config.lua) failed.");
    return 1;
  }

  lua_getglobal(L, "WS2801_HOST");
  if(!lua_isstring(L, -1)) return 2;
  const char *host = lua_tostring(L, -1);

  lua_getglobal(L, "WS2801_PORT");
  if(!lua_isnumber(L, -1)) return 2;
  unsigned short port = lua_tointeger(L, -1);

  lua_getglobal(L, "GAMMA");
  if(!lua_isnumber(L, -1)) return 2;
  double gamma = lua_tonumber(L, -1);

  lua_getglobal(L, "NUM_MODULES");
  if(!lua_isnumber(L, -1)) return 2;
  num_modules = lua_tointeger(L, -1);

  lua_getglobal(L, "NUM_STRIPS");
  if(!lua_isnumber(L, -1)) return 2;
  num_strips = lua_tointeger(L, -1);

  lua_getglobal(L, "CENTER_MODULE");
  if(!lua_isnumber(L, -1)) return 2;
  int center_module = lua_tointeger(L, -1);

  // export some global constants
  lua_pushnumber(L, SAMPLE_RATE);
  lua_setglobal(L, "SAMPLE_RATE");

  lua_pushnumber(L, BLOCK_LEN);
  lua_setglobal(L, "BLOCK_LEN");

  lua_pushnumber(L, DATALEN);
  lua_setglobal(L, "DATALEN");

  // allocate arrays
  red = malloc(num_strips * num_modules * sizeof(double));
  green = malloc(num_strips * num_modules * sizeof(double));
  blue = malloc(num_strips * num_modules * sizeof(double));
  white = malloc(num_strips * num_modules * sizeof(double));

  // load and initialize the script
  if(luaL_loadfile(L, argv[1])) {
    lua_showerror(L, "luaL_loadfile(cmdline_argument) failed.");
  }

  // priming call: read the lua file to make functions known
  if(lua_pcall(L, 0, 0, 0)) {
    lua_showerror(L, "lua_pcall failed.");
  }

  // call the init function
  lua_getglobal(L, "init");
  lua_pushnumber(L, num_strips);
  lua_pushnumber(L, num_modules);
  lua_pushnumber(L, center_module);
  if(lua_pcall(L, 3, 1, 0)) {
    lua_showerror(L, "lua_pcall(init) failed.");
  }

  fadeStep = lua_tointeger(L, -1);
  useFading = fadeStep > 0;
  if(useFading) {
    sk6812_set_fadestep(&sk6812, fadeStep);
    printf("Fading enabled with fadestep %i.\n", fadeStep);
  }

  // initialize the WS2801 library
  printf("Connecting to %s:%i\n", host, port);
  sk6812_init(&sk6812, host, port);

  // create semaphores
  sem_init(&fftSemaphore, 0, 1);

  // run the fft thread
  pthread_create(&fftThread, NULL, fft_thread, NULL);

  while(running) {
    if(active) {
      // call the periodic() function from LUA
      lua_getglobal(L, "periodic");
      if(lua_pcall(L, 0, 4, 0)) { // no arguments, 4 return values
        lua_showerror(L, "lua_pcall(periodic) failed.");
      }

      // read the return values (reverse order, as lua uses a stack)
      lua_readdoublearray(L, white, num_strips*num_modules);
      lua_readdoublearray(L, blue, num_strips*num_modules);
      lua_readdoublearray(L, green, num_strips*num_modules);
      lua_readdoublearray(L, red, num_strips*num_modules);

      for(int s = 0; s < num_strips; s++) {
        if(useFading) {
          for(i = 0; i < num_modules; i++) {
            int lidx = idx(s, i);
            sk6812_fade_color(&sk6812, s, i,
                255 * gamma_correct(red[lidx], gamma),
                255 * gamma_correct(green[lidx], gamma),
                255 * gamma_correct(blue[lidx], gamma),
                255 * gamma_correct(white[lidx], gamma));
          }
        } else {
          for(i = 0; i < num_modules; i++) {
            int lidx = idx(s, i);
            sk6812_set_color(&sk6812, s, i,
                255 * gamma_correct(red[lidx], gamma),
                255 * gamma_correct(green[lidx], gamma),
                255 * gamma_correct(blue[lidx], gamma),
                255 * gamma_correct(white[lidx], gamma));
          }
        }
      }

      sk6812_commit(&sk6812);

      if(lastUpdateTime < nextFrame - 1) {
        printf("Idle for 1 second -> stopping updates.\n");

        for(i = 0; i < num_modules; i++) {
          for(int s = 0; s < num_strips; s++) {
            sk6812_fade_color(&sk6812, s, i, 0, 0, 0, 20);
          }
        }

        sk6812_commit(&sk6812);

        active = 0;
      }
    } else if(lastUpdateTime > nextFrame - 1) {
      printf("Resuming updates.\n");
      active = 1;
    }

    nextFrame += LED_INTERVAL;
    sleep_until(nextFrame);
  }

  sk6812_shutdown(&sk6812);

  // free arrays
  free(red);
  free(green);
  free(blue);
  free(white);

  pthread_join(fftThread, NULL);

  lua_close(L);

  return 0;
}