use crate::{H, NX, NY, PEREODIC};
use ndarray::Array2;
use pdifflib::field::Grid;
use pdifflib::finit_diff::poisson_relax;
use pdifflib::{dx, dx_b, dx_f, dy, dy_b, dy_f, jacobian, lap, mx_b, my_b};
use serde::{Deserialize, Serialize};

#[derive(Serialize, Deserialize, Copy, Clone, Debug)]
pub struct Params {
    pub rel: f64,
    pub rel_c: f64,
    pub le: f64,
    pub sc: f64,
    pub pe: f64,
    pub ma: f64,
    pub time: f64,
}

pub struct TemperatureEquation {
    grid: Grid,
    delta: Array2<f64>,
    le: f64,
}

impl TemperatureEquation {
    pub fn new(nx: usize, ny: usize, le: f64) -> Self {
        Self {
            grid: Grid::new(nx, ny),
            delta: Array2::<f64>::zeros((nx, ny)),
            le,
        }
    }

    pub fn step(&mut self, temp: &mut Array2<f64>, psi: &Array2<f64>, dt: f64) {
        let inv_h2 = 1.0 / (H * H);
        for (i, j) in self.grid.inner_nodes() {
            self.delta[[i, j]] =
                dt * inv_h2 * (self.le * lap!(temp, i, j) + jacobian!(psi, temp, i, j));
        }
        *temp += &self.delta;
    }
}

pub struct PhiEquation {
    grid: Grid,
    delta: Array2<f64>,
    rel: f64,
    rel_c: f64,
    sc: f64,
    le: f64,
}

impl PhiEquation {
    pub fn new(nx: usize, ny: usize, rel: f64, rel_c: f64, sc: f64, le: f64) -> Self {
        Self {
            grid: Grid::new(nx, ny),
            delta: Array2::<f64>::zeros((nx, ny)),
            rel,
            rel_c,
            sc,
            le,
        }
    }

    pub fn step(
        &mut self,
        phi: &mut Array2<f64>,
        psi: &Array2<f64>,
        temp: &Array2<f64>,
        conc: &Array2<f64>,
        dt: f64,
    ) {
        let inv_h = 1.0 / H;
        let inv_h2 = inv_h * inv_h;
        for (i, j) in self.grid.inner_nodes() {
            let archim = self.le * self.rel * dx!(temp, i, j) - self.rel_c * dx!(conc, i, j);
            let tmp = self.sc * (lap!(phi, i, j) * inv_h2 + archim * inv_h)
                + jacobian!(psi, phi, i, j) * inv_h2;
            self.delta[[i, j]] = tmp * dt;
        }
        *phi += &self.delta;
    }
}

pub struct PsiEquation;

impl PsiEquation {
    pub fn new() -> Self {
        Self
    }

    pub fn step(&mut self, psi: &mut Array2<f64>, phi: &Array2<f64>) {
        poisson_relax(&phi, psi, H, PEREODIC);
    }
}

pub struct ConcentrationEquation {
    grid: Grid,
    qew: Array2<f64>,
    qsn: Array2<f64>,
    vx: Array2<f64>,
    vy: Array2<f64>,
    pe: f64,
}

impl ConcentrationEquation {
    pub fn new(nx: usize, ny: usize, pe: f64) -> Self {
        Self {
            grid: Grid::new(nx, ny),
            qew: Array2::<f64>::zeros((nx, ny)),
            qsn: Array2::<f64>::zeros((nx, ny)),
            vx: Array2::<f64>::zeros((nx, ny)),
            vy: Array2::<f64>::zeros((nx, ny)),
            pe,
        }
    }

    pub fn step(
        &mut self,
        conc: &mut Array2<f64>,
        psi: &Array2<f64>,
        _temp: &Array2<f64>,
        dt: f64,
    ) {
        self.calc_v(psi);

        let grid = self.grid;
        for (i, k) in grid.x_faces() {
            let mut tmp = self.vx[[i, k]] * mx_b!(conc, i, k);
            tmp += -dx_b!(conc, i, k);
            self.qew[[i, k]] = tmp / H;
        }
        for (i, k) in grid.y_faces() {
            let mut tmp = ((-self.vy[[i, k]]) + self.pe * H) * my_b!(conc, i, k);
            tmp += -dy_b!(conc, i, k);
            self.qsn[[i, k]] = tmp / H;
        }

        for (i, k) in grid.inner_nodes() {
            let q = -(dx_f!(self.qew, i, k) + dy_f!(self.qsn, i, k));
            conc[[i, k]] += dt / H * q;
        }
        for i in 1..NX - 1 {
            let k = 0;
            conc[[i, k]] += dt / H * (-dx_f!(self.qew, i, k) - 2.0 * self.qsn[[i, k + 1]]);

            let k = NY - 1;
            conc[[i, k]] += dt / H * (-dx_f!(self.qew, i, k) + 2.0 * self.qsn[[i, k]]);
        }
        if PEREODIC {
            for k in 0..NY {
                conc[[0, k]] = conc[[NX - 2, k]];
                conc[[NX - 1, k]] = conc[[1, k]];
            }
        } else {
            for k in 1..NY - 1 {
                let i = 0;
                conc[[i, k]] += dt / H * (-dy_f!(self.qsn, i, k) - 2.0 * self.qew[[i + 1, k]]);

                let i = NX - 1;
                conc[[i, k]] += dt / H * (-dy_f!(self.qsn, i, k) + 2.0 * self.qew[[i, k]]);
            }
            conc[[0, 0]] += dt / H * 2.0 * (-self.qew[[1, 0]] - self.qsn[[0, 1]]);
            conc[[0, NY - 1]] += dt / H * 2.0 * (-self.qew[[1, NY - 1]] + self.qsn[[0, NY - 1]]);
            conc[[NX - 1, 0]] += dt / H * 2.0 * (self.qew[[NX - 1, 0]] - self.qsn[[NX - 1, 1]]);
            conc[[NX - 1, NY - 1]] +=
                dt / H * 2.0 * (self.qew[[NX - 1, NY - 1]] + self.qsn[[NX - 1, NY - 1]]);
        }
    }

    fn calc_v(&mut self, psi: &Array2<f64>) {
        for i in 1..NX {
            let mut k = 0;
            self.vx[[i, k]] = 0.5 * (psi[[i, k + 1]] + psi[[i - 1, k + 1]]);

            for k in 1..NY - 1 {
                self.vx[[i, k]] = 0.5 * (dy!(psi, i, k) + dy!(psi, i - 1, k));
            }
            k = NY - 1;
            self.vx[[i, k]] = -0.5 * (psi[[i, k - 1]] + psi[[i - 1, k - 1]]);
        }

        for i in 1..NX - 1 {
            for k in 1..NY {
                self.vy[[i, k]] = 0.5 * (dx!(psi, i, k) + dx!(psi, i, k - 1));
            }
        }

        for k in 1..NY {
            let i = 0;
            self.vy[[i, k]] = 0.5 * (psi[[i + 1, k]] + psi[[i + 1, k - 1]]);

            let i = NX - 1;
            self.vy[[i, k]] = -0.5 * (psi[[i - 1, k]] + psi[[i - 1, k - 1]]);
        }
    }
}