//! Functions for computational electromagnetics.

#![feature(step_by)]
#![feature(inclusive_range_syntax)]

extern crate rulinalg;

pub mod topo;

use rulinalg::matrix::{BaseMatrix, BaseMatrixMut, Matrix, DiagOffset};
use rulinalg::vector::Vector;
use rulinalg::error::Error;

/// Permittivity of free space.
pub static EPS_FREE_SPACE: &'static f64 = &8.854187817620e-12;

/// Computes the potential over a grid represented by a field matrix using a
/// finite difference approximation to Poisson's equation. Takes as input a
/// matrix representing the physical grid you're modeling, with known starting
/// values (or zero).
///
/// Returns a `Matrix` of the same dimensions with the potentials computed.
pub fn finite_diff_rect(pot_mat: &Matrix<f64>) -> Result<Matrix<f64>, Error> {
  
  /* algorithm
   *
   * 1. compute node counts
   * 2. create boundary conditions vector
   * 3. create node coefficient matrix
   * 4. solve
   * 5. plots (possibly out of rust)
   * 5a. coefficient node matrix
   * 5b. solved potential matrix
   */
  
  // 1. compute node counts
  let row_count = pot_mat.rows();
  let col_count = pot_mat.cols();
  let node_count = row_count * col_count;
  
  // 3. create boundary conditions vector from potential matrix
  // This is done by draining the matrix by columns into one long column vector.
  // We have to transpose, as rulinalg stores data in row major form.
  let boundary_vec = Vector::<f64>::from(
    pot_mat.transpose()
           .into_vec()
           .iter()
           .map(|el| { el * -1.0})
           .collect::<Vec<f64>>());
  
  // 4. create node coefficient matrix
  let main_diag = Vector::<f64>::ones(node_count).apply(&|el| { el * -4.0 });
  let mut coef_mat = Matrix::from_diag(main_diag.data());
  // add neighboring diagonal ones
  for el in coef_mat.iter_diag_mut(DiagOffset::Above(1)) { *el = 1.0; }
  for el in coef_mat.iter_diag_mut(DiagOffset::Below(1)) { *el = 1.0; }
  for el in coef_mat.iter_diag_mut(DiagOffset::Above(col_count)) { *el = 1.0; }
  for el in coef_mat.iter_diag_mut(DiagOffset::Below(col_count)) { *el = 1.0; }

  // remove wrapped nodes. These are places where the coefficient is actually
  // met by a boundary condition (corners and edges)
  let wrapped_coords: Vec<[usize; 2]> = ((col_count)..node_count)
    .step_by(col_count)
    .zip(((col_count - 1)..node_count).step_by(col_count))
    .map(|(val_1, val_2)| [val_1, val_2]) // convert to slice
    .collect();

  for coord in wrapped_coords {
    coef_mat[coord] = 0.0;
    let transpose_slice = [coord[1], coord[0]];
    coef_mat[transpose_slice] = 0.0;
  }

  // println!("Coefficient matrix: ");
  // println!("");
  // for row in coef_mat.iter_rows() {
  //   println!("{:?}", row);
  // }
  // println!("");
  // 
  // println!("Constant vector: ");
  // for el in boundary_vec.iter() {
  //   println!("{}", el);
  // }

  coef_mat.solve(boundary_vec).and_then(|pot_vec| {
    Ok(Matrix::new(row_count, col_count, pot_vec).transpose())
  })
}

/// Computes the integral (sum) over a rectangular contour, starting from the
/// bottom edge and proceeding counterclockwise.
pub fn contour_int_rect(field_mat: &Matrix<f64>,
                        rect_coords: &[[usize; 2]; 4],
                        pad_inward: bool) -> f64 {

  let mut padded_coords = rect_coords.clone();
  
  // Check coords to see if they fit inside the matrix
  if rect_coords.iter().any(|coord| {
    coord[0] > field_mat.rows() || coord[1] > field_mat.cols()
  }) {
    if pad_inward {
      for coord in padded_coords.iter_mut() {
        if coord[0] > field_mat.rows() { coord[0] = field_mat.rows(); }
        if coord[1] > field_mat.cols() { coord[1] = field_mat.cols(); }
      }
    } else {
      panic!("Contour integral coordinates didn't fit inside the matrix. \
              They were: {:?}", rect_coords);
    }
  }

  ccw_rect_from_corners(padded_coords)
    .iter()
    .fold(0.0, |acc, &coord| {
      acc + field_mat[coord]
    })
}

/// Creates a vector of coordinates proceeding in counterclockwise order,
/// starting from the bottom edge.
/// TODO: make edge order irrelevant
fn ccw_rect_from_corners(rect_coords: [[usize; 2]; 4]) -> Vec<[usize; 2]> {
  let bottom_horiz = rect_coords[0][0]...rect_coords[1][0];
  // reverse here because Rust still doesn't do decreasing ranges
  let right_vert = (rect_coords[2][1]...rect_coords[1][1]).rev();
  let top_horiz = (rect_coords[3][0]...rect_coords[2][0]).rev();
  // non-inclusive range because we expect the coordinates to loop back
  let left_vert = rect_coords[2][1]..rect_coords[0][1];

  let mut bottom_vec: Vec<[usize; 2]> = bottom_horiz
    .map(|bottom_x| [bottom_x, rect_coords[0][1]])
    .collect::<Vec<[usize; 2]>>();
  let mut right_vec: Vec<[usize; 2]> = right_vert
    .skip(1)
    .map(|right_y| [rect_coords[1][0], right_y])
    .collect();
  let mut top_vec: Vec<[usize; 2]> = top_horiz
    .skip(1)
    .map(|top_x| [top_x, rect_coords[2][1]])
    .collect();
  let mut left_vec: Vec<[usize; 2]> = left_vert
    .skip(1)
    .map(|left_y| [rect_coords[3][0], left_y])
    .collect();

  bottom_vec.append(&mut right_vec);
  bottom_vec.append(&mut top_vec);
  bottom_vec.append(&mut left_vec);

  bottom_vec
}

/*
 * Private function tests
*/

#[test]
fn rect_small_square() {
  /* 
   * [0, 0]              [1, 0]
   *
   *
   *
   * [0, 1]              [1, 1]
   */
  let coords = [[0, 1], [1, 1], [1, 0], [0, 0]];
  let rect = vec![[0, 1], [1, 1], [1, 0], [0, 0]];
  assert_eq!(rect, ccw_rect_from_corners(coords));
}

#[test]
fn rect_larger() {
  /* 
   * [2, 1]              [8, 1]
   *
   *
   *
   * [2, 5]              [8, 5]
   */

  
  let coords = [[2, 5], [8, 5], [8, 1], [2, 1]];
  let rect = vec![[2, 5], [3, 5], [4, 5], [5, 5], [6, 5], [7, 5], [8, 5],
                  [8, 4], [8, 3], [8, 2], [8, 1],
                  [7, 1], [6, 1], [5, 1], [4, 1], [3, 1], [2, 1],
                  [2, 2], [2, 3], [2, 4]];

  assert_eq!(rect, ccw_rect_from_corners(coords));
}

#[cfg(test)]
mod tests {
  use rulinalg::matrix::Matrix;
  use contour_int_rect;

  #[test]
  fn contour_int_small_square() {
    let ones_mat = Matrix::<f64>::ones(2,2);
    let rect_coords = [[0, 1], [1, 1], [1, 0], [0, 0]];

    assert_eq!(contour_int_rect(&ones_mat, &rect_coords, false), 4.0);
  }

  #[test]
  fn contour_int_larger_rect() {
    let zeros_mat = Matrix::<f64>::zeros(4,5);
    let rect_coords = [[1, 4], [3, 4], [3, 2], [1, 2]];

    assert_eq!(contour_int_rect(&zeros_mat, &rect_coords, false), 0.0);
  }
}