#!/usr/bin/perl
#
# Fluffy Ball Signal Analyzer (c) 2011 Petr Baudis <pasky@ucw.cz>
# MIT licence
#
# http://brmlab.cz/project/fluffyball
#
# Example: fb.pl -r sound hw:1,0
# ...will show rectangle labels instead of frequency, and take input
# from the second soundcard instead of the first.

use warnings;
use strict;

use Math::FFT;
use List::Util qw(min max sum);
use List::MoreUtils qw(pairwise);
use Tk;


our ($render, $label_recti, $delay, $debug);
$render = 1;
if ($ARGV[0] eq '-n') {
	shift @ARGV;
	$render = 0;
}
if ($ARGV[0] eq '-r') {
	shift @ARGV;
	$label_recti = 1;
}
if ($ARGV[0] eq '-s') {
	shift @ARGV;
	$delay = shift @ARGV;
}
if ($ARGV[0] eq '-d') {
	shift @ARGV;
	$debug = 1;
}

my $mode = shift @ARGV;
my ($devname, $playfile);
if (not $mode or $mode eq 'sound') {
	($devname) = @ARGV;
	$devname ||= 'default';
} elsif ($mode eq 'replay') {
	($playfile) = @ARGV;
	$playfile or die "no file specified";
} else {
	print STDERR "Usage: $0 [-n] [-r] [-s USEC] [-d] [sound [DEVNAME]] | replay FILENAME\n";
	exit 1;
}


our ($mw, $canvas, $text);
our ($dsp);
our ($sampleno, $sampling);

our ($chan, $fmt, $rate, $sps) = (1, 16, 16384, 8);
our $bps = ($chan * $fmt * $rate) / 8;
our $buf = $bps / $sps;

our @avgsteps = (2, 4);

# How much to amplify the spectrum
our $amplifier = 1;

# Rectangles cover averages in various bands. They also include averages
# of surrounding of the rectangle picture. E.g. with step 20, over 10,
# the average includes average of 20 of the covered area + 10 on each side.
our $rect_step = 10;
our $rect_over = 5;

our $draw_lines = 0; # this makes for real pretty graphs but it is a huge slowdown
our $draw_rect = 1;


our $state = 'ticho';

sub assess {
	my ($finfo) = @_;

	# Jingle inside of our ball. If we did not have it,
	# alternative detection methods could be possible too.
	my ($jingle) = $finfo->{rect_2}->[29];

	# Look at mid-frequency average level over longer time
	# and check that it is universal in that range.
	my @midrange = 25..50;
	my ($fidgetl, $fidgetm);
	$fidgetl = sum (@{$finfo->{rect_2}}[@midrange]) / @midrange;
	$fidgetm = min (@{$finfo->{rect_2}}[@midrange]);

	# Look at dynamics of the sound; too much change means
	# this is more likely buch fallout.
	my ($dynamics);
	my (@mid_rect_now) = @{$finfo->{rect_now}}[@midrange];
	my (@mid_rect_2) = @{$finfo->{rect_4}}[@midrange];
	$dynamics = sqrt(sum (pairwise { ($a - $b) * ($a - $b) } @mid_rect_now, @mid_rect_2)) / @midrange;

	if ($debug) {
		print "(jingle $jingle) (fidget $fidgetl $fidgetm) (dynamics $dynamics)\n";
		#print join(' ', @{$finfo->{rect_4}}[29, 46, 47, 48])."\n";
	}

	my $newstate = $state;
	if ($state ne 'buch' and $jingle > 0.15) {
		$newstate = 'buch';

	# buch -> hlazeni not allowed
	} elsif ($state eq 'ticho' and ($fidgetl > 0.03 and $fidgetm > 0.02 and $dynamics < 0.005)) {
		$newstate = 'hlazeni';

	} elsif ($state eq 'hlazeni' and ($fidgetl < 0.02 or $fidgetm < 0.015 or $dynamics > 0.02)) {
		$newstate = 'ticho';

	} elsif ($state eq 'buch' and $jingle < 0.05) {
		$newstate = 'ticho';
	}

	$newstate ne $state and print "$newstate\n";
	$state = $newstate;

	if ($render) {
		$text->delete('all');
		if ($newstate ne 'ticho') {
			$text->createText(0, 20, -anchor => 'nw', -font => 'Small', -text => "[$newstate]", -fill => 'blue');
		}
	}
}

sub render_spectrum {
	my ($finfo, $label, $reflabel, $fir, $srows, $xofs, $yofs, $w, $h, $hs) = @_;
	my @fi = @$fir;
	my @f = @{$finfo->{"freq_$label"}};
	my @fref; defined $reflabel and @fref = @{$finfo->{"freq_$reflabel"}};
	my @r = @{$finfo->{"rect_$label"}};
	my @rref; defined $reflabel and @rref = @{$finfo->{"rect_$reflabel"}};
	for my $y (0..($srows-1)) {
		for my $x (0..($w-1)) {
			my $hb = $yofs + $h * ($y + 1) + $hs * $y;
			my $i = $w * $y + $x;

			if ($draw_lines) {
				my $barh = $h * ($draw_rect ? 2/3 : 1);
				my $bar = $h * $f[$i] * $amplifier;
				$bar = min ($barh, $bar);
				$canvas->createLine($xofs + $x, $hb, $xofs + $x, $hb - $bar);
				if (defined $reflabel) {
					my $refbar = $h * $fref[$i] * $amplifier;
					$refbar = min ($barh, $refbar);
					if ($refbar > $bar) {
						$canvas->createLine($xofs + $x, $hb - $bar, $xofs + $x, $hb - $refbar, -fill => 'darkred');
					} else {
						$canvas->createLine($xofs + $x, $hb - $refbar, $xofs + $x, $hb - $bar, -fill => 'red');
					}
				}
			}

			if ($draw_rect and !($i % $rect_step) and $i >= $rect_over and $i <= @f - $rect_step - $rect_over) {
				my $recth = $h * ($draw_lines ? 2/3 : 1);
				my $avg = $recth * $amplifier * $r[$i / $rect_step];
				$avg = min ($h, $avg);
				my $rhb = $hb - $h;
				$canvas->createRectangle($xofs + $x, $rhb, $xofs + $x + $rect_step, $rhb + $avg, -fill => 'black');
				if (defined $reflabel) {
					my $refavg = $recth * $amplifier * $rref[$i / $rect_step];
					$refavg = min ($h, $refavg);
					if ($refavg > $avg) {
						$canvas->createRectangle($xofs + $x, $rhb + $avg, $xofs + $x + $rect_step, $rhb + $refavg, -fill => 'darkred', -outline => 'darkred');
					} else {
						$canvas->createRectangle($xofs + $x, $rhb + $refavg, $xofs + $x + $rect_step, $rhb + $avg, -fill => 'red', -outline => 'red');
					}
				}
			}

			if (!($x % ($w/4))) {
				$canvas->createLine($xofs + $x, $hb + 0, $xofs + $x, $hb + 5, -fill => 'blue');
				$canvas->createText($xofs + $x, $hb + 15, -fill => 'blue', -font => 'small', -text => $label_recti ? $i/$rect_step : $fi[$i]);
			}
		}
	}
	$canvas->createText($xofs, $yofs, -fill => 'darkgreen', -font => 'small', -text => $label);
}

sub render {
	my ($finfo) = @_;
	my @fi = @{$finfo->{freqi}};

	# main plot
	my $srows = 32/$sps; # spectrum rows; "density" of data depends on sps
	my $w = @fi / $srows; # width
	my $ws = 20; # spacing
	my $h = 150; # height
	my $hs = 20; # spacing

	my $htot = $hs + ($h + $hs) * $srows;

	# averages plots
	my $acols = 1;
	my $arows = @avgsteps / $acols; # rows of plots; each plot is further divided to $srows rows of spectrum
	my $aw = $w;
	my $ahtot = ($htot - 3 * $hs) / $arows;
	my $ah = $ahtot / $srows - $hs;

	unless ($canvas) {
		$canvas = $mw->Canvas(-width => $ws + $w + $ws*3 + ($aw + $ws) * $acols + $ws, -height => $htot);
		$canvas->pack;
	}
	$canvas->delete('all');

	render_spectrum($finfo, 'now', $avgsteps[0], $finfo->{freqi}, $srows, $ws, $hs, $w, $h, $hs);
	for my $i (0..$#avgsteps) {
		my $ax = $i % $acols;
		my $ay = sprintf('%d', $i / $acols);
		my $su = $avgsteps[$i];
		render_spectrum($finfo, $su, $i < $#avgsteps ? $avgsteps[$i+1] : undef, $finfo->{freqi}, $srows,
				$ws + $w + $ws*3 + ($ws + $aw) * $ax,
				($ahtot + $hs*2) * $ay + $hs,
				$aw, $ah, $hs);
	}

	$mw->update();
}

sub getdata {
	my $w;
	read $dsp, $w, $buf or die "read: $!";
	return $w;
}

# Array with frequencies corresponding to elements of fftsig
sub fftfreq {
	my @freqs;
	my $dft_size = $rate / $sps;
	for (my $i = 0; $i < $dft_size / 2; $i++) {
		$freqs[$i] = $i / $dft_size * $rate;
	}
	return @freqs;
}

sub fftsig {
	my ($bytes) = @_;

	my @samples;
	if ($fmt == 8) {
		while (length($bytes) > 0) {
			my $sample = unpack('c', substr($bytes, 0, 1, ''));
			push(@samples, $sample);
		}
	} elsif ($fmt == 16) {
		while (length($bytes) > 0) {
			my $sample = unpack('s<', substr($bytes, 0, 2, ''));
			push(@samples, $sample);
		}
	} else {
		die "unsupported $fmt bits per sample\n";
	}

	# Magic from Audio::Analyze:
	my $fft = Math::FFT->new(\@samples);
	my $coeff = $fft->rdft;
	my $size = scalar(@$coeff);
	my $k = 0;
	my @mag;

	$mag[$k] = sqrt($coeff->[$k*2]**2);
	for($k = 1; $k < $size / 2; $k++) {
		$mag[$k] = sqrt(($coeff->[$k * 2] ** 2) + ($coeff->[$k * 2 + 1] ** 2));
	}

	# Rescale to 0..1
	my $avgmag = sum (@mag) / @mag; #[1000..$#mag];
	@mag = map { $_ / $avgmag * 0.3 } @mag;
	return @mag;
}

sub rectsof {
	my ($finfo, $l) = @_;
	my @f = @{$finfo->{"freq_$l"}};
	my @r;
	for (my ($fi, $ri) = ($rect_step, 1); $fi <= $#f - $rect_step - $rect_over; $fi += $rect_step, $ri++) {
		$r[$ri] = sum (@f[$fi - $rect_over .. $fi + $rect_step + $rect_over]) / ($rect_step + 2 * $rect_over);
	}
	$finfo->{"rect_$l"} = \@r;
}

sub ticks {
	my $s = 0;
	my $finfo;

	while (1) {
		my $w = getdata();
		defined $sampling and $sampling .= $w;
		$finfo->{freqi} = [fftfreq()];
		$finfo->{freq_now} = [fftsig($w)];

		# Create averaged rects
		rectsof($finfo, 'now');

		# Update floating freq. averages
		for my $su (@avgsteps) {
			for my $i (0..$#{$finfo->{freqi}}) {
				my $fn = $finfo->{freq_now}->[$i];
				my $fs = $finfo->{"freq_$su"}->[$i];
				$fs ||= $fn;
				$finfo->{"freq_$su"}->[$i] = ($fs * ($su - 1) + $fn) / $su;
			}
			rectsof($finfo, $su);
		}

		assess($finfo);
		$render and render($finfo);

		if ($delay) {
			use Time::HiRes qw(usleep);
			usleep($delay);
		}

		$render and $mw->idletasks();

		$s++;
	}
}


sub recStart {
	if (defined $sampling) {
		recCancel();
		return;
	}
	$sampleno++;
	$sampling = '';
	$text->delete('all');
	$text->createText(0, 20, -anchor => 'nw', -font => 'Small', -text => "Recording $sampleno", -fill => 'red');
}

sub recCancel {
	$sampleno--;
	$sampling = undef;
	$text->delete('all');
}

sub recSave {
	open my $fh, '>'.sprintf('sample-%04d.raw', $sampleno) or die "$!";
	print $fh $sampling;
	close $fh;
	$sampling = undef;
	$text->delete('all');
}


if ($mode eq 'replay') {
	open ($dsp, $playfile) or die "$playfile: $!";
} else {
	# '-c', $chan,
	print join(' ', 'arecord', '-D', $devname, '-t', 'raw', '-r', $rate, '-f', 'S'.$fmt)."\n";
	open ($dsp, '-|', 'arecord', '-D', $devname, '-t', 'raw', '-r', $rate, '-f', 'S'.$fmt) or die "arecord: $!";
	use IO::Handle;
	$dsp->autoflush(1);
}

if ($render) {
	$mw = MainWindow->new;
	$mw->bind('<space>' => \&recStart);
	$mw->bind('<Escape>' => \&recCancel);
	$mw->bind('<Return>' => \&recSave);
	$text = $mw->Canvas(-width => 400, -height => 40);
	$text->pack;

	$mw->after(1, \&ticks);
	MainLoop;
} else {
	ticks();
}

