use std::{convert, ops};
use num_traits::identities;
use crate::gec::Gec;
use crate::scalar::{Scalar, ScalarFloat};
#[derive(Clone, Debug)]
pub struct Mat<S: Scalar, const NROWS: usize, const NCOLS: usize>([[S; NCOLS]; NROWS]);
impl<S: Scalar, const NROWS: usize, const NCOLS: usize> Mat<S, NROWS, NCOLS> {
pub fn zeros() -> Self {
Self::from([[identities::zero(); NCOLS]; NROWS])
}
pub fn ones() -> Self {
Self::from([[identities::one(); NCOLS]; NROWS])
}
}
impl<S: Scalar, const DIM: usize> Mat<S, DIM, DIM> {
pub fn identity() -> Self {
let mut res = Self::zeros();
for i in 0..DIM {
res[[i,i]] = identities::one();
}
res
}
}
impl<S: Scalar, const NROWS: usize, const NCOLS: usize> convert::From<[[S; NCOLS]; NROWS]>
for Mat<S, NROWS, NCOLS>
{
fn from(value: [[S; NCOLS]; NROWS]) -> Self {
Self(value)
}
}
impl<S: Scalar, const NROWS: usize, const NCOLS: usize> ops::Index<[usize; 2]>
for Mat<S, NROWS, NCOLS>
{
type Output = S;
fn index(&self, [row, col]: [usize; 2]) -> &Self::Output {
&self.0[row][col]
}
}
impl<S: Scalar, const NROWS: usize, const NCOLS: usize> ops::IndexMut<[usize; 2]>
for Mat<S, NROWS, NCOLS>
{
fn index_mut(&mut self, [row, col]: [usize; 2]) -> &mut Self::Output {
&mut self.0[row][col]
}
}
impl<S: Scalar, const NROWS: usize, const NCOLS: usize> ops::Mul<Gec<S, NCOLS>>
for &Mat<S, NROWS, NCOLS>
{
type Output = Gec<S, NROWS>;
fn mul(self, rhs: Gec<S, NCOLS>) -> Self::Output {
let mut res = Gec::zero();
for row in 0..NROWS {
for col in 0..NCOLS {
res[row] = res[row] + self[[row, col]] * rhs[col];
}
}
res
}
}
impl<S: Scalar, const NROWS: usize, const NCOLS: usize> ops::Mul<Gec<S, NCOLS>>
for Mat<S, NROWS, NCOLS>
{
type Output = Gec<S, NROWS>;
fn mul(self, rhs: Gec<S, NCOLS>) -> Self::Output {
&self * rhs
}
}
impl<S: Scalar, const DIM1: usize, const DIM2: usize, const DIM3: usize>
ops::Mul<&Mat<S, DIM2, DIM3>> for &Mat<S, DIM1, DIM2>
{
type Output = Mat<S,DIM1,DIM3>;
fn mul(self, rhs: &Mat<S, DIM2, DIM3>) -> Self::Output {
let mut res = Self::Output::from([[identities::zero(); DIM3]; DIM1]);
for i in 0..DIM1 {
for j in 0..DIM2 {
for k in 0..DIM3 {
res[[i,k]] = res[[i,k]] + self[[i,j]]*rhs[[j,k]];
}
}
}
res
}
}
impl<S: Scalar, const DIM1: usize, const DIM2: usize, const DIM3: usize>
ops::Mul<&Mat<S, DIM2, DIM3>> for Mat<S, DIM1, DIM2>
{
type Output = Mat<S,DIM1,DIM3>;
fn mul(self, rhs: &Mat<S, DIM2, DIM3>) -> Self::Output {
&self * rhs
}
}
impl<S: Scalar, const DIM1: usize, const DIM2: usize, const DIM3: usize>
ops::Mul<Mat<S, DIM2, DIM3>> for &Mat<S, DIM1, DIM2>
{
type Output = Mat<S,DIM1,DIM3>;
fn mul(self, rhs: Mat<S, DIM2, DIM3>) -> Self::Output {
self * &rhs
}
}
impl<S: Scalar, const DIM1: usize, const DIM2: usize, const DIM3: usize>
ops::Mul<Mat<S, DIM2, DIM3>> for Mat<S, DIM1, DIM2>
{
type Output = Mat<S,DIM1,DIM3>;
fn mul(self, rhs: Mat<S, DIM2, DIM3>) -> Self::Output {
&self * &rhs
}
}
impl<S: ScalarFloat> Mat<S,3,3> {
pub fn rotate_xyz(rot_x: S, rot_y: S, rot_z: S) -> Self {
let xcos = rot_x.cos();
let ycos = rot_y.cos();
let zcos = rot_z.cos();
let xsin = rot_x.sin();
let ysin = rot_y.sin();
let zsin = rot_z.sin();
let xsin_ysin = xsin * ysin;
let xcos_ysin = xcos * ysin;
Self::from([
[ycos * zcos, xsin_ysin * zcos - xcos * zsin, xcos_ysin * zcos + xsin * zsin],
[ycos * zsin, xsin_ysin * zsin - xcos * zcos, xcos_ysin * zsin - xsin * zcos],
[ -ysin, xsin * ycos, xcos * ycos],
])
}
}
impl<S: ScalarFloat> Mat<S,4,4> {
pub fn rotate_xyz_translate(rot_x: S, rot_y: S, rot_z: S, translate: Gec<S,3>) -> Self {
let xcos = rot_x.cos();
let ycos = rot_y.cos();
let zcos = rot_z.cos();
let xsin = rot_x.sin();
let ysin = rot_y.sin();
let zsin = rot_z.sin();
let xsin_ysin = xsin * ysin;
let xcos_ysin = xcos * ysin;
let zero = identities::zero();
let one = identities::one();
Self::from([
[ycos * zcos, xsin_ysin * zcos - xcos * zsin, xcos_ysin * zcos + xsin * zsin, translate[0]],
[ycos * zsin, xsin_ysin * zsin - xcos * zcos, xcos_ysin * zsin - xsin * zcos, translate[1]],
[ -ysin, xsin * ycos, xcos * ycos, translate[2]],
[ zero, zero, zero, one],
])
}
}