use crate::saftvrqmie::eos::FeynmanHibbsOrder;
use core::cmp::max;
use feos_core::parameter::{Parameter, ParameterError, PureRecord};
use ndarray::{Array, Array1, Array2};
use num_traits::Zero;
use quantity::si::{SINumber, ANGSTROM, CALORIE, GRAM, KELVIN, KILO, KILOGRAM, MOL, NAV, RGAS};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::convert::TryFrom;
use std::fmt::Write;
use std::fs::File;
use std::io::BufWriter;

/// SAFT-VRQ Mie pure-component parameters.
#[derive(Serialize, Deserialize, Clone, Default)]
pub struct SaftVRQMieRecord {
    /// Segment number
    pub m: f64,
    /// Segment diameter in units of Angstrom
    pub sigma: f64,
    /// Energetic parameter in units of Kelvin
    pub epsilon_k: f64,
    /// Repulsive Mie exponent
    pub lr: f64,
    /// Attractive Mie exponent
    pub la: f64,
    /// Feynman-Hibbs order
    pub fh: usize,
    /// Entropy scaling coefficients for the viscosity
    #[serde(skip_serializing_if = "Option::is_none")]
    pub viscosity: Option<[f64; 4]>,
    /// Entropy scaling coefficients for the diffusion coefficient
    #[serde(skip_serializing_if = "Option::is_none")]
    pub diffusion: Option<[f64; 5]>,
    /// Entropy scaling coefficients for the thermal conductivity
    #[serde(skip_serializing_if = "Option::is_none")]
    pub thermal_conductivity: Option<[f64; 4]>,
}

impl std::fmt::Display for SaftVRQMieRecord {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "SaftVRQMieRecord(m={}", self.m)?;
        write!(f, ", 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 SaftVRQMieRecord {
    pub fn new(
        m: f64,
        sigma: f64,
        epsilon_k: f64,
        lr: f64,
        la: f64,
        fh: usize,
        viscosity: Option<[f64; 4]>,
        diffusion: Option<[f64; 5]>,
        thermal_conductivity: Option<[f64; 4]>,
    ) -> SaftVRQMieRecord {
        SaftVRQMieRecord {
            m,
            sigma,
            epsilon_k,
            lr,
            la,
            fh,
            viscosity,
            diffusion,
            thermal_conductivity,
        }
    }
}

/// SAFT-VRQ Mie binary mixture parameters.
#[derive(Serialize, Deserialize, Clone, Default)]
pub struct SaftVRQMieBinaryRecord {
    /// correction to energy parameters
    pub k_ij: f64,
    /// correction to diameter
    pub l_ij: f64,
}

impl std::fmt::Display for SaftVRQMieBinaryRecord {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "SaftVRQMieBinaryParameters(")?;
        write!(f, "\n\tk_ij={}", self.k_ij)?;
        write!(f, "\n\tl_ij={}", self.l_ij)?;
        write!(f, "\n)")
    }
}

impl TryFrom<f64> for SaftVRQMieBinaryRecord {
    type Error = ParameterError;

    fn try_from(_f: f64) -> Result<Self, Self::Error> {
        Err(ParameterError::IncompatibleParameters(
            "Cannot infer k_ij and l_ij from single float.".to_string(),
        ))
    }
}

impl TryFrom<SaftVRQMieBinaryRecord> for f64 {
    type Error = ParameterError;

    fn try_from(_f: SaftVRQMieBinaryRecord) -> Result<Self, Self::Error> {
        Err(ParameterError::IncompatibleParameters(
            "Cannot infer k_ij and l_ij from single float.".to_string(),
        ))
    }
}

/// Parameter set required for the SAFT-VRQ Mie equation of state and Helmholtz energy functional.
pub struct SaftVRQMieParameters {
    pub molarweight: Array1<f64>,
    pub m: Array1<f64>,
    pub sigma: Array1<f64>,
    pub epsilon_k: Array1<f64>,
    pub k_ij: Array2<f64>,
    pub l_ij: Array2<f64>,
    pub sigma_ij: Array2<f64>,
    pub epsilon_k_ij: Array2<f64>,
    pub e_k_ij: Array2<f64>,
    pub lr: Array1<f64>,
    pub la: Array1<f64>,
    pub fh: Array1<usize>,
    pub c_ij: Array2<f64>,
    pub lambda_r_ij: Array2<f64>,
    pub lambda_a_ij: Array2<f64>,
    pub mass_ij: Array2<f64>,
    pub viscosity: Option<Array2<f64>>,
    pub diffusion: Option<Array2<f64>>,
    pub thermal_conductivity: Option<Array2<f64>>,
    pub pure_records: Vec<PureRecord<SaftVRQMieRecord>>,
    pub binary_records: Option<Array2<SaftVRQMieBinaryRecord>>,
    pub fh_ij: Array2<FeynmanHibbsOrder>,
}

impl Parameter for SaftVRQMieParameters {
    type Pure = SaftVRQMieRecord;
    type Binary = SaftVRQMieBinaryRecord;

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

        let mut fh = Array1::<usize>::zeros(n);
        let mut molarweight = Array::zeros(n);
        let mut m = Array::zeros(n);
        let mut sigma = Array::zeros(n);
        let mut epsilon_k = Array::zeros(n);
        let mut lr = Array::zeros(n);
        let mut la = 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;
            m[i] = r.m;
            sigma[i] = r.sigma;
            epsilon_k[i] = r.epsilon_k;
            lr[i] = r.lr;
            la[i] = r.la;
            fh[i] = r.fh;
            viscosity.push(r.viscosity);
            diffusion.push(r.diffusion);
            thermal_conductivity.push(r.thermal_conductivity);
            molarweight[i] = record.molarweight;
        }

        let mut fh_ij: Array2<FeynmanHibbsOrder> =
            Array2::from_shape_fn((n, n), |(_i, _j)| FeynmanHibbsOrder::FH0);
        let br = binary_records.as_ref();
        let k_ij = br.map_or_else(|| Array2::zeros([n; 2]), |br| br.mapv(|br| br.k_ij));
        let l_ij = br.map_or_else(|| Array2::zeros([n; 2]), |br| br.mapv(|br| br.l_ij));
        let mut epsilon_k_ij = Array::zeros((n, n));
        let mut sigma_ij = Array::zeros((n, n));
        let mut e_k_ij = Array::zeros((n, n));
        let mut lambda_r_ij = Array::zeros((n, n));
        let mut lambda_a_ij = Array::zeros((n, n));
        let mut c_ij = Array::zeros((n, n));
        let mut mass_ij = Array::zeros((n, n));
        let to_mass_per_molecule = (GRAM / MOL / NAV).to_reduced(KILOGRAM).unwrap();
        for i in 0..n {
            for j in 0..n {
                sigma_ij[[i, j]] = (1.0 - l_ij[[i, j]]) * 0.5 * (sigma[i] + sigma[j]);
                e_k_ij[[i, j]] = (sigma[i].powi(3) * sigma[j].powi(3)).sqrt()
                    / sigma_ij[[i, j]].powi(3)
                    * (epsilon_k[i] * epsilon_k[j]).sqrt();
                epsilon_k_ij[[i, j]] = (1.0 - k_ij[[i, j]]) * e_k_ij[[i, j]];
                lambda_r_ij[[i, j]] = ((lr[i] - 3.0) * (lr[j] - 3.0)).sqrt() + 3.0;
                lambda_a_ij[[i, j]] = ((la[i] - 3.0) * (la[j] - 3.0)).sqrt() + 3.0;
                c_ij[[i, j]] = lambda_r_ij[[i, j]] / (lambda_r_ij[[i, j]] - lambda_a_ij[[i, j]])
                    * (lambda_r_ij[[i, j]] / lambda_a_ij[[i, j]])
                        .powf(lambda_a_ij[[i, j]] / (lambda_r_ij[[i, j]] - lambda_a_ij[[i, j]]));
                mass_ij[[i, j]] = 2.0 * molarweight[i] * molarweight[j]
                    / (molarweight[i] + molarweight[j])
                    * to_mass_per_molecule;
                fh_ij[[i, j]] = FeynmanHibbsOrder::try_from(max(fh[i], fh[j]))?;
                if fh[i] * fh[j] == 2 {
                    return Err(
                        ParameterError::IncompatibleParameters(
                            format!(
                                "cannot combine Feynman-Hibbs orders 1 and 2. Component {} has order {} and component {} has order {}.", 
                                i, fh[i], j, fh[j]
                            )
                        )
                    );
                }
            }
        }

        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,
            m,
            sigma,
            epsilon_k,
            k_ij,
            l_ij,
            sigma_ij,
            epsilon_k_ij,
            e_k_ij,
            lr,
            la,
            fh,
            c_ij,
            lambda_r_ij,
            lambda_a_ij,
            mass_ij,
            viscosity: viscosity_coefficients,
            diffusion: diffusion_coefficients,
            thermal_conductivity: thermal_conductivity_coefficients,
            pure_records,
            binary_records,
            fh_ij,
        })
    }

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

impl SaftVRQMieParameters {
    pub fn to_markdown(&self) -> String {
        let mut output = String::new();
        let o = &mut output;
        write!(
            o,
            "|component|molarweight|$\\sigma$|$\\varepsilon$|$\\lambda_r$|$\\lambda_a$|fh|\n|-|-|-|-|-|-|-|"
        )
        .unwrap();
        for i in 0..self.m.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],
                self.lr[i],
                self.la[i],
                self.fh[i]
            )
            .unwrap();
        }

        output
    }

    /// Generate energy and force tables to be used with LAMMPS' `pair_style table` command.
    ///
    /// For a given `temperature`, `n` values between `r_min` and `r_max` (both including) are tabulated.
    ///
    /// Files for all pure substances and all unique pairs are generated,
    /// where filenames use either the "name" field of the identifier or the index if no name is present.
    ///
    /// # Example
    ///
    /// For a hydrogen-neon mixture at 30 K, three files will be created.
    ///
    /// - "hydrogen_30K.table" for H-H interactions,
    /// - "neon_30K.table" for Ne-Ne interactions,
    /// - "hydrogen_neon_30K.table" for H-Ne interactions.
    pub fn lammps_tables(
        &self,
        temperature: SINumber,
        n: usize,
        r_min: SINumber,
        r_max: SINumber,
    ) -> std::io::Result<()> {
        let t = temperature.to_reduced(KELVIN).unwrap();
        let rs = Array1::linspace(
            r_min.to_reduced(ANGSTROM).unwrap(),
            r_max.to_reduced(ANGSTROM).unwrap(),
            n,
        );
        let energy_conversion = (KELVIN * RGAS / (KILO * CALORIE / MOL))
            .into_value()
            .unwrap();
        let force_conversion = (KELVIN * RGAS / (KILO * CALORIE / MOL))
            .into_value()
            .unwrap();

        let n_components = self.sigma.len();
        for i in 0..n_components {
            for j in i..n_components {
                let name_i = self.pure_records[i]
                    .identifier
                    .name
                    .clone()
                    .unwrap_or_else(|| i.to_string());
                let name_j = self.pure_records[j]
                    .identifier
                    .name
                    .clone()
                    .unwrap_or_else(|| j.to_string());

                let name = if i == j {
                    name_i
                } else {
                    format!("{}_{}", name_i, name_j)
                };
                let f = File::create(format!("{}_{}K.table", name, t))?;
                let mut stream = BufWriter::new(f);

                std::io::Write::write(
                    &mut stream,
                    b"# DATE: YYYY-MM-DD UNITS: real CONTRIBUTOR: YOUR NAME\n",
                )?;
                std::io::Write::write(
                    &mut stream,
                    format!("# FH1 potential for {} at T = {}\n", name, temperature).as_bytes(),
                )?;
                std::io::Write::write(&mut stream, format!("FH1_{}\n", name).as_bytes())?;
                std::io::Write::write(&mut stream, format!("N {}\n\n", n).as_bytes())?;

                for (k, &r) in rs.iter().enumerate() {
                    let [u, du, _] = self.qmie_potential_ij(i, j, r, t);
                    std::io::Write::write(
                        &mut stream,
                        format!(
                            "{} {:12.8} {:12.8} {:12.8}\n",
                            k + 1,
                            r,
                            u * energy_conversion,
                            -du * force_conversion
                        )
                        .as_bytes(),
                    )?;
                }
                std::io::Write::flush(&mut stream)?;
            }
        }
        Ok(())
    }
}

impl std::fmt::Display for SaftVRQMieParameters {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "SaftVRQMieParameters(")?;
        write!(f, "\n\tmolarweight={}", self.molarweight)?;
        write!(f, "\n\tm={}", self.m)?;
        write!(f, "\n\tsigma={}", self.sigma)?;
        write!(f, "\n\tepsilon_k={}", self.epsilon_k)?;

        if !self.k_ij.iter().all(|k| k.is_zero()) {
            write!(f, "\n\tk_ij=\n{}", self.k_ij)?;
        }
        write!(f, "\n)")
    }
}

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

    pub fn hydrogen_fh(fh: &str) -> Arc<SaftVRQMieParameters> {
        let hydrogen_json = &(r#"
            {
                "identifier": {
                    "cas": "1333-74-0",
                    "name": "hydrogen",
                    "iupac_name": "hydrogen",
                    "smiles": "[HH]",
                    "inchi": "InChI=1S/H2/h1H",
                    "formula": "H2"
                },
                "model_record": {
                    "m": 1.0,
                    "sigma": 3.0243,
                    "epsilon_k": 26.706,
                    "lr": 9.0,
                    "la": 6.0,
                    "fh": "#
            .to_owned()
            + fh
            + r#"
                },
                "molarweight": 2.0157309551872
            }"#);
        let hydrogen_record: PureRecord<SaftVRQMieRecord> =
            serde_json::from_str(hydrogen_json).expect("Unable to parse json.");
        Arc::new(SaftVRQMieParameters::new_pure(hydrogen_record).unwrap())
    }

    #[allow(dead_code)]
    pub fn helium_fh1() -> Arc<SaftVRQMieParameters> {
        let helium_json = r#"
            {
                "identifier": {
                    "cas": "1333-74-0",
                    "name": "helium",
                    "iupac_name": "helium",
                    "smiles": "[HH]",
                    "inchi": "InChI=1S/H2/h1H",
                    "formula": "He"
                },
                "model_record": {
                    "m": 1.0,
                    "sigma": 2.7443,
                    "epsilon_k": 5.4195,
                    "lr": 9.0,
                    "la": 6.0,
                    "fh": 1
                },
                "molarweight": 4.002601643881807
            }"#;
        let helium_record: PureRecord<SaftVRQMieRecord> =
            serde_json::from_str(helium_json).expect("Unable to parse json.");
        Arc::new(SaftVRQMieParameters::new_pure(helium_record).unwrap())
    }

    #[allow(dead_code)]
    pub fn neon_fh1() -> Arc<SaftVRQMieParameters> {
        let neon_json = r#"
            {
                "identifier": {
                    "cas": "1333-74-0",
                    "name": "neon",
                    "iupac_name": "neon",
                    "smiles": "[HH]",
                    "inchi": "InChI=1S/H2/h1H",
                    "formula": "Ne"
                },
                "model_record": {
                    "m": 1.0,
                    "sigma": 2.7778,
                    "epsilon_k": 37.501,
                    "lr": 13.0,
                    "la": 6.0,
                    "fh": 1
                },
                "molarweight": 20.17969806457545
            }"#;
        let neon_record: PureRecord<SaftVRQMieRecord> =
            serde_json::from_str(neon_json).expect("Unable to parse json.");
        Arc::new(SaftVRQMieParameters::new_pure(neon_record).unwrap())
    }

    pub fn h2_ne_fh(fh: &str) -> Arc<SaftVRQMieParameters> {
        let binary_json = &(r#"[
            {
                "identifier": {
                    "cas": "1333-74-0",
                    "name": "hydrogen",
                    "iupac_name": "hydrogen",
                    "smiles": "[HH]",
                    "inchi": "InChI=1S/H2/h1H",
                    "formula": "H2"
                },
                "model_record": {
                    "m": 1.0,
                    "sigma": 3.0243,
                    "epsilon_k": 26.706,
                    "lr": 9.0,
                    "la": 6.0,
                    "fh": "#
            .to_owned()
            + fh
            + r#"
                },
                "molarweight": 2.0157309551872
            },
            {
                "identifier": {
                    "cas": "1333-74-0",
                    "name": "neon",
                    "iupac_name": "neon",
                    "smiles": "[HH]",
                    "inchi": "InChI=1S/H2/h1H",
                    "formula": "H2"
                },
                "model_record": {
                    "m": 1.0,
                    "sigma": 2.7778,
                    "epsilon_k": 37.501,
                    "lr": 13.0,
                    "la": 6.0,
                    "fh": "#
            + fh
            + r#"
                },
                "molarweight": 20.17969806457545
            }
        ]"#);
        let binary_record: Vec<PureRecord<SaftVRQMieRecord>> =
            serde_json::from_str(binary_json).expect("Unable to parse json.");
        Arc::new(
            SaftVRQMieParameters::new_binary(
                binary_record,
                Some(SaftVRQMieBinaryRecord {
                    k_ij: 0.105,
                    l_ij: 0.0,
                }),
            )
            .unwrap(),
        )
    }
}

#[cfg(test)]
mod test {
    use super::utils::{helium_fh1, hydrogen_fh};
    use super::SaftVRQMieParameters;
    use feos_core::parameter::Parameter;

    #[test]
    #[should_panic(
        expected = "cannot combine Feynman-Hibbs orders 1 and 2. Component 0 has order 1 and component 1 has order 2."
    )]
    fn incompatible_order() {
        let order1 = helium_fh1();
        let order2 = hydrogen_fh("2");
        SaftVRQMieParameters::new_binary(
            vec![
                order1.pure_records[0].clone(),
                order2.pure_records[0].clone(),
            ],
            None,
        )
        .unwrap();
    }
}