#[cfg(feature = "ffi")]
mod ffi;

use opendp_derive::bootstrap;

use crate::core::{Domain, Function, StabilityMap, Transformation, MetricSpace};
use crate::domains::{AtomDomain, OptionDomain, VectorDomain};
use crate::error::Fallible;
use crate::traits::samplers::SampleUniform;
use crate::traits::{CheckAtom, CheckNull, Float, InherentNull};
use crate::transformations::{make_row_by_row, make_row_by_row_fallible};

use super::DatasetMetric;

#[bootstrap(
    features("contrib"), 
    generics(M(suppress), TA(suppress)),
    derived_types(TA = "$get_atom(get_type(input_domain))")
)]
/// Make a Transformation that replaces NaN values in `Vec<TA>` with uniformly distributed floats within `bounds`.
///
/// # Arguments
/// * `input_domain` - Domain of the input.
/// * `input_metric` - Metric of the input.
/// * `bounds` - Tuple of inclusive lower and upper bounds.
///
/// # Generics
/// * `M` - Metric Type. A dataset metric.
/// * `TA` - Atomic Type of data being imputed. One of `f32` or `f64`
pub fn make_impute_uniform_float<M, TA>(
    input_domain: VectorDomain<AtomDomain<TA>>,
    input_metric: M,
    bounds: (TA, TA),
) -> Fallible<
    Transformation<
        VectorDomain<AtomDomain<TA>>,
        VectorDomain<AtomDomain<TA>>,
        M,
        M,
    >,
>
where
    TA: Float + SampleUniform,
    M: DatasetMetric,
    (VectorDomain<AtomDomain<TA>>, M): MetricSpace,
{
    let (lower, upper) = bounds;
    if lower.is_nan() {
        return fallible!(MakeTransformation, "lower may not be nan");
    }
    if upper.is_nan() {
        return fallible!(MakeTransformation, "upper may not be nan");
    }
    if lower > upper {
        return fallible!(MakeTransformation, "lower may not be greater than upper");
    }
    let scale = upper - lower;

    make_row_by_row_fallible(
        input_domain,
        input_metric,
        AtomDomain::default(),
        move |v| {
            if v.is_null() {
                TA::sample_standard_uniform(false).map(|v| v * scale + lower)
            } else {
                Ok(*v)
            }
        },
    )
}

/// Utility trait to impute with a constant, regardless of the representation of nullity.
pub trait ImputeConstantDomain: Domain {
    /// This is the type of `Self::Carrier` after imputation.
    ///
    /// On any type `D` for which the `ImputeConstantDomain` trait is implemented,
    /// the syntax `D::Imputed` refers to this associated type.
    /// For example, consider `D` to be `OptionDomain<T>`, the domain of all `Option<T>`.
    /// The implementation of this trait for `OptionDomain<T>` designates that `type Imputed = T`.
    /// Thus `OptionDomain<T>::Imputed` is `T`.
    ///
    /// # Proof Definition
    /// `Self::Imputed` can represent the set of possible output values after imputation.
    type Imputed;

    /// A function that replaces a potentially-null carrier type with a non-null imputed type.
    ///
    /// # Proof Definition
    /// For any setting of the input parameters, where `constant` is non-null,
    /// the function returns a non-null value.
    fn impute_constant<'a>(
        default: &'a Self::Carrier,
        constant: &'a Self::Imputed,
    ) -> &'a Self::Imputed;
}
// how to impute, when null represented as `Option<T>`
impl<T: CheckAtom> ImputeConstantDomain for OptionDomain<AtomDomain<T>> {
    type Imputed = T;
    fn impute_constant<'a>(
        default: &'a Self::Carrier,
        constant: &'a Self::Imputed,
    ) -> &'a Self::Imputed {
        default.as_ref().unwrap_or(constant)
    }
}
// how to impute, when null represented as T with internal nullity
impl<T: CheckAtom + InherentNull> ImputeConstantDomain for AtomDomain<T> {
    type Imputed = Self::Carrier;
    fn impute_constant<'a>(
        default: &'a Self::Carrier,
        constant: &'a Self::Imputed,
    ) -> &'a Self::Imputed {
        if default.is_null() {
            constant
        } else {
            default
        }
    }
}

#[bootstrap(
    features("contrib"),
    arguments(
        constant(
            rust_type = "$get_atom(get_type(input_domain))",
            c_type = "AnyObject *"
        )
    ),
    generics(DIA(suppress), M(suppress))
)]
/// Make a Transformation that replaces null/None data with `constant`.
///
/// If chaining after a `make_cast`, the input type is `Option<Vec<TA>>`.
/// If chaining after a `make_cast_inherent`, the input type is `Vec<TA>`, where `TA` may take on float NaNs.
///
/// | input_domain                                    |  Input Data Type  |
/// | ----------------------------------------------- | ----------------- |
/// | `vector_domain(option_domain(atom_domain(TA)))` | `Vec<Option<TA>>` |
/// | `vector_domain(atom_domain(TA))`                | `Vec<TA>`         |
///
/// # Arguments
/// * `input_domain` - Domain of the input data. See table above.
/// * `input_metric` - Metric of the input data. A dataset metric.
/// * `constant` - Value to replace nulls with.
///
/// # Generics
/// * `DIA` - Atomic Input Domain of data being imputed.
/// * `M` - Dataset Metric.
pub fn make_impute_constant<DIA, M>(
    input_domain: VectorDomain<DIA>,
    input_metric: M,
    constant: DIA::Imputed,
) -> Fallible<Transformation<VectorDomain<DIA>, VectorDomain<AtomDomain<DIA::Imputed>>, M, M>>
where
    DIA: ImputeConstantDomain + Default,
    DIA::Imputed: 'static + Clone + CheckAtom,
    DIA::Carrier: 'static,
    M: DatasetMetric,
    (VectorDomain<DIA>, M): MetricSpace,
    (VectorDomain<AtomDomain<DIA::Imputed>>, M): MetricSpace
{
    let output_atom_domain = AtomDomain::default();
    if !output_atom_domain.member(&constant)? {
        return fallible!(MakeTransformation, "Constant may not be null.");
    }

    make_row_by_row(
        input_domain,
        input_metric,
        output_atom_domain,
        move |v| DIA::impute_constant(v, &constant).clone(),
    )
}

/// Utility trait to drop null values from a dataset, regardless of the representation of nullity.
pub trait DropNullDomain: Domain {
    /// This is the type of `Self::Carrier` after dropping null.
    ///
    /// On any type `D` for which the `DropNullDomain` trait is implemented,
    /// the syntax `D::Imputed` refers to this associated type.
    /// For example, consider `D` to be `OptionDomain<T>`, the domain of all `Option<T>`.
    /// The implementation of this trait for `DropNullDomain<T>` designates that `type Imputed = T`.
    /// Thus `DropNullDomain<T>::Imputed` is `T`.
    type Imputed;

    /// Standardizes `D::Carrier` into an `Option<D::Imputed>`, where `D::Imputed` is never null.
    ///
    /// `Self::Imputed` may have the capacity to represent null (like `f64`),
    /// but implementations of this function must guarantee that `Self::Imputed` is never null.
    fn option(value: &Self::Carrier) -> Option<Self::Imputed>;
}

/// how to standardize into an option, when null represented as `Option<T>`
impl<T: CheckAtom + Clone> DropNullDomain for OptionDomain<AtomDomain<T>> {
    type Imputed = T;
    fn option(value: &Self::Carrier) -> Option<T> {
        if value.is_null() {
            None
        } else {
            value.clone()
        }
    }
}
/// how to standardize into an option, when null represented as T with internal nullity
impl<T: CheckAtom + InherentNull + Clone> DropNullDomain for AtomDomain<T> {
    type Imputed = T;
    fn option(value: &Self::Carrier) -> Option<T> {
        if value.is_null() {
            None
        } else {
            Some(value.clone())
        }
    }
}

#[bootstrap(features("contrib"),
    generics(DIA(suppress), M(suppress)))]
/// Make a Transformation that drops null values.
///
///
/// | input_domain                                    |
/// | ----------------------------------------------- |
/// | `vector_domain(option_domain(atom_domain(TA)))` |
/// | `vector_domain(atom_domain(TA))`                |
///
/// # Generics
/// * `M` - Dataset Metric.
/// * `DIA` - atomic domain of input data that contains nulls.
pub fn make_drop_null<M, DIA>(
    input_domain: VectorDomain<DIA>,
    input_metric: M,
) -> Fallible<
    Transformation<
        VectorDomain<DIA>,
        VectorDomain<AtomDomain<DIA::Imputed>>,
        M,
        M,
    >,
>
where
    DIA: DropNullDomain + Default,
    DIA::Imputed: CheckAtom,
    M: DatasetMetric,
    (VectorDomain<DIA>, M): MetricSpace,
    (VectorDomain<AtomDomain<DIA::Imputed>>, M): MetricSpace,
{
    Transformation::new(
        input_domain,
        VectorDomain::new(AtomDomain::default()),
        Function::new(|arg: &Vec<DIA::Carrier>| arg.iter().filter_map(DIA::option).collect()),
        input_metric.clone(),
        input_metric,
        StabilityMap::new_from_constant(1),
    )
}

#[cfg(test)]
mod tests {
    use crate::metrics::SymmetricDistance;

    use super::*;

    #[test]
    fn test_impute_uniform() -> Fallible<()> {
        let imputer = make_impute_uniform_float(
            VectorDomain::new(AtomDomain::default()),
            SymmetricDistance::default(),
            (2.0, 2.0))?;

        let result = imputer.invoke(&vec![1.0, f64::NAN])?;

        assert_eq!(result, vec![1., 2.]);
        assert!(imputer.check(&1, &1)?);
        Ok(())
    }

    #[test]
    fn test_impute_constant_option() -> Fallible<()> {
        let imputer = make_impute_constant(
            VectorDomain::new(OptionDomain::new(AtomDomain::default())),
            SymmetricDistance::default(),
            "IMPUTED".to_string(),
        )?;

        let result = imputer.invoke(&vec![Some("A".to_string()), None])?;

        assert_eq!(result, vec!["A".to_string(), "IMPUTED".to_string()]);
        assert!(imputer.check(&1, &1)?);
        Ok(())
    }

    #[test]
    fn test_impute_constant_inherent() -> Fallible<()> {
        let imputer = make_impute_constant(
            VectorDomain::new(AtomDomain::new_nullable()), 
            SymmetricDistance::default(),
            12.)?;

        let result = imputer.invoke(&vec![f64::NAN, 23.])?;

        assert_eq!(result, vec![12., 23.]);
        assert!(imputer.check(&1, &1)?);
        Ok(())
    }

    #[test]
    fn test_impute_drop_option() -> Fallible<()> {
        let imputer = make_drop_null(VectorDomain::new(OptionDomain::default()), SymmetricDistance::default())?;

        let result = imputer.invoke(&vec![Some(f64::NAN), Some(23.), None])?;

        assert_eq!(result, vec![23.]);
        assert!(imputer.check(&1, &1)?);
        Ok(())
    }
    #[test]
    fn test_impute_drop_inherent() -> Fallible<()> {
        let imputer = make_drop_null(VectorDomain::new(AtomDomain::new_nullable()), SymmetricDistance::default())?;

        let result = imputer.invoke(&vec![f64::NAN, 23.])?;

        assert_eq!(result, vec![23.]);
        assert!(imputer.check(&1, &1)?);
        Ok(())
    }
}
