Minilase/src/config/pen.rs

use std::path::PathBuf;

use ezcad::{
    pen::Pen,
    types::{PulseWidth, Rgba},
};
use itertools::Itertools;
use log::debug;
use rand::{seq::SliceRandom, Rng};
use serde::{Deserialize, Serialize};
use strum::IntoEnumIterator;

const SPEED_MIN: f64 = 0.0;
const SPEED_MAX: f64 = 100000.0;
const POWER_MIN: f64 = 0.0;
const POWER_MAX: f64 = 100.0;
const FREQUENCY_MIN: u32 = 20_000;
const FREQUENCY_MAX: u32 = 4_000_000;

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "PascalCase")]
pub struct Patch {
    color: Option<(u8, u8, u8)>,
    enabled: Option<bool>,
    loop_count: Option<u32>,
    speed: Option<f64>,
    power: Option<f64>,
    frequency: Option<u32>,
    pulse_width: Option<PulseWidth>,
}

impl Patch {
    fn patch(&self, pen: &mut Pen) {
        self.color.map(|color| {
            debug!("Patching pen color to {:?}", color);
            *pen.color = color.into()
        });

        self.enabled.map(|enabled| {
            debug!("Patching pen enablement to {}", enabled);
            *pen.disabled = !enabled as u32;
        });

        self.loop_count.map(|loop_count| {
            debug!("Patching pen loop count to {}", loop_count);
            assert!(loop_count > 0, "Pen loop count must be greater than zero");
            *pen.loop_count = loop_count;
        });

        self.speed.map(|speed| {
            debug!("Patching pen speed to {}", speed);
            assert!(
                speed > SPEED_MIN && speed <= SPEED_MAX,
                "Pen speed must be between {} and {}",
                SPEED_MIN,
                SPEED_MAX
            );
            *pen.speed = speed;
        });

        self.power.map(|power| {
            debug!("Patching pen power to {}", power);
            assert!(
                power > POWER_MIN && power <= POWER_MAX,
                "Pen power must be between {} and {}",
                POWER_MIN,
                POWER_MAX
            );
            *pen.power = power;
        });

        self.frequency.map(|frequency| {
            debug!("Patching pen frequency to {}", frequency);
            assert!(
                frequency >= FREQUENCY_MIN && frequency <= FREQUENCY_MAX,
                "Pen frequency must be between {} and {}",
                FREQUENCY_MIN,
                FREQUENCY_MAX
            );
            *pen.frequency = frequency;
            *pen.frequency_2 = frequency.try_into().unwrap();
        });

        self.pulse_width.map(|width| {
            let width: u32 = width.into();
            debug!("Patching pen pulse width to {}ns", width);
            *pen.pulse_width = width.into();
            *pen.pulse_width_2 = width.try_into().unwrap();
        });

        // Always enable custom settings for pen
        *pen.use_default = 0;
    }
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "PascalCase")]
pub struct PatchPen {
    pen: usize,
    #[serde(flatten)]
    patch: Patch,
}

impl PatchPen {
    pub fn patch(&self, pens: &mut Vec<Pen>) {
        debug!("Patching pen #{}", self.pen);
        let pen: &mut Pen = pens.get_mut(self.pen).expect("Invalid pen index");
        self.patch.patch(pen);
    }
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "PascalCase")]
pub struct ClonePen {
    from: usize,
    to: usize,
    inclusive: Option<bool>,
    patch: Option<Patch>,
}

impl ClonePen {
    pub fn clone(&self, pens: &mut Vec<Pen>) {
        debug!(
            "Cloning pen #{} to #{}{}",
            self.from,
            self.to,
            match self.inclusive {
                Some(true) => format!(" (inclusive)"),
                _ => format!(""),
            }
        );

        // Clone pen
        let src: Pen = pens.get(self.from).expect("Invalid pen index").clone();

        match self.inclusive {
            Some(true) => {
                assert!(
                    self.to > self.from,
                    "Target pen(s) must be greater than source pen"
                );

                // Clone pen (randomize color)
                for idx in (self.from..=self.to).skip(1) {
                    let dst: &mut Pen = pens.get_mut(idx).expect("Invalid pen index");
                    *dst = src.clone();
                    *dst.color = Rgba::random().into();

                    // Patch pen if needed
                    self.patch.as_ref().map(|patch| {
                        debug!("Patching pen #{}", idx);
                        patch.patch(dst);
                    });
                }
            }
            _ => {
                let dst: &mut Pen = pens.get_mut(self.to).expect("Invalid pen index");
                *dst = src;

                // Patch pen if needed
                self.patch.as_ref().map(|patch| {
                    debug!("Patching pen #{}", self.to);
                    patch.patch(dst);
                });
            }
        }
    }
}

#[derive(Debug, Serialize, Deserialize)]
pub enum PatternField {
    Loops(i32),
    Speed(f64),
    Power(f64),
    Frequency(i32),
    PulseWidth(u32),
}

impl PatternField {
    pub fn pattern(&self, pens: &mut dyn Iterator<Item = (usize, &mut Pen)>) {
        // Obtain settings from source (first) pen
        let (src_idx, src) = pens.next().expect("Pattern must involve at least one pen");

        let mut setting: PatternField = match self {
            PatternField::Loops(_) => {
                debug!(
                    "Initial loop count from pen #{} is {}",
                    src_idx, *src.loop_count
                );
                PatternField::Loops((*src.loop_count).try_into().unwrap())
            }
            PatternField::Speed(_) => {
                debug!("Initial speed from pen #{} is {}", src_idx, *src.speed);
                PatternField::Speed(*src.speed)
            }
            PatternField::Power(_) => {
                debug!("Initial power from pen #{} is {}", src_idx, *src.power);
                PatternField::Power(*src.power)
            }
            PatternField::Frequency(_) => {
                debug!(
                    "Initial frequency from pen #{} is {}",
                    src_idx, *src.frequency
                );
                PatternField::Frequency((*src.frequency).try_into().unwrap())
            }
            PatternField::PulseWidth(_) => {
                debug!(
                    "Initial pulse width from pen #{} is {}ns",
                    src_idx, *src.pulse_width
                );
                PatternField::PulseWidth(*src.pulse_width)
            }
        };

        for (idx, dst) in pens {
            // Calculate new setting
            setting = match (setting, self) {
                (PatternField::Loops(prev), PatternField::Loops(incr)) => {
                    let value: i32 = prev + incr;
                    debug!("Patching loop count for pen #{} to {}", idx, value);
                    assert!(value > 0, "Pen loop count must be greater than zero");
                    PatternField::Loops(value)
                }
                (PatternField::Speed(prev), PatternField::Speed(incr)) => {
                    let value: f64 = prev + incr;
                    debug!("Patching speed for pen #{} to {}", idx, value);
                    assert!(
                        value > SPEED_MIN && value <= SPEED_MAX,
                        "Pen speed must be between {} and {}",
                        SPEED_MIN,
                        SPEED_MAX
                    );
                    PatternField::Speed(value)
                }
                (PatternField::Power(prev), PatternField::Power(incr)) => {
                    let value: f64 = prev + incr;
                    debug!("Patching power for pen #{} to {}", idx, value);
                    assert!(
                        value > POWER_MIN && value <= POWER_MAX,
                        "Pen power must be between {} and {}",
                        POWER_MIN,
                        POWER_MAX
                    );
                    PatternField::Power(value)
                }
                (PatternField::Frequency(prev), PatternField::Frequency(incr)) => {
                    let value: i32 = prev + incr;
                    debug!("Patching frequency for pen #{} to {}", idx, value);
                    assert!(
                        value >= FREQUENCY_MIN.try_into().unwrap()
                            && value <= FREQUENCY_MAX.try_into().unwrap(),
                        "Pen frequency must be between {} and {}",
                        FREQUENCY_MIN,
                        FREQUENCY_MAX
                    );
                    PatternField::Frequency(value)
                }
                (PatternField::PulseWidth(prev), PatternField::PulseWidth(incr)) => {
                    let mut pw = PulseWidth::iter();
                    let _ = pw
                        .find(|x| u32::from(*x) == prev)
                        .expect("Unknown pulse width");

                    let mut pw = pw.skip((*incr - 1).try_into().unwrap());
                    let next: u32 = pw.next().expect("Pulse width out of bounds").into();
                    debug!("Patching pulse width for pen #{} to {}ns", idx, next);
                    PatternField::PulseWidth(next)
                }
                _ => unreachable!(),
            };

            // Patch updated value
            match setting {
                PatternField::Loops(x) => *dst.loop_count = x.try_into().unwrap(),
                PatternField::Speed(x) => *dst.speed = x,
                PatternField::Power(x) => *dst.power = x,
                PatternField::Frequency(x) => {
                    *dst.frequency = x.try_into().unwrap();
                    *dst.frequency_2 = x.try_into().unwrap();
                }
                PatternField::PulseWidth(x) => {
                    *dst.pulse_width = x;
                    *dst.pulse_width_2 = x.try_into().unwrap();
                }
            }

            // Randomize pen color
            *dst.color = Rgba::random().into();

            // Always enable custom settings for pen
            *dst.use_default = 0;
        }
    }
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "PascalCase")]
pub struct PatternPen {
    index: usize,
    count: usize,
    field: PatternField,
}

impl PatternPen {
    pub fn pattern(&self, pens: &mut Vec<Pen>) {
        debug!(
            "Patterning from pen #{} to #{}",
            self.index,
            self.index + self.count - 1
        );
        self.field.pattern(
            &mut pens
                .iter_mut()
                .enumerate()
                .skip(self.index)
                .take(self.count),
        )
    }
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "PascalCase")]
pub struct RandomizePen {
    index: usize,
    count: usize,
    speed: Option<(f64, f64, f64)>,
    power: Option<(f64, f64, f64)>,
    frequency: Option<(u32, u32, u32)>,
    pulse_width: Option<(PulseWidth, PulseWidth)>,
}

impl RandomizePen {
    pub fn random(&self, pens: &mut Vec<Pen>) {
        debug!(
            "Randomizing from pen #{} to #{}",
            self.index,
            self.index + self.count - 1
        );

        for (index, pen) in pens
            .iter_mut()
            .skip(self.index)
            .take(self.count)
            .enumerate()
        {
            if let Some((min, max, incr)) = self.speed {
                let offset: usize = rand::thread_rng().gen_range(0..=((max - min) / incr) as usize);
                let value: f64 = min + incr * offset as f64;
                debug!("Randomizing speed for pen #{} to {}", index, value);
                *pen.speed = value;
            }

            if let Some((min, max, incr)) = self.power {
                let offset: usize = rand::thread_rng().gen_range(0..=((max - min) / incr) as usize);
                let value: f64 = min + incr * offset as f64;
                debug!("Randomizing power for pen #{} to {}", index, value);
                *pen.power = value;
            }

            if let Some((min, max, incr)) = self.frequency {
                let offset: usize = rand::thread_rng().gen_range(0..=((max - min) / incr) as usize);
                let value: u32 = min + incr * offset as u32;
                debug!("Randomizing frequency for pen #{} to {}", index, value);
                *pen.frequency = value;
                *pen.frequency_2 = value.try_into().unwrap();
            }

            if let Some((min, max)) = self.pulse_width {
                let mut pw = PulseWidth::iter();
                let mut v: Vec<PulseWidth> = vec![pw.find(|x| *x == min).unwrap()];
                v.extend(pw.take_while_inclusive(|x| *x != max).collect_vec());

                let width: &PulseWidth = v
                    .choose_multiple(&mut rand::thread_rng(), 1)
                    .next()
                    .unwrap();
                let value: u32 = (*width).into();
                debug!("Randomizing pulse width for pen #{} to {}", index, value);
                *pen.pulse_width = value;
                *pen.pulse_width_2 = value.try_into().unwrap();
            }
        }
    }
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "PascalCase")]
pub struct ImportExportPen {
    index: usize,
    path: PathBuf,
}

impl ImportExportPen {
    pub fn export(&self, pens: &mut Vec<Pen>) {
        debug!(
            "Exporting pen #{} to '{}'",
            self.index,
            self.path.to_string_lossy()
        );

        let pen = pens.get(self.index).expect("Invalid pen index");

        pen.write_to_file(&self.path);
    }

    pub fn import(&self, pens: &mut Vec<Pen>) {
        debug!(
            "Importing pen #{} from '{}'",
            self.index,
            self.path.to_string_lossy()
        );

        let pen: Pen = Pen::read_from_file(&self.path);

        let dst: &mut Pen = pens.get_mut(self.index).expect("Invalid pen index");
        *dst = pen;
    }
}