use super::parameter::GcPcSaftFunctionalParameters;
use crate::gc_pcsaft::eos::dispersion::{A0, A1, A2, B0, B1, B2};
use crate::hard_sphere::HardSphereProperties;
use feos_core::FeosError;
use feos_dft::{FunctionalContribution, WeightFunction, WeightFunctionInfo, WeightFunctionShape};
use ndarray::*;
use num_dual::DualNum;
use std::f64::consts::{FRAC_PI_6, PI};

#[derive(Clone)]
pub struct AttractiveFunctional<'a> {
    parameters: &'a GcPcSaftFunctionalParameters,
}

impl<'a> AttractiveFunctional<'a> {
    pub fn new(parameters: &'a GcPcSaftFunctionalParameters) -> Self {
        Self { parameters }
    }
}

impl<'a> FunctionalContribution for AttractiveFunctional<'a> {
    fn name(&self) -> &'static str {
        "Attractive functional (GC)"
    }

    fn weight_functions<N: DualNum<f64> + Copy>(&self, temperature: N) -> WeightFunctionInfo<N> {
        let p = &self.parameters;

        let d = p.hs_diameter(temperature);
        WeightFunctionInfo::new(p.component_index.clone(), false).add(
            WeightFunction::new_scaled(
                d.zip_map(&p.psi_dft, |d, p| d * p),
                WeightFunctionShape::Theta,
            ),
            false,
        )
    }

    fn helmholtz_energy_density<N: DualNum<f64> + Copy>(
        &self,
        temperature: N,
        density: ArrayView2<N>,
    ) -> Result<Array1<N>, FeosError> {
        // auxiliary variables
        let p = &self.parameters;
        let n = p.m.len();

        // temperature dependent segment diameter
        let d = p.hs_diameter(temperature);

        // packing fraction
        let d3m = d.zip_map(&p.m, |d, m| d.powi(3) * (m * FRAC_PI_6));
        let eta = density
            .outer_iter()
            .zip(&d3m)
            .map(|(rho, &d3m)| &rho * d3m)
            .reduce(|a, b| a + b)
            .unwrap();

        // mean segment number
        let mut m_i: Array1<f64> = Array::zeros(p.component_index[n - 1] + 1);
        let mut s_i: Array1<f64> = Array::zeros(p.component_index[n - 1] + 1);
        for (&c, &m) in p.component_index.iter().zip(p.m.iter()) {
            m_i[c] += m;
            s_i[c] += 1.0;
        }
        let mut rhog: Array1<N> = Array::zeros(eta.raw_dim());
        let mut m_bar: Array1<N> = Array::zeros(eta.raw_dim());
        for (rho, &c) in density.outer_iter().zip(p.component_index.iter()) {
            m_bar += &(&rho * m_i[c] / s_i[c]);
            rhog += &(&rho / s_i[c]);
        }
        m_bar.iter_mut().zip(rhog.iter()).for_each(|(m, &r)| {
            if r.re() > f64::EPSILON {
                *m /= r
            } else {
                *m = N::one()
            }
        });

        // mixture densities, crosswise interactions of all segments on all chains
        let mut rho1mix: Array1<N> = Array::zeros(eta.raw_dim());
        let mut rho2mix: Array1<N> = Array::zeros(eta.raw_dim());
        for i in 0..n {
            for j in 0..n {
                let eps_ij = temperature.recip() * p.epsilon_k_ij[(i, j)];
                let sigma_ij = p.sigma_ij[(i, j)].powi(3);
                rho1mix = rho1mix
                    + (&density.index_axis(Axis(0), i) * &density.index_axis(Axis(0), j))
                        .mapv(|x| x * (eps_ij * sigma_ij * p.m[i] * p.m[j]));
                rho2mix = rho2mix
                    + (&density.index_axis(Axis(0), i) * &density.index_axis(Axis(0), j))
                        .mapv(|x| x * (eps_ij * eps_ij * sigma_ij * p.m[i] * p.m[j]));
            }
        }

        // I1, I2 and C1
        let mut i1: Array1<N> = Array::zeros(eta.raw_dim());
        let mut i2: Array1<N> = Array::zeros(eta.raw_dim());
        let mut eta_i: Array1<N> = Array::ones(eta.raw_dim());
        let m1 = (m_bar.clone() - 1.0) / &m_bar;
        let m2 = (m_bar.clone() - 2.0) / &m_bar * &m1;
        for i in 0..=6 {
            i1 = i1 + (&m2 * A2[i] + &m1 * A1[i] + A0[i]) * &eta_i;
            i2 = i2 + (&m2 * B2[i] + &m1 * B1[i] + B0[i]) * &eta_i;
            eta_i = &eta_i * &eta;
        }
        let c1 = Zip::from(&eta).and(&m_bar).map_collect(|&eta, &m| {
            (m * (eta * 8.0 - eta.powi(2) * 2.0) / (eta - 1.0).powi(4)
                + (eta * (eta * (eta * (eta * 2.0 - 12.0) + 27.0) - 20.0))
                    / ((eta - 1.0) * (eta - 2.0)).powi(2)
                    * (m - 1.0)
                + 1.0)
                .recip()
        });

        // Helmholtz energy density
        Ok((-rho1mix * i1 * 2.0 - rho2mix * m_bar * c1 * i2) * PI)
    }
}
