GS8208.rs

use std::{
    f64::consts::PI,
    io::Write,
    time::{Duration, Instant},
};

use humantime::parse_duration;
use spidev::{SpiModeFlags, Spidev, SpidevOptions};
use structopt::StructOpt;

#[derive(StructOpt)]
struct Options {
    /// Intensity of the red color.
    #[structopt(short, long)]
    red: u8,

    /// Intensity of the green color.
    #[structopt(short, long)]
    green: u8,

    /// Intensity of the blue color.
    #[structopt(short, long)]
    blue: u8,

    /// Optional multiplier to make all the colors dimmer.
    #[structopt(long, default_value = "1")]
    multiplier: f64,

    /// Pulsation interval.
    #[structopt(long, short, default_value = "500 ms", parse(try_from_str = parse_duration))]
    interval: Duration,

    /// Frequency
    #[structopt(long, default_value = "800000")]
    frequency: u32,

    /// Leds count.
    #[structopt(long, default_value = "50")]
    leds_count: usize,

    /// How fast the wave moves.
    #[structopt(long, short, default_value = "100 ms", parse(try_from_str = parse_duration))]
    moving_interval: Duration,

    /// Wave width.
    #[structopt(long, default_value = "10")]
    wave_width: usize,
}

/// # Safety
///
/// The behavior is undefined when `output` is shorter than 4 bytes.
unsafe fn write_spi(output: &mut [u8], byte: u8) {
    debug_assert!(output.len() >= 4);

    const ZERO_ZERO: u8 = 0b1000_1000;
    const ZERO_ONE: u8 = 0b1000_1110;
    const ONE_ZERO: u8 = 0b1110_1000;
    const ONE_ONE: u8 = 0b1110_1110;

    const MAP: [u8; 4] = [ZERO_ZERO, ZERO_ONE, ONE_ZERO, ONE_ONE];

    let first = (byte & 0b_1100_0000) >> 6;
    let first = *MAP.get_unchecked(usize::from(first));

    let second = (byte & 0b_0011_0000) >> 4;
    let second = *MAP.get_unchecked(usize::from(second));

    let third = (byte & 0b_0000_1100) >> 2;
    let third = *MAP.get_unchecked(usize::from(third));

    let forth = byte & 0b_0000_0011;
    let forth = *MAP.get_unchecked(usize::from(forth));

    *output.get_unchecked_mut(0) = first;
    *output.get_unchecked_mut(1) = second;
    *output.get_unchecked_mut(2) = third;
    *output.get_unchecked_mut(3) = forth;
}

/// # Safety
///
/// The behavior is undefined when output is shorter than `4 * bytes.len()` bytes.
unsafe fn write_spi_array(output: &mut [u8], bytes: &[u8]) {
    debug_assert!(output.len() >= 4 * bytes.len());
    for i in 0..bytes.len() {
        let byte = *bytes.get_unchecked(i);
        let output = output.get_unchecked_mut(i * 4..(i + 1) * 4);
        write_spi(output, byte);
    }
}

/// # Safety
///
/// The behavior is undefined when output is shorter than 12 bytes.
unsafe fn write_spi_rgb(output: &mut [u8], red: u8, green: u8, blue: u8) {
    // Blue, Green, Red
    write_spi_array(output, &[blue, green, red])
}

/// # Safety
///
/// The behavior is undefined when output is shorter than 12 bytes.
fn get_bytes_for_led(output: &mut [u8], led_number: usize) -> Option<&mut [u8]> {
    // It takes 4 spi bytes to write a single byte, and we've got 3 bytes per
    // led, so ...
    let begin = led_number * 12;
    let end = (led_number + 1) * 12;
    output.get_mut(begin..end)
}

fn main() -> anyhow::Result<()> {
    let Options {
        red,
        green,
        blue,
        multiplier,
        interval,
        frequency,
        leds_count,
        moving_interval,
        wave_width,
    } = Options::from_args();

    if multiplier < 0. || multiplier > 1. {
        panic!("Multiplier must be 0..1");
    }

    let red = multiply(red, multiplier);
    let green = multiply(green, multiplier);
    let blue = multiply(blue, multiplier);

    let mut spi = Spidev::open("/dev/spidev0.0").expect("Open spidev0.0");
    let options = SpidevOptions::new()
        .bits_per_word(8)
        .max_speed_hz(frequency)
        .mode(SpiModeFlags::SPI_MODE_0)
        .build();
    spi.configure(&options).unwrap();

    let mut buf = [0u8; 600];

    // let color_buffer = get_bytes_for_led(&mut buf[..], 0).unwrap();
    // unsafe {
    //     write_spi_rgb(color_buffer, red, green, blue);
    // }
    // spi.write_all(&buf[..12]).unwrap();

    let begin = Instant::now();

    let mut first_led = 0;
    let mut last_update = Instant::now();
    while begin.elapsed() < interval {
        if last_update.elapsed() >= moving_interval {
            first_led = (first_led + 1) % leds_count;
            last_update = Instant::now();
        }

        for i in 0..leds_count {
            let color_buffer = get_bytes_for_led(&mut buf[..], i).unwrap();
            assert_eq!(color_buffer.len(), 12);
            unsafe { write_spi_rgb(color_buffer, 0, 0, 0) }
        }

        for i in 0..wave_width {
            let color_buffer = match get_bytes_for_led(&mut buf[..], first_led + i) {
                Some(buf) => buf,
                None => break,
            };
            let local_multiplier = (i as f64 * PI / wave_width as f64).sin().powi(6);
            let multiplier = multiplier * local_multiplier;

            let red = multiply(red, multiplier);
            let green = multiply(green, multiplier);
            let blue = multiply(blue, multiplier);

            // We are sure that color_buffer is large enough.
            unsafe {
                write_spi_rgb(color_buffer, red, green, blue);
            }
        }
        spi.write_all(&buf[..]).unwrap();
        
        std::thread::sleep(Duration::from_micros(300));
    }

    Ok(())
}

fn multiply(base: u8, scale: f64) -> u8 {
    let scale = if scale < 0. {
        0.
    } else if scale > 1. {
        1.
    } else {
        scale
    };
    (base as f64 * scale) as u8
}