use crate::hard_sphere::{HardSphereProperties, MonomerShape};
use feos_core::parameter::{Parameter, ParameterError, PureRecord};
use ndarray::{Array, Array1, Array2};
use num_dual::DualNum;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fmt::Write;

/// PeTS parameters for a pure substance.
#[derive(Serialize, Deserialize, Debug, Clone, Default)]
pub struct PetsRecord {
    /// Segment diameter in units of Angstrom
    pub sigma: f64,
    /// Energetic parameter in units of Kelvin
    pub epsilon_k: f64,
    /// Entropy scaling parameters for viscosity
    #[serde(skip_serializing_if = "Option::is_none")]
    pub viscosity: Option<[f64; 4]>,
    /// Entropy scaling parameters for self-diffusion coefficient
    #[serde(skip_serializing_if = "Option::is_none")]
    pub diffusion: Option<[f64; 5]>,
    /// Entropy scaling parameters for thermal conductivity
    #[serde(skip_serializing_if = "Option::is_none")]
    pub thermal_conductivity: Option<[f64; 4]>,
}

impl std::fmt::Display for PetsRecord {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "PetsRecord(sigma={}", self.sigma)?;
        write!(f, ", epsilon_k={}", self.epsilon_k)?;
        if let Some(n) = &self.viscosity {
            write!(f, ", viscosity={:?}", n)?;
        }
        if let Some(n) = &self.diffusion {
            write!(f, ", diffusion={:?}", n)?;
        }
        if let Some(n) = &self.thermal_conductivity {
            write!(f, ", thermal_conductivity={:?}", n)?;
        }
        write!(f, ")")
    }
}

impl PetsRecord {
    /// New PeTS parameters for a pure substance.
    ///
    /// # Example
    ///
    /// ```
    /// use feos::pets::PetsRecord;
    /// let record = PetsRecord::new(3.7, 120.0, None, None, None);
    /// ```
    pub fn new(
        sigma: f64,
        epsilon_k: f64,
        viscosity: Option<[f64; 4]>,
        diffusion: Option<[f64; 5]>,
        thermal_conductivity: Option<[f64; 4]>,
    ) -> PetsRecord {
        PetsRecord {
            sigma,
            epsilon_k,
            viscosity,
            diffusion,
            thermal_conductivity,
        }
    }
}

/// Parameters that modify binary interactions.
///
/// $\varepsilon_{k,ij} = (1 - k_{ij})\sqrt{\varepsilon_{k,i} \varepsilon_{k,j}}$
#[derive(Serialize, Deserialize, Clone, Default, Debug)]
pub struct PetsBinaryRecord {
    k_ij: f64,
}

impl From<f64> for PetsBinaryRecord {
    fn from(k_ij: f64) -> Self {
        Self { k_ij }
    }
}

impl From<PetsBinaryRecord> for f64 {
    fn from(binary_record: PetsBinaryRecord) -> Self {
        binary_record.k_ij
    }
}

impl std::fmt::Display for PetsBinaryRecord {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "PetsBinaryRecord(k_ij={})", self.k_ij)
    }
}

/// Parameter set for the PeTS equation of state and Helmholtz energy functional.
#[derive(Serialize, Deserialize, Debug, Clone, Default)]
pub struct PetsParameters {
    /// molar weight in gram per mole
    pub molarweight: Array1<f64>,
    /// Lennard-Jones diameter in Angstrom
    pub sigma: Array1<f64>,
    /// Lennard-Jones energy parameter in Kelvin
    pub epsilon_k: Array1<f64>,
    /// binary interaction parameter
    pub k_ij: Option<Array2<f64>>,
    /// diameter matrix
    pub sigma_ij: Array2<f64>,
    /// energy parameter matrix including k_ij
    pub epsilon_k_ij: Array2<f64>,
    /// energy parameter matrix not including k_ij
    pub e_k_ij: Array2<f64>,
    /// viscosity parameters for entropy scaling
    pub viscosity: Option<Array2<f64>>,
    /// diffusion parameters for entropy scaling
    pub diffusion: Option<Array2<f64>>,
    /// thermal conductivity parameters for entropy scaling
    pub thermal_conductivity: Option<Array2<f64>>,
    /// records of all pure substances of the system
    pub pure_records: Vec<PureRecord<PetsRecord>>,
    /// records of all binary interaction parameters
    pub binary_records: Option<Array2<PetsBinaryRecord>>,
}

impl Parameter for PetsParameters {
    type Pure = PetsRecord;
    type Binary = PetsBinaryRecord;

    fn from_records(
        pure_records: Vec<PureRecord<Self::Pure>>,
        binary_records: Option<Array2<PetsBinaryRecord>>,
    ) -> Result<Self, ParameterError> {
        let n = pure_records.len();

        let mut molarweight = Array::zeros(n);
        let mut sigma = Array::zeros(n);
        let mut epsilon_k = Array::zeros(n);
        let mut viscosity = Vec::with_capacity(n);
        let mut diffusion = Vec::with_capacity(n);
        let mut thermal_conductivity = Vec::with_capacity(n);

        let mut component_index = HashMap::with_capacity(n);

        for (i, record) in pure_records.iter().enumerate() {
            component_index.insert(record.identifier.clone(), i);
            let r = &record.model_record;
            sigma[i] = r.sigma;
            epsilon_k[i] = r.epsilon_k;
            viscosity.push(r.viscosity);
            diffusion.push(r.diffusion);
            thermal_conductivity.push(r.thermal_conductivity);
            molarweight[i] = record.molarweight;
        }

        let k_ij = binary_records.as_ref().map(|br| br.map(|br| br.k_ij));
        let mut sigma_ij = Array::zeros((n, n));
        let mut e_k_ij = Array::zeros((n, n));
        for i in 0..n {
            for j in 0..n {
                e_k_ij[[i, j]] = (epsilon_k[i] * epsilon_k[j]).sqrt();
                sigma_ij[[i, j]] = 0.5 * (sigma[i] + sigma[j]);
            }
        }
        let mut epsilon_k_ij = e_k_ij.clone();
        if let Some(k_ij) = k_ij.as_ref() {
            epsilon_k_ij *= &(1.0 - k_ij);
        }

        let viscosity_coefficients = if viscosity.iter().any(|v| v.is_none()) {
            None
        } else {
            let mut v = Array2::zeros((4, viscosity.len()));
            for (i, vi) in viscosity.iter().enumerate() {
                v.column_mut(i).assign(&Array1::from(vi.unwrap().to_vec()));
            }
            Some(v)
        };

        let diffusion_coefficients = if diffusion.iter().any(|v| v.is_none()) {
            None
        } else {
            let mut v = Array2::zeros((5, diffusion.len()));
            for (i, vi) in diffusion.iter().enumerate() {
                v.column_mut(i).assign(&Array1::from(vi.unwrap().to_vec()));
            }
            Some(v)
        };

        let thermal_conductivity_coefficients = if thermal_conductivity.iter().any(|v| v.is_none())
        {
            None
        } else {
            let mut v = Array2::zeros((4, thermal_conductivity.len()));
            for (i, vi) in thermal_conductivity.iter().enumerate() {
                v.column_mut(i).assign(&Array1::from(vi.unwrap().to_vec()));
            }
            Some(v)
        };

        Ok(Self {
            molarweight,
            sigma,
            epsilon_k,
            k_ij,
            sigma_ij,
            epsilon_k_ij,
            e_k_ij,
            viscosity: viscosity_coefficients,
            diffusion: diffusion_coefficients,
            thermal_conductivity: thermal_conductivity_coefficients,
            pure_records,
            binary_records,
        })
    }

    fn records(&self) -> (&[PureRecord<PetsRecord>], Option<&Array2<PetsBinaryRecord>>) {
        (&self.pure_records, self.binary_records.as_ref())
    }
}

impl HardSphereProperties for PetsParameters {
    fn monomer_shape<N: DualNum<f64>>(&self, _: N) -> MonomerShape<N> {
        MonomerShape::Spherical(self.sigma.len())
    }

    fn hs_diameter<D: DualNum<f64> + Copy>(&self, temperature: D) -> Array1<D> {
        let ti = temperature.recip() * -3.052785558;
        Array::from_shape_fn(self.sigma.len(), |i| {
            -((ti * self.epsilon_k[i]).exp() * 0.127112544 - 1.0) * self.sigma[i]
        })
    }
}

impl PetsParameters {
    pub fn to_markdown(&self) -> String {
        let mut output = String::new();
        let o = &mut output;
        write!(
            o,
            "|component|molarweight|$\\sigma$|$\\varepsilon$|\n|-|-|-|-|"
        )
        .unwrap();
        for i in 0..self.sigma.len() {
            let component = self.pure_records[i].identifier.name.clone();
            let component = component.unwrap_or(format!("Component {}", i + 1));
            write!(
                o,
                "\n|{}|{}|{}|{}|",
                component, self.molarweight[i], self.sigma[i], self.epsilon_k[i],
            )
            .unwrap();
        }

        output
    }
}

impl std::fmt::Display for PetsParameters {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "PetsParameters(")?;
        write!(f, "\n\tmolarweight={}", self.molarweight)?;
        write!(f, "\n\tsigma={}", self.sigma)?;
        write!(f, "\n\tepsilon_k={}", self.epsilon_k)?;
        if let Some(k_ij) = self.k_ij.as_ref() {
            write!(f, "\n\tk_ij=\n{}", k_ij)?;
        }
        write!(f, "\n)")
    }
}

#[cfg(test)]
pub mod utils {
    use super::*;
    use std::sync::Arc;

    pub fn argon_parameters() -> Arc<PetsParameters> {
        let argon_json = r#"
            {
                "identifier": {
                    "cas": "7440-37-1",
                    "name": "argon",
                    "iupac_name": "argon",
                    "smiles": "[Ar]",
                    "inchi": "InChI=1/Ar",
                    "formula": "Ar"
                },
                "model_record": {
                    "sigma": 3.4050,
                    "epsilon_k": 119.8,
                    "viscosity": [0.0, 0.0, 0.0, 0.0],
                    "thermal_conductivity": [0.0, 0.0, 0.0, 0.0],
                    "diffusion": [0.0, 0.0, 0.0, 0.0, 0.0]
                },
                "molarweight": 39.948
            }"#;
        let argon_record: PureRecord<PetsRecord> =
            serde_json::from_str(argon_json).expect("Unable to parse json.");
        Arc::new(PetsParameters::new_pure(argon_record).unwrap())
    }

    pub fn krypton_parameters() -> Arc<PetsParameters> {
        let krypton_json = r#"
            {
                "identifier": {
                    "cas": "7439-90-9",
                    "name": "krypton",
                    "iupac_name": "krypton",
                    "smiles": "[Kr]",
                    "inchi": "InChI=1S/Kr",
                    "formula": "Kr"
                },
                "model_record": {
                    "sigma": 3.6300,
                    "epsilon_k": 163.10
                },
                "molarweight": 83.798
            }"#;
        let krypton_record: PureRecord<PetsRecord> =
            serde_json::from_str(krypton_json).expect("Unable to parse json.");
        Arc::new(PetsParameters::new_pure(krypton_record).unwrap())
    }

    pub fn argon_krypton_parameters() -> Arc<PetsParameters> {
        let binary_json = r#"[
            {
                "identifier": {
                    "cas": "7440-37-1",
                    "name": "argon",
                    "iupac_name": "argon",
                    "smiles": "[Ar]",
                    "inchi": "1/Ar",
                    "formula": "Ar"
                },
                "model_record": {
                    "sigma": 3.4050,
                    "epsilon_k": 119.8,
                    "viscosity": [0.0, 0.0, 0.0, 0.0],
                    "thermal_conductivity": [0.0, 0.0, 0.0, 0.0],
                    "diffusion": [0.0, 0.0, 0.0, 0.0, 0.0]
                },
                "molarweight": 39.948
            },
            {
                "identifier": {
                    "cas": "7439-90-9",
                    "name": "krypton",
                    "iupac_name": "krypton",
                    "smiles": "[Kr]",
                    "inchi": "InChI=1S/Kr",
                    "formula": "Kr"
                },
                "model_record": {
                    "sigma": 3.6300,
                    "epsilon_k": 163.10,
                    "viscosity": [0.0, 0.0, 0.0, 0.0],
                    "thermal_conductivity": [0.0, 0.0, 0.0, 0.0],
                    "diffusion": [0.0, 0.0, 0.0, 0.0, 0.0]
                },
                "molarweight": 83.798
            }
        ]"#;
        let binary_record: Vec<PureRecord<PetsRecord>> =
            serde_json::from_str(binary_json).expect("Unable to parse json.");
        Arc::new(PetsParameters::new_binary(binary_record, None).unwrap())
    }
}