// Copyright (c) 2021-2021 Thomas Kramer.
// SPDX-FileCopyrightText: 2022 Thomas Kramer <code@tkramer.ch>
//
// SPDX-License-Identifier: AGPL-3.0-or-later

//! Helper functions to interpret the liberty library which comes from the parser.

use std::ops::{Add, Div, Mul, Sub};

use crate::interp1d::Interp1D;
use crate::interp2d::Interp2D;
use itertools::Itertools;
use liberty_io::Group;
use ndarray::OwnedRepr;

/// Find timing groups within a `pin` group by the name of the related pin and the timing type.
pub fn select_timing_groups<'a>(
    pin_group: &'a Group,
    related_pin: &str,
    timing_type: &str,
) -> Vec<&'a Group> {
    assert_eq!(pin_group.name, "pin", "Must be a `pin` group.");

    let timing_groups: Vec<_> = pin_group
        .find_groups_by_name("timing")
        .filter(|g| {
            g.get_simple_attribute("related_pin")
                .and_then(|v| v.as_str())
                == Some(related_pin)
        })
        .collect();

    if timing_groups.is_empty() {
        // Print the available related pin names.
        let related_pins = pin_group
            .find_groups_by_name("timing")
            .flat_map(|g| {
                g.get_simple_attribute("related_pin")
                    .and_then(|a| a.as_str())
            })
            .unique()
            .sorted();
        log::warn!(
            "No timing group found. Related pin name must be one of: {}",
            related_pins.into_iter().join(", ")
        );
    }

    // Filter by timing type.
    let timing_groups: Vec<_> = timing_groups
        .into_iter()
        .filter(|g| {
            g.get_simple_attribute("timing_type")
                .and_then(|v| v.as_str())
                == Some(timing_type)
        })
        .collect();

    timing_groups
}

/// Describes errors that happen when accessing a liberty library.
#[derive(Debug, Clone)]
pub enum LibertyErr {
    /// An attribute could not be found.
    AttributeNotFound(String),
    /// A table was malformed: Not a float array or wrong sizes of rows for 2D arrays.
    TableMalformed,
    /// Table with this name was not found.
    TableNotFound(String),
    /// Something is wrong with a lookup table template.
    InvalidLuTableTemplate(String),
    /// The specified delay model is not supported or no delay model is specified.
    UnsupportedDelayModel(String),
    /// The library does not define a required unit (time, capacitance, ...).
    UnitNotDefined(&'static str),
    /// Unspecified error.
    Other(String),
}

impl std::fmt::Display for LibertyErr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            LibertyErr::AttributeNotFound(attr) => write!(f, "Attribute not found: {}", attr),
            LibertyErr::TableNotFound(name) => write!(f, "Table not found: {}", name),
            LibertyErr::TableMalformed => write!(f, "Table malformed."),
            LibertyErr::InvalidLuTableTemplate(name) => {
                writeln!(f, "Lookup table template is malformed: {}", name)
            }
            LibertyErr::UnsupportedDelayModel(model_name) => write!(
                f,
                "The specified delay model '{}' is not supported or no delay model is specified.",
                model_name
            ),
            LibertyErr::Other(msg) => write!(f, "Error while preparing timing library: {}", msg),
            LibertyErr::UnitNotDefined(unit_name) => {
                write!(f, "Unit is not defined: {}", unit_name)
            }
        }
    }
}

/// Interpolated timing table with two, one or zero variables.
#[derive(Clone, Debug)]
pub enum Interp<Z = f64, X = f64, Y = X> {
    /// A scalar (constant) value.
    Scalar(Z),
    /// A one-dimensional table. Value depends only on the first variable.
    Interp1D1(Interp1D<X, Z>),
    /// A one-dimensional table. Value depends only on the second variable.
    Interp1D2(Interp1D<Y, Z>),
    /// A two-dimensional table
    Interp2D(Interp2D<X, Y, Z, OwnedRepr<Z>>),
}

impl<Z, X, Y> Interp<Z, X, Y> {
    /// Evaluate the interpolated function at `(x, y)`.
    /// For 1D functions `y` is ignored.
    /// For scalars both `x` and `y` is ignored.
    pub fn eval2d(&self, (x, y): (X, Y)) -> Z
    //where
    //    f64: Mul<Z, Output = Z>,
    //    Z: Num + Copy + Mul<f64, Output = Z>,
    where
        X: Copy + Sub<Output = X> + Div + PartialOrd,
        Y: Copy + Sub<Output = Y> + Div<Output = <X as Div>::Output> + PartialOrd,
        Z: Copy + Mul<<X as Div>::Output, Output = Z> + Add<Output = Z> + Sub<Output = Z>,
        <X as Div>::Output:
            Copy + Add<Output = <X as Div>::Output> + Sub<Output = <X as Div>::Output>,
    {
        match self {
            Self::Scalar(s) => *s,
            Self::Interp1D1(i) => i.eval(x),
            Self::Interp1D2(i) => i.eval(y),
            Self::Interp2D(i) => i.eval((x, y)),
        }
    }

    /// Swap the variables
    pub fn swap_variables(self) -> Interp<Z, Y, X>
    where
        Z: Clone,
    {
        match self {
            Interp::Scalar(s) => Interp::Scalar(s),
            Interp::Interp1D1(i) => Interp::Interp1D2(i),
            Interp::Interp1D2(i) => Interp::Interp1D1(i),
            Interp::Interp2D(i) => Interp::Interp2D(i.swap_variables()),
        }
    }

    pub fn map_values<Z2>(&self, f: impl Fn(&Z) -> Z2) -> Interp<Z2, X, Y>
    where
        X: PartialOrd + Clone,
        Y: PartialOrd + Clone,
    {
        match self {
            Interp::Scalar(s) => Interp::Scalar(f(s)),
            Interp::Interp1D1(i) => Interp::Interp1D1(i.map_values(f)),
            Interp::Interp1D2(i) => Interp::Interp1D2(i.map_values(f)),
            Interp::Interp2D(i) => Interp::Interp2D(i.map_values(f)),
        }
    }
    pub fn map_x_axis<Xnew>(self, f: impl Fn(X) -> Xnew) -> Interp<Z, Xnew, Y>
    where
        Xnew: PartialOrd,
    {
        match self {
            Interp::Scalar(s) => Interp::Scalar(s),
            Interp::Interp1D1(i) => Interp::Interp1D1(i.map_axis(f)),
            Interp::Interp1D2(i) => Interp::Interp1D2(i),
            Interp::Interp2D(i) => Interp::Interp2D(i.map_x_axis(f)),
        }
    }

    pub fn map_y_axis<Ynew>(self, f: impl Fn(Y) -> Ynew) -> Interp<Z, X, Ynew>
    where
        Ynew: PartialOrd,
    {
        match self {
            Interp::Scalar(s) => Interp::Scalar(s),
            Interp::Interp1D1(i) => Interp::Interp1D1(i),
            Interp::Interp1D2(i) => Interp::Interp1D2(i.map_axis(f)),
            Interp::Interp2D(i) => Interp::Interp2D(i.map_y_axis(f)),
        }
    }
}

/// Read a two dimensional table with indices and create an `Interp2D` struct out of them.
/// The `timing_group` must be an actual `timing` liberty group. Otherwise this function panics.
/// `table_name` is the liberty group name of the table (`cell_rise`, `rise_transition`, ...).
///
/// The indices must be attributes with names `index_1_name` and `index_2_name`, the values must be an attribute
/// with name `values_name`.
pub fn get_timing_table(
    timing_group: &Group,
    table_name: &str,
    index_1_name: &str,
    index_2_name: &str,
    values_name: &str,
) -> Result<Interp<f64>, LibertyErr> {
    assert_eq!(timing_group.name, "timing", "Must be a `timing` group.");

    if let Some(table_group) = timing_group.find_groups_by_name(table_name).next() {
        let index1 = table_group.get_simple_attribute(index_1_name);
        let index2 = table_group.get_simple_attribute(index_2_name);
        let values = table_group.get_attribute(values_name).ok_or_else(|| {
            log::error!("`{}` not found.", values_name);
            LibertyErr::AttributeNotFound(values_name.to_string())
        })?;

        // Try to convert the attributes to float arrays.

        let index1 = index1
            .and_then(|v| v.as_str())
            .map(liberty_io::util::parse_float_array)
            .map_or(Ok(None), |v| v.map(Some))
            .map_err(|e| {
                log::error!("Index `{}` is not well formatted: {}", index_1_name, e);
                LibertyErr::TableMalformed
            })?;

        let index2 = index2
            .and_then(|v| v.as_str())
            .map(liberty_io::util::parse_float_array)
            .map_or(Ok(None), |v| v.map(Some))
            .map_err(|e| {
                log::error!("Index `{}` is not well formatted: {}", index_2_name, e);
                LibertyErr::TableMalformed
            })?;

        let maybe_values: Result<Vec<Vec<f64>>, _> = values
            .iter()
            .map(|v| {
                v.as_str()
                    .map(liberty_io::util::parse_float_array)
                    .ok_or_else(|| {
                        log::error!("Values `{}` is not well formatted.", values_name);
                        LibertyErr::TableMalformed
                    })?
                    .map_err(|e| {
                        log::error!("Values `{}` is not well formatted: {}", values_name, e);
                        LibertyErr::TableMalformed
                    })
            })
            .collect();
        let values = maybe_values?;

        // Check if 2D array is well formed.
        if !values.iter().map(|v| v.len()).all_equal() {
            log::error!("Each row in table must have the same amount of elements.");
            return Err(LibertyErr::TableMalformed);
        }

        // There can be two dimensional tables, one dimensional tables
        // and scalar values.
        // Handle the correct case:
        match (index1, index2) {
            (Some(index1), Some(index2)) => {
                // 2D
                if index1.len() != values.len() {
                    log::error!(
                        "Table dimension does not match length of index {}.",
                        index_1_name
                    );
                    Err(LibertyErr::TableMalformed)?;
                }
                if values
                    .first()
                    .map(|v| v.len() != index2.len())
                    .unwrap_or(false)
                {
                    log::error!(
                        "Table dimension does not match length of index {}.",
                        index_2_name
                    );
                    Err(LibertyErr::TableMalformed)?;
                }

                // Convert the table into a 2D array.
                let mut arr = ndarray::Array2::zeros((index1.len(), index2.len()));
                for (i, vs) in values.iter().enumerate() {
                    for (j, v) in vs.iter().enumerate() {
                        arr[[i, j]] = *v;
                    }
                }

                // Wrap into interpolator.
                let interp = Interp2D::new(index1, index2, arr);

                Ok(Interp::Interp2D(interp))
            }
            (Some(index1), None) => {
                // 1D. Only the first index is present.
                if values.len() != 1 {
                    log::error!("Values must be one dimensional for one-dimensional tables.");
                    Err(LibertyErr::TableMalformed)?;
                }
                let interp = Interp1D::new(index1, values[0].clone());
                Ok(Interp::Interp1D1(interp))
            }
            (None, Some(_)) => {
                log::error!(
                    "Table has second index defined ({}) but not first.",
                    index_2_name
                );
                Err(LibertyErr::TableMalformed)
            }
            (None, None) => {
                // Scalar
                let scalar_value =
                    values
                        .get(0)
                        .and_then(|v| v.first())
                        .copied()
                        .ok_or_else(|| {
                            log::error!("Scalar table must contain exactly one value.");
                            LibertyErr::TableMalformed
                        })?;
                Ok(Interp::Scalar(scalar_value))
            }
        }
    } else {
        log::warn!("No such table found: {}", table_name);
        Err(LibertyErr::TableNotFound(table_name.to_string()))
    }
}

#[test]
fn test_get_timing_table_scalar() {
    // Read a timing group with and then get the constant scalar.

    use liberty_io::read_liberty_chars;

    let data = r#"
    timing () {
        cell_rise(scalar) {
            values ("42.0")
        }
    }
    "#;

    let timing_group = read_liberty_chars(data.chars()).unwrap();

    let timing_table =
        get_timing_table(&timing_group, "cell_rise", "index_1", "index_2", "values").unwrap();

    assert!((timing_table.eval2d((0.0, 0.0)) - 42.0) < 1e-6);
    assert!((timing_table.eval2d((7.0, 123.0)) - 42.0) < 1e-6);
}

#[test]
fn test_get_timing_table_1d() {
    // Read a timing group and then get the one dimensional interpolated table.

    use liberty_io::read_liberty_chars;

    let data = r#"

    timing () {
        cell_rise() {
            index_1 ("0.0, 1.0, 2.0");
            values ("0.0, 1.0, 0.0")
        }
    }

    "#;

    let timing_group = read_liberty_chars(data.chars()).unwrap();

    let timing_table =
        get_timing_table(&timing_group, "cell_rise", "index_1", "index_2", "values").unwrap();

    assert!((timing_table.eval2d((0.0, 0.0)) - 0.0) < 1e-6);
    assert!((timing_table.eval2d((1.0, 0.0)) - 1.0) < 1e-6);
    assert!((timing_table.eval2d((2.0, 0.0)) - 0.0) < 1e-6);
}

#[test]
fn test_get_timing_table_2d() {
    // Read a timing group and then get the two-dimensional interpolated table.

    use liberty_io::read_liberty_chars;

    let data = r#"

    timing () {
        cell_rise() {
            index_1 ("0.0, 1.0");
            index_2 ("2.0, 3.0, 4.0");
            values ("0.0, 1.0, 0.0", \
            "1.0, 0.0, 1.0")
        }
    }

    "#;

    let timing_group = read_liberty_chars(data.chars()).unwrap();

    let timing_table =
        get_timing_table(&timing_group, "cell_rise", "index_1", "index_2", "values").unwrap();

    assert!((timing_table.eval2d((0.0, 2.0)) - 0.0) < 1e-6);
    assert!((timing_table.eval2d((0.0, 3.0)) - 1.0) < 1e-6);
    assert!((timing_table.eval2d((0.0, 4.0)) - 0.0) < 1e-6);

    assert!((timing_table.eval2d((1.0, 2.0)) - 1.0) < 1e-6);
    assert!((timing_table.eval2d((1.0, 3.0)) - 0.0) < 1e-6);
}