use ndarray::arr1;
use num_complex::Complex64;

fn main() {
    let complex_array = arr1(&[Complex64::from(1.0), Complex64::from(2.0)]);
    let complex_scalar = Complex64::from(1.0);
    let f_scalar = 1.0;
    let f_array = arr1(&[1.0, 2.0]);
    println!("{}", &f_array * complex_scalar * f_scalar);
    // println!("{}", f_scalar * &f_array * complex_scalar); // not working: type mismatch
    println!("{}", &(f_scalar * &f_array) * complex_scalar); // fixes above
    println!("{}", complex_scalar * &f_array * f_scalar); // this works as well

    println!("{}", &complex_array * &f_array); // this works
    // println!("{}", &f_array * &complex_array); // this will not
}

impl<A, S, D, B> Mul<B> for ArrayBase<S, D>
where
    A: Clone + Mul<B, Output = A>,
    S: DataOwned<Elem = A> + DataMut,
    D: Dimension,
    B: ScalarOperand,

impl<'a, A, B, S, S2, D, E> Mul<&'a ArrayBase<S2, E>> for &'a ArrayBase<S, D>
where
    A: Clone + Mul<B, Output = A>,
    B: Clone,
    S: Data<Elem = A>,
    S2: Data<Elem = B>,
    D: Dimension,
    E: Dimension,

use ndarray::prelude::*;
use num_complex::Complex64;

fn main() {
    let complex_array = arr1(&[Complex64::from(1.0), Complex64::from(2.0)]);
    let f_array = arr1(&[1.0, 2.0]);
    let mut result = Array1::zeros(2);
    azip!((res in &mut result, &f in &f_array, &c in &complex_array) *res = f * c);
}

use ndarray::prelude::*;
use ndarray::{Data, DataMut, DataOwned, ScalarOperand};
use std::ops::Add;

#[derive(Clone, Copy, Debug, PartialEq)]
pub struct MyType(f32);

impl Add<f32> for MyType {
    type Output = MyType;

    fn add(self, rhs: f32) -> MyType {
        MyType(self.0 + rhs)
    }
}

impl Add<MyType> for f32 {
    type Output = MyType;

    fn add(self, rhs: MyType) -> MyType {
        MyType(self + rhs.0)
    }
}

impl Add<MyType> for MyType {
    type Output = MyType;

    fn add(self, rhs: MyType) -> MyType {
        MyType(self.0 + rhs.0)
    }
}

impl ScalarOperand for MyType {}

macro_rules! impl_scalar_lhs_op {
    ($scalar:ty, $operator:tt, $trt:ident, $mth:ident, $doc:expr) => (
        /// Perform elementwise
        #[doc=$doc]
        /// between the scalar `self` and array `rhs`,
        /// and return the result (based on `self`).
        impl<A, S, D> $trt<ArrayBase<S, D>> for $scalar
        where
            $scalar: Clone + $trt<A, Output=A>,
            A: Clone,
            S: DataOwned<Elem=A> + DataMut,
            D: Dimension,
        {
            type Output = ArrayBase<S, D>;
            fn $mth(self, mut rhs: ArrayBase<S, D>) -> ArrayBase<S, D> {
                rhs.map_inplace(move |elt| *elt = self.clone() $operator elt.clone());
                rhs
            }
        }

        /// Perform elementwise
        #[doc=$doc]
        /// between the scalar `self` and array `rhs`,
        /// and return the result as a new `Array`.
        impl<'a, A, S, D, B> $trt<&'a ArrayBase<S, D>> for $scalar
        where
            $scalar: Clone + $trt<A, Output=B>,
            A: Clone,
            S: Data<Elem=A>,
            D: Dimension,
        {
            type Output = Array<B, D>;
            fn $mth(self, rhs: &ArrayBase<S, D>) -> Array<B, D> {
                rhs.map(move |elt| self.clone() $operator elt.clone())
            }
        }
    );
}

impl_scalar_lhs_op!(MyType, +, Add, add, "addition");

fn main() {
    let f = 1.;
    let my = MyType(2.);
    let arr_f = array![10., 20., 30.];
    let arr_my = array![MyType(1.), MyType(2.), MyType(3.)];

    // f32 and MyType
    assert_eq!(f + my, MyType(3.));
    assert_eq!(my + f, MyType(3.));

    // f32 and array of MyType
    assert_eq!(f + arr_my.clone(), array![MyType(2.), MyType(3.), MyType(4.)]);
    assert_eq!(f + &arr_my, array![MyType(2.), MyType(3.), MyType(4.)]);
    assert_eq!(arr_my.clone() + f, array![MyType(2.), MyType(3.), MyType(4.)]);
    assert_eq!(&arr_my + f, array![MyType(2.), MyType(3.), MyType(4.)]);

    // MyType and array of f32
    // assert_eq!(my + arr_f.clone(), array![MyType(12.), MyType(22.), MyType(32.)]); // doesn't work
    assert_eq!(my + &arr_f, array![MyType(12.), MyType(22.), MyType(32.)]);
    // assert_eq!(arr_f.clone() + my, array![MyType(12.), MyType(22.), MyType(32.)]); // doesn't work
    assert_eq!(&arr_f + my, array![MyType(12.), MyType(22.), MyType(32.)]);

    // MyType and array of MyType
    assert_eq!(my + arr_my.clone(), array![MyType(3.), MyType(4.), MyType(5.)]);
    assert_eq!(my + &arr_my, array![MyType(3.), MyType(4.), MyType(5.)]);
    assert_eq!(arr_my.clone() + my, array![MyType(3.), MyType(4.), MyType(5.)]);
    assert_eq!(&arr_my + my, array![MyType(3.), MyType(4.), MyType(5.)]);

    // array of f32 and array of MyType
    // assert_eq!(arr_f.clone() + arr_my.clone(), array![MyType(11.), MyType(22.), MyType(33.)]); // doesn't work
    // assert_eq!(arr_f.clone() + &arr_my, array![MyType(11.), MyType(22.), MyType(33.)]); // doesn't work
    // assert_eq!(&arr_f + &arr_my, array![MyType(11.), MyType(22.), MyType(33.)]); // doesn't work
    assert_eq!(arr_my.clone() + arr_f.clone(), array![MyType(11.), MyType(22.), MyType(33.)]);
    assert_eq!(arr_my.clone() + &arr_f, array![MyType(11.), MyType(22.), MyType(33.)]);
    assert_eq!(&arr_my + &arr_f, array![MyType(11.), MyType(22.), MyType(33.)]);

    // array of MyType and array of MyType
    assert_eq!(arr_my.clone() + arr_my.clone(), array![MyType(2.), MyType(4.), MyType(6.)]);
    assert_eq!(arr_my.clone() + &arr_my, array![MyType(2.), MyType(4.), MyType(6.)]);
    assert_eq!(&arr_my + &arr_my, array![MyType(2.), MyType(4.), MyType(6.)]);
}