#![no_std]

use core::{
    fmt::Display,
    ops::{Add, Sub},
};
use num_traits::{SaturatingAdd, SaturatingSub};

pub trait Scaled {
    #[must_use]
    fn saturating_mul_unscaled(self, rhs: Self) -> Self;
    #[must_use]
    fn rescale(self) -> Self;
}

#[derive(Clone, Copy, Default, Eq, PartialEq)]
pub struct Fixed<const F: usize> {
    v: i32,
}

impl<const F: usize> Fixed<F> {
    #[must_use]
    pub fn from_int(v: i32) -> Self {
        Self {
            v: v.saturating_mul(1 << F),
        }
    }

    #[must_use]
    pub fn from_raw(v: i32) -> Self {
        Self { v }
    }

    #[must_use]
    pub fn round(self) -> i32 {
        self.rescale().v
    }

    #[must_use]
    pub fn to_raw(self) -> i32 {
        self.v
    }

    const OFFSET: i32 = 1 << (F - 1);
}

impl<const F: usize> core::fmt::Debug for Fixed<F> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.write_fmt(format_args!("{:?}", f64::from(self.v) / f64::from(1 << F)))
    }
}

impl<const F: usize> Add<Self> for Fixed<F> {
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        Self { v: self.v + rhs.v }
    }
}

impl<const F: usize> SaturatingAdd for Fixed<F> {
    fn saturating_add(&self, rhs: &Self) -> Self::Output {
        Self {
            v: self.v.saturating_add(rhs.v),
        }
    }
}

impl<const F: usize> Sub<Self> for Fixed<F> {
    type Output = Self;

    fn sub(self, rhs: Self) -> Self::Output {
        Self { v: self.v - rhs.v }
    }
}

impl<const F: usize> SaturatingSub for Fixed<F> {
    fn saturating_sub(&self, rhs: &Self) -> Self::Output {
        Self {
            v: self.v.saturating_sub(rhs.v),
        }
    }
}

impl<const F: usize> Scaled for Fixed<F> {
    fn saturating_mul_unscaled(self, rhs: Self) -> Self {
        Self {
            v: self.v.saturating_mul(rhs.v),
        }
    }

    fn rescale(self) -> Self {
        if self.v >= 0 {
            Self {
                v: self.v.saturating_add(Self::OFFSET) >> F,
            }
        } else {
            Self {
                v: self.v.saturating_add(Self::OFFSET - 1) >> F,
            }
        }
    }
}

impl<const F: usize> Display for Fixed<F> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.write_fmt(format_args!("{}", f64::from(self.v) / f64::from(1 << F)))
    }
}

#[derive(Default)]
pub struct Pid<T> {
    pub k_p: T,
    pub k_i: T,
    pub k_d: T,
    pub prev_error: T,
    pub integrated_error: T,
}

impl<T: Default> Pid<T> {
    pub fn new(k_p: T, k_i: T, k_d: T) -> Self {
        Self {
            k_p,
            k_i,
            k_d,
            prev_error: Default::default(),
            integrated_error: Default::default(),
        }
    }

    pub fn reset_d(&mut self) {
        self.prev_error = Default::default();
    }

    pub fn reset_i(&mut self) {
        self.integrated_error = Default::default();
    }

    pub fn reset(&mut self) {
        self.reset_d();
        self.reset_i();
    }
}

impl<T: SaturatingAdd + SaturatingSub + Scaled + Copy> Pid<T> {
    pub fn command_from_error(&mut self, error: T) -> T {
        let deriv_error = error.saturating_sub(&self.prev_error);
        self.prev_error = error;
        self.integrated_error = self.integrated_error.saturating_add(&error);
        let p = error.saturating_mul_unscaled(self.k_p);
        let i = self.integrated_error.saturating_mul_unscaled(self.k_i);
        let d = deriv_error.saturating_mul_unscaled(self.k_d);
        (p + i + d).rescale()
    }

    pub fn command(&mut self, target: T, measured: T) -> T {
        self.command_from_error(target.saturating_sub(&measured))
    }
}