use super::functional::{HelmholtzEnergyFunctional, DFT};
use super::functional_contribution::FunctionalContribution;
use super::weight_functions::WeightFunctionInfo;
use feos_core::{Components, Contributions, EosResult, EosUnit, PhaseEquilibrium};
use ndarray::*;
use num_dual::Dual2_64;
use quantity::si::{SIArray1, SIArray2, SINumber, SIUnit};
use std::ops::AddAssign;

impl WeightFunctionInfo<Dual2_64> {
    fn pdgt_weight_constants(&self) -> (Array2<f64>, Array2<f64>, Array2<f64>) {
        let k = Dual2_64::from(0.0).derivative();
        let w = self.weight_constants(k, 1);
        (w.mapv(|w| w.re), w.mapv(|w| -w.v1), w.mapv(|w| -0.5 * w.v2))
    }
}

impl dyn FunctionalContribution {
    pub fn pdgt_properties(
        &self,
        temperature: f64,
        density: &Array2<f64>,
        helmholtz_energy_density: &mut Array1<f64>,
        first_partial_derivatives: Option<&mut Array2<f64>>,
        second_partial_derivatives: Option<&mut Array3<f64>>,
        influence_diagonal: Option<&mut Array2<f64>>,
        influence_matrix: Option<&mut Array3<f64>>,
    ) -> EosResult<()> {
        // calculate weighted densities
        let weight_functions = self.weight_functions_pdgt(Dual2_64::from(temperature));
        let (w0, w1, w2) = weight_functions.pdgt_weight_constants();
        let weighted_densities = w0.dot(density);

        // calculate Helmholtz energy and derivatives
        let w = weighted_densities.shape()[0]; // # of weighted densities
        let s = density.shape()[0]; // # of segments
        let n = density.shape()[1]; // # of grid points
        let mut df = Array::zeros((w, n));
        let mut d2f = Array::zeros((w, w, n));
        self.second_partial_derivatives(
            temperature,
            weighted_densities.view(),
            helmholtz_energy_density.view_mut(),
            df.view_mut(),
            d2f.view_mut(),
        )?;

        // calculate partial derivatives w.r.t. density
        if let Some(df_drho) = first_partial_derivatives {
            df_drho.assign(&df.t().dot(&w0));
        }

        // calculate second partial derivatives w.r.t. density
        if let Some(d2f_drho2) = second_partial_derivatives {
            for i in 0..s {
                for j in 0..s {
                    for alpha in 0..w {
                        for beta in 0..w {
                            d2f_drho2
                                .index_axis_mut(Axis(0), i)
                                .index_axis_mut(Axis(0), j)
                                .add_assign(
                                    &(&d2f.index_axis(Axis(0), alpha).index_axis(Axis(0), beta)
                                        * w0[(alpha, i)]
                                        * w0[(beta, j)]),
                                );
                        }
                    }
                }
            }
        }

        // calculate influence diagonal
        if let Some(c) = influence_diagonal {
            for i in 0..s {
                for alpha in 0..w {
                    for beta in 0..w {
                        c.index_axis_mut(Axis(0), i).add_assign(
                            &(&d2f.index_axis(Axis(0), alpha).index_axis(Axis(0), beta)
                                * (w1[(alpha, i)] * w1[(beta, i)]
                                    - w0[(alpha, i)] * w2[(beta, i)]
                                    - w2[(alpha, i)] * w0[(beta, i)])),
                        );
                    }
                }
            }
        }

        // calculate influence matrix
        if let Some(c) = influence_matrix {
            for i in 0..s {
                for j in 0..s {
                    for alpha in 0..w {
                        for beta in 0..w {
                            c.index_axis_mut(Axis(0), i)
                                .index_axis_mut(Axis(0), j)
                                .add_assign(
                                    &(&d2f.index_axis(Axis(0), alpha).index_axis(Axis(0), beta)
                                        * (w1[(alpha, i)] * w1[(beta, j)]
                                            - w0[(alpha, i)] * w2[(beta, j)]
                                            - w2[(alpha, i)] * w0[(beta, j)])),
                                );
                        }
                    }
                }
            }
        }
        Ok(())
    }

    pub fn influence_diagonal(
        &self,
        temperature: SINumber,
        density: &SIArray2,
    ) -> EosResult<(SIArray1, SIArray2)> {
        let t = temperature.to_reduced(SIUnit::reference_temperature())?;
        let n = density.shape()[1];
        let mut f = Array::zeros(n);
        let mut c = Array::zeros(density.raw_dim());
        self.pdgt_properties(
            t,
            &density.to_reduced(SIUnit::reference_density())?,
            &mut f,
            None,
            None,
            Some(&mut c),
            None,
        )?;
        Ok((
            f * t * SIUnit::reference_pressure(),
            c * t * SIUnit::reference_influence_parameter(),
        ))
    }
}

impl<T: HelmholtzEnergyFunctional> DFT<T> {
    pub fn solve_pdgt(
        &self,
        vle: &PhaseEquilibrium<Self, 2>,
        n_grid: usize,
        reference_component: usize,
        z: Option<(&mut SIArray1, &mut SINumber)>,
    ) -> EosResult<(SIArray2, SINumber)> {
        // calculate density profile
        let density = if self.components() == 1 {
            let delta_rho = (vle.vapor().density - vle.liquid().density) / (n_grid + 1) as f64;
            SIArray1::linspace(
                vle.liquid().density + delta_rho,
                vle.vapor().density - delta_rho,
                n_grid,
            )?
            .insert_axis(Axis(0))
        } else {
            self.pdgt_density_profile_mix(vle, n_grid, reference_component)?
        };

        // calculate Helmholtz energy density and influence parameter
        let mut delta_omega = Array::zeros(n_grid) * SIUnit::reference_pressure();
        let mut influence_diagonal =
            Array::zeros(density.raw_dim()) * SIUnit::reference_influence_parameter();
        for contribution in self.contributions() {
            let (f, c) = contribution.influence_diagonal(vle.vapor().temperature, &density)?;
            delta_omega += &f;
            influence_diagonal += &c;
        }
        delta_omega += &self
            .ideal_chain_contribution()
            .helmholtz_energy_density::<Ix1>(vle.vapor().temperature, &density)?;

        // calculate excess grand potential density
        let mu_res = vle.vapor().residual_chemical_potential();
        for i in 0..self.components() {
            let rhoi = density.index_axis(Axis(0), i).to_owned();
            let rhoi_b = vle.vapor().partial_density.get(i);
            let mui_res = mu_res.get(i);
            let kt = SIUnit::gas_constant() * vle.vapor().temperature;
            delta_omega +=
                &(&rhoi * (kt * (rhoi.to_reduced(rhoi_b)?.mapv(f64::ln) - 1.0) - mui_res));
        }
        delta_omega += vle.vapor().pressure(Contributions::Total);

        // calculate density gradients w.r.t. reference density
        let dx = density.get((reference_component, 0)) - density.get((reference_component, 1));
        let drho = density.gradient(
            -dx,
            &vle.liquid().partial_density,
            &vle.vapor().partial_density,
        );

        // calculate integrand
        let gamma_int =
            ((influence_diagonal * delta_omega.clone() * 2.0).sqrt()? * drho).sum_axis(Axis(0));

        // calculate z-axis
        if let Some((z, w)) = z {
            let z_int = gamma_int.clone() / (delta_omega * 2.0);
            *z = z_int.integrate_trapezoidal_cumulative(dx);

            // calculate equimolar surface
            let rho_v = density.index_axis(Axis(1), 0).sum();
            let rho_l = density.index_axis(Axis(1), n_grid - 1).sum();
            let rho_r = (density.sum_axis(Axis(0)) - rho_v) / (rho_l - rho_v);
            let ze = (rho_r.clone() * z_int.clone()).integrate_trapezoidal(dx);

            // calculate interfacial width
            *w = (rho_r * z.clone() * z_int).integrate_trapezoidal(dx);
            *w = (24.0 * (*w - 0.5 * ze.powi(2))).sqrt()?;

            // shift density profile
            *z -= ze;
        }

        // integration weights (First and last points are 0)
        let mut weights = Array::ones(n_grid);
        weights[0] = 7.0 / 6.0;
        weights[1] = 23.0 / 24.0;
        weights[n_grid - 2] = 23.0 / 24.0;
        weights[n_grid - 1] = 7.0 / 6.0;
        let weights = weights * dx;

        // calculate surface tension
        Ok((density, gamma_int.integrate(&[weights])))
    }

    fn pdgt_density_profile_mix(
        &self,
        _vle: &PhaseEquilibrium<Self, 2>,
        _n_grid: usize,
        _reference_component: usize,
    ) -> EosResult<SIArray2> {
        unimplemented!()
    }
}