use crate::{
	model::{Model, Settings},
	solver::Solver,
	utils::*,
};

use nalgebra::ComplexField;
use std::{collections::HashMap, marker::PhantomData};

#[derive(Clone)]
pub struct Point<T: Copy, const D: usize> {
	pub pos: Vect<T, D>,
	pub diffusion: Vect<T, D>,
}

#[derive(Clone)]
pub struct Source<T: Copy, const D: usize> {
	pub amplitude: T,
	pub fonction: fn(T) -> Vect<T, D>,
	pub freq: T,
	pub phase: T,
}

pub struct Space<T: Copy, S: Settings, M: Model<T, S, D>, V: Solver<T, S, M, D>, const D: usize> {
	pub model: M,
	pub old_points: Vec<Point<T, D>>,
	pub points: Vec<Point<T, D>>,
	pub size: (usize, usize),
	pub solver: V,
	pub sources: HashMap<usize, Source<T, D>>,
	pub time: T,
	pub _p: PhantomData<S>,
}

impl<
		T: Copy + ComplexField<RealField = T>,
		S: Settings,
		M: Model<T, S, D>,
		V: Solver<T, S, M, D>,
		const D: usize,
	> Space<T, S, M, V, D>
{
	pub fn simulate(&mut self, delta_time: T) {
		std::mem::swap(&mut self.old_points, &mut self.points);
		for (i, (point, old_point)) in self
			.points
			.iter_mut()
			.zip(self.old_points.iter())
			.enumerate()
		{
			if let Some(source) = self.sources.get(&i) {
				let t = self.time * source.freq + source.phase;
				//Point{pos:t.sin()*source.amplitude, speed: t.cos()*source.amplitude}
				point.pos = (source.fonction)(t) * source.amplitude;
			} else {
				*point = old_point.clone()
			};
			point.pos = self.solver.f(&self.model, point.pos);
		}

		let mut i = 0usize;
		for y in 0..self.size.1 {
			for x in 0..self.size.0 {
				let old_point = self.old_points[i].clone();
				let point = &mut self.points[i];

				if y > 0 {
					point.pos += (self.old_points[i - self.size.0].pos - old_point.pos)
						.component_mul(&point.diffusion)
						* delta_time;
				}
				if y < self.size.1 - 1 {
					point.pos += (self.old_points[i + self.size.0].pos - old_point.pos)
						.component_mul(&point.diffusion)
						* delta_time;
				}
				if x > 0 {
					point.pos += (self.old_points[i - 1].pos - old_point.pos)
						.component_mul(&point.diffusion)
						* delta_time;
				}
				if x < self.size.0 - 1 {
					point.pos += (self.old_points[i + 1].pos - old_point.pos)
						.component_mul(&point.diffusion)
						* delta_time;
				}

				i += 1;
			}
		}

		self.time += delta_time;
	}
}