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);
        pen.valid_settings();
    }
}

#[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;

            dst.valid_settings();
        }
    }
}

#[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(Copy, Clone, Default, PartialEq)]
struct RandomizedSetting {
    speed: Option<f64>,
    power: Option<f64>,
    frequency: Option<u32>,
    pulse_width: Option<PulseWidth>,
}

impl RandomizedSetting {
    fn apply(&self, pen: &mut Pen) {
        self.speed.map(|speed| *pen.speed = speed);
        self.power.map(|power| *pen.power = power);
        self.frequency.map(|freq| {
            *pen.frequency = freq;
            *pen.frequency_2 = freq.try_into().unwrap();
        });
        self.pulse_width.map(|pw| {
            let pw: u32 = pw.into();
            *pen.pulse_width = pw;
            *pen.pulse_width_2 = pw.try_into().unwrap();
        });
    }
}

#[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, usize)>,
}

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

        let mut generated: Vec<RandomizedSetting> = vec![];
        const MAX_ATTEMPTS: usize = 1000;

        for (index, pen) in pens
            .iter_mut()
            .skip(self.index)
            .take(self.count)
            .enumerate()
        {
            for attempt in 0..=MAX_ATTEMPTS {
                let mut setting: RandomizedSetting = RandomizedSetting::default();

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

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

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

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

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

                if !generated.contains(&setting) {
                    generated.push(setting);
                    setting.apply(pen);
                    if !pen.valid_settings() {
                        debug!("Retrying..");
                    } else {
                        break;
                    }
                } else {
                    debug!("Duplicate random setting");
                }

                if attempt == MAX_ATTEMPTS {
                    panic!(
                        "Exceeded maximum number of {} randommization attempts",
                        MAX_ATTEMPTS
                    );
                }
            }

            pen.valid_settings();
        }
    }
}

#[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;
    }
}