//! # Matrix of transfer functions
//!
//! `TfMatrix` allow the definition of a matrix of transfer functions. The
//! numerators are stored in a matrix, while the denominator is stored once,
//! since it is equal for every transfer function.
//! * evaluation of a vector of inputs
//! * conversion from a generic state space representation

use std::{
    fmt,
    fmt::{Debug, Display},
    ops::{Add, AddAssign, Div, Mul, Neg, Sub},
};

use ndarray::{Array2, Axis, Zip};

use polynomen::Poly;

use crate::{
    enums::Time,
    linear_system::{self, SsGen},
    polynomial_matrix::{MatrixOfPoly, PolyMatrix},
    wrappers::{PWrapper, PP},
    Abs, Complex, Inv, NumCast, One, Polynomial, Zero,
};

/// Matrix of transfer functions
#[derive(Debug)]
pub struct TfMatrix<T> {
    /// Polynomial matrix of the numerators
    num: MatrixOfPoly<T>,
    /// Common polynomial denominator
    den: PP<T>,
}

/// Implementation of transfer function matrix
impl<T> TfMatrix<T>
where
    T: Clone + PartialEq + Zero,
{
    /// Create a new transfer function matrix
    ///
    /// # Arguments
    ///
    /// * `num` - Polynomial matrix
    /// * `den` - Characteristic polynomial of the system
    #[allow(dead_code)]
    fn new<R>(num: MatrixOfPoly<T>, den: R) -> Self
    where
        R: AsRef<[T]>,
    {
        Self {
            num,
            den: PP::new_from_coeffs(den.as_ref()),
        }
    }

    /// Create a new transfer function matrix
    ///
    /// # Arguments
    ///
    /// * `num` - Polynomial matrix
    /// * `den` - Characteristic polynomial of the system
    fn new_from_pp(num: MatrixOfPoly<T>, den: PP<T>) -> Self {
        Self { num, den }
    }
}

impl<T> TfMatrix<T>
where
    T: Clone + PartialEq,
{
    /// Retrieve the characteristic polynomial of the system.
    #[must_use]
    pub fn den(&self) -> impl Polynomial<T> {
        self.den.to_unwrapped_vec()
    }
}

impl<T> TfMatrix<T>
where
    T: Abs
        + Add<Output = T>
        + Clone
        + Div<Output = T>
        + Inv<Output = T>
        + Mul<Output = T>
        + Neg<Output = T>
        + PartialOrd
        + Sub<Output = T>
        + Zero,
{
    /// Evaluate the matrix transfers function.
    ///
    /// # Arguments
    ///
    /// * `s` - Array at which the Matrix of transfer functions is evaluated.
    #[must_use]
    pub fn eval(&self, s: &[Complex<T>]) -> Vec<Complex<T>> {
        //
        // ┌  ┐ ┌┌         ┐ ┌     ┐┐┌  ┐
        // │y1│=││1/pc 1/pc│*│n1 n2│││s1│
        // │y2│ ││1/pc 1/pc│ │n3 n4│││s2│
        // └  ┘ └└         ┘ └     ┘┘└  ┘
        // `*` is the element by element multiplication
        // ┌     ┐ ┌┌         ┐ ┌     ┐┐ ┌┌     ┐ ┌     ┐┐
        // │y1+y2│=││1/pc 1/pc│.│s1 s2││*││n1 n2│.│s1 s2││
        // │y3+y4│ ││1/pc 1/pc│ │s1 s2││ ││n3 n4│ │s1 s2││
        // └     ┘ └└         ┘ └     ┘┘ └└     ┘ └     ┘┘
        // `.` means 'evaluated at'

        // Create a matrix to contain the result of the evaluation.
        let mut res = Array2::from_elem(
            self.num.matrix().dim(),
            Complex::<T>::new(T::zero(), T::zero()),
        );

        // Zip the result and the numerator matrix row by row.
        Zip::from(res.rows_mut())
            .and(self.num.matrix().rows())
            .for_each(|mut res_row, matrix_row| {
                // Zip the row of the result matrix.
                Zip::from(&mut res_row)
                    .and(s) // The vector of the input.
                    .and(matrix_row) // The row of the numerator matrix.
                    .for_each(|r, s, n| {
                        let div =
                            n.0.eval(&PWrapper(s.clone())) / self.den.0.eval(&PWrapper(s.clone()));
                        *r = div.0;
                    });
            });

        res.fold_axis(Axis(1), Complex::<T>::zero(), |x, y| x + y)
            .to_vec()
    }
}

impl<T, U> From<SsGen<T, U>> for TfMatrix<T>
where
    T: Add<Output = T>
        + AddAssign
        + Clone
        + Inv<Output = T>
        + Mul<Output = T>
        + Neg<Output = T>
        + NumCast
        + One
        + PartialEq
        + Zero,
    U: Time,
{
    /// Convert a state-space representation into a matrix of transfer functions
    ///
    /// # Arguments
    ///
    /// `ss` - state space linear system
    fn from(ss: SsGen<T, U>) -> Self {
        let (pc, a_inv) = linear_system::leverrier(&ss.a).unwrap();
        let g = a_inv.left_mul(&ss.c).right_mul(&ss.b);
        let rest = PolyMatrix::multiply(&pc, &ss.d);
        let tf = g + rest;
        Self::new_from_pp(MatrixOfPoly::from(tf), pc)
    }
}

impl<T> TfMatrix<T>
where
    T: Clone,
{
    /// Get a copy of a polynomial at index `i` of the numerator matrix
    ///
    /// # Panics
    ///
    /// Panics for out of bounds access.
    #[must_use]
    pub fn get_unchecked(&self, i: [usize; 2]) -> impl Polynomial<T> {
        self.num.get_unchecked(i)
    }
}

impl<T> TfMatrix<T>
where
    T: Clone + PartialEq + Zero,
{
    /// Set a polynomial at index `i` of the numerator matrix
    ///
    /// # Panics
    ///
    /// Panics for out of bounds access.
    pub fn set_unchecked(&mut self, i: [usize; 2], value: &impl Polynomial<T>) {
        self.num.set_unchecked(i, value);
    }
}

/// Implementation of transfer function matrix printing
impl<T> Display for TfMatrix<T>
where
    T: Display + PartialOrd + Zero,
    Poly<T>: Display,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let s_num = self.num.to_string();
        let s_den = self.den.to_string();

        let length = s_den.len();
        let dash = "\u{2500}".repeat(length);

        write!(f, "{}\n{}\n{}", s_num, dash, s_den)
    }
}

#[cfg(test)]
mod tests {
    use crate::{Ss, Ssd};

    use super::*;

    #[test]
    #[allow(clippy::float_cmp)]
    fn tf_matrix_new() {
        let sys = Ssd::new_from_slice(
            2,
            2,
            2,
            &[-2.0_f32, 0., 0., -1.],
            &[0., 1., 1., 2.],
            &[1., 2., 3., 4.],
            &[1., 0., 0., 1.],
        );
        let tfm = TfMatrix::from(sys);
        assert_eq!(tfm.get_unchecked([0, 0]).as_slice(), [6., 5., 1.]);
        assert_eq!(tfm.get_unchecked([0, 1]).as_slice(), [9., 5.]);
        assert_eq!(tfm.get_unchecked([1, 0]).as_slice(), [8., 4.]);
        assert_eq!(tfm.get_unchecked([1, 1]).as_slice(), [21., 14., 1.]);
        assert_eq!(tfm.den().as_slice(), [2., 3., 1.]);
    }

    #[test]
    fn tf_matrix_eval() {
        let sys = Ss::new_from_slice(
            2,
            2,
            2,
            &[-2., 0., 0., -1.],
            &[0., 1., 1., 2.],
            &[1., 2., 3., 4.],
            &[1., 0., 0., 1.],
        );
        let tfm = TfMatrix::from(sys);
        let i = Complex::new(0., 1.);
        let res = tfm.eval(&[i, i]);
        assert_relative_eq!(res[0].re, 4.4, max_relative = 1e-15);
        assert_relative_eq!(res[0].im, -3.2, max_relative = 1e-15);
        assert_relative_eq!(res[1].re, 8.2, max_relative = 1e-15);
        assert_relative_eq!(res[1].im, -6.6, max_relative = 1e-15);
    }

    #[test]
    #[allow(clippy::float_cmp)]
    fn tf_matrix_index_mut() {
        let sys = Ss::new_from_slice(
            2,
            2,
            2,
            &[-2., 0., 0., -1.],
            &[0., 1., 1., 2.],
            &[1., 2., 3., 4.],
            &[1., 0., 0., 1.],
        );
        let mut tfm = TfMatrix::from(sys);
        tfm.set_unchecked([0, 0], &vec![1., 2., 3.]);
        assert_eq!([1., 2., 3.], tfm.get_unchecked([0, 0]).as_slice());
    }

    #[test]
    fn tf_matrix_print() {
        let sys = Ss::new_from_slice(
            2,
            2,
            2,
            &[-2., 0., 0., -1.],
            &[0., 1., 1., 2.],
            &[1., 2., 3., 4.],
            &[1., 0., 0., 1.],
        );
        let tfm = TfMatrix::from(sys);
        assert_eq!(
            "[[6 +5s +1s^2, 9 +5s],\n [8 +4s, 21 +14s +1s^2]]\n\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\u{2500}\n2 +3s +1s^2",
            format!("{}", tfm)
        );
    }
}