pineappl_cli 1.4.0

use anyhow::{Result, bail};
use cxx::{UniquePtr, let_cxx_string};
use float_cmp::approx_eq;
use itertools::izip;
use lhapdf::Pdf;
use ndarray::{Axis, s};
use pineappl::boc::{Channel, Kinematics, Order};
use pineappl::grid::Grid;
use pineappl::interpolation::{Interp, InterpMeth, Map, ReweightMeth};
use pineappl::pids::PidBasis;
use pineappl::subgrid::{self, Subgrid};
use pineappl_applgrid::ffi::{self, grid};
use std::f64::consts::TAU;
use std::iter;
use std::path::Path;
use std::pin::Pin;

fn reconstruct_subgrid_params(grid: &Grid, order: usize, bin: usize) -> Result<Vec<Interp>> {
    if grid
        .kinematics()
        .iter()
        .filter(|kin| matches!(kin, Kinematics::Scale(_)))
        .count()
        > 1
    {
        bail!("APPLgrid does not support grids with more than one scale");
    }

    let mut mu2_grid: Vec<_> = grid
        .subgrids()
        .slice(s![order, bin, ..])
        .iter()
        .filter(|subgrid| !subgrid.is_empty())
        .flat_map(|subgrid| {
            grid.scales()
                .fac
                .calc(&subgrid.node_values(), grid.kinematics())
                .into_owned()
        })
        .collect();
    mu2_grid.dedup_by(subgrid::node_value_eq_ref_mut);

    // TODO: implement the general case
    let mut result = vec![
        Interp::new(
            2e-7,
            1.0,
            50,
            3,
            ReweightMeth::ApplGridX,
            Map::ApplGridF2,
            InterpMeth::Lagrange,
        ),
        Interp::new(
            2e-7,
            1.0,
            50,
            3,
            ReweightMeth::ApplGridX,
            Map::ApplGridF2,
            InterpMeth::Lagrange,
        ),
    ];

    if let &[fac] = mu2_grid.as_slice() {
        result.insert(
            0,
            Interp::new(
                fac,
                fac,
                1,
                0,
                ReweightMeth::NoReweight,
                Map::ApplGridH0,
                InterpMeth::Lagrange,
            ),
        );
    } else {
        result.insert(
            0,
            Interp::new(
                1e2,
                1e8,
                40,
                3,
                ReweightMeth::NoReweight,
                Map::ApplGridH0,
                InterpMeth::Lagrange,
            ),
        );
    }

    Ok(result)
}

pub fn convert_into_applgrid(
    grid: &mut Grid,
    output: &Path,
    discard_non_matching_values: bool,
) -> Result<(UniquePtr<grid>, Vec<bool>)> {
    let dim = grid.bwfl().dimensions();

    if dim > 1 {
        bail!(
            "grid has {dim} dimensions, but APPLgrid only supports one-dimensional distributions"
        );
    }

    // APPLgrid can only be used with one-dimensional consecutive bin limits
    if grid.bwfl().slices().len() != 1 {
        bail!("grid has non-consecutive bin limits, which APPLgrid does not support");
    }

    match grid.convolutions() {
        [_] => {}
        [a, b] => {
            if (a != b) && (a.cc() == *b) {
                // use charge conjugate to map hadron-anti-hadron grids to use the same single
                // convolution function
                let index = usize::from(a.pid() < 0);
                grid.charge_conjugate(index);
            }
        }
        _ => bail!("APPLgrid does not support grids with more than two convolutions"),
    }

    // APPLgrid only understands PDG PIDs
    grid.rotate_pid_basis(PidBasis::Pdg);

    let non_trivial_factors = grid
        .channels()
        .iter()
        .flat_map(Channel::entry)
        .any(|&(_, factor)| !approx_eq!(f64, factor, 1.0, ulps = 4));

    // APPLgrid doesn't support non-trivial factors
    if non_trivial_factors {
        // TODO: this isn't the most efficient strategy, since we also split up channels with
        // trivial factors
        grid.split_channels();
        grid.merge_channel_factors();
        grid.optimize();
    }

    let combinations: Vec<_> = iter::once(grid.channels().len().try_into().unwrap())
        .chain(
            grid.channels()
                .iter()
                .enumerate()
                .flat_map(|(index, entry)| {
                    [
                        index.try_into().unwrap(),
                        entry.entry().len().try_into().unwrap(),
                    ]
                    .into_iter()
                    .chain(entry.entry().iter().flat_map(|(pids, _)| {
                        pids.iter()
                            .copied()
                            .chain(iter::repeat(0))
                            .take(2)
                            .collect::<Vec<_>>()
                    }))
                }),
        )
        .collect();

    // `id` must end with '.config' for APPLgrid to know its type is `lumi_pdf`
    let id = format!(
        "{}.config",
        output
            .file_stem()
            // UNWRAP: because we write to that file in the end, there always must be a file name
            .unwrap()
            .to_string_lossy()
    );
    // this object is managed by APPLgrid internally
    ffi::make_lumi_pdf(&id, &combinations).into_raw();

    let limits: Vec<_> = grid
        .bwfl()
        .bins()
        .iter()
        .map(|bin| {
            // TODO: instead of `bin.limits()[0]` we should use `bin.fill_limits()`, but this
            // requires changing the normalization
            bin.limits()[0].0
        })
        .chain(Some(
            grid.bwfl()
                .bins()
                .last()
                // UNWRAP: every `grid` should have at least one bin
                .unwrap()
                .limits()[0]
                .1,
        ))
        .collect();

    let order_mask = Order::create_mask(grid.orders(), 3, 0, false);
    let orders_with_mask: Vec<_> = grid
        .orders()
        .iter()
        .cloned()
        .zip(order_mask.iter().copied())
        .collect();
    let lo_alphas = orders_with_mask
        .iter()
        .filter_map(|&(Order { alphas, .. }, keep)| keep.then_some(alphas))
        .min()
        .unwrap();
    let loops = orders_with_mask
        .iter()
        .filter_map(|&(Order { alphas, .. }, keep)| keep.then_some(alphas))
        .max()
        .unwrap()
        - lo_alphas;

    let mut applgrid =
        ffi::make_empty_grid(&limits, &id, lo_alphas.into(), loops.into(), "f2", "h0");

    // APPLgrid has either two or one convolution(s)
    let convolutions = grid.convolutions().len();

    for order in order_mask
        .iter()
        .enumerate()
        .filter_map(|(index, keep)| keep.then_some(index))
    {
        let appl_order = grid.orders()[order].alphas - lo_alphas;
        let factor = TAU.powi(grid.orders()[order].alphas.into());

        for (bin, subgrids) in grid
            .subgrids()
            .index_axis(Axis(0), order)
            .axis_iter(Axis(0))
            .enumerate()
        {
            let interps = reconstruct_subgrid_params(grid, order, bin)?;
            // TODO: support DIS case
            assert_eq!(interps.len(), 3);

            // TODO: make sure interps[1] is the same as interps[2]

            let mut igrid = ffi::make_igrid(
                interps[0].nodes().try_into().unwrap(),
                interps[0].min(),
                interps[0].max(),
                interps[0].order().try_into().unwrap(),
                interps[1].nodes().try_into().unwrap(),
                interps[1].min(),
                interps[1].max(),
                interps[1].order().try_into().unwrap(),
                match interps[1].map() {
                    Map::ApplGridF2 => "f2",
                    map @ Map::ApplGridH0 => panic!("export does not support {map:?}"),
                },
                match interps[0].map() {
                    Map::ApplGridH0 => "h0",
                    map @ Map::ApplGridF2 => panic!("export does not support {map:?}"),
                },
                grid.channels().len().try_into().unwrap(),
                grid.convolutions().len() == 1,
            );
            let appl_grids: Vec<Vec<_>> = vec![
                (0..igrid.Ntau()).map(|i| igrid.getQ2(i)).collect(),
                (0..igrid.Ny1()).map(|i| igrid.getx1(i)).collect(),
                (0..igrid.Ny2()).map(|i| igrid.getx2(i)).collect(),
            ];

            for (channel, subgrid) in subgrids
                .iter()
                .enumerate()
                .filter(|(_, subgrid)| !subgrid.is_empty())
            {
                let grids = vec![
                    grid.scales()
                        .fac
                        .calc(&subgrid.node_values(), grid.kinematics())
                        .into_owned(),
                    grid.kinematics()
                        .iter()
                        .zip(subgrid.node_values())
                        .find_map(|(kin, node_values)| {
                            matches!(kin, &Kinematics::X(idx) if idx == 0).then_some(node_values)
                        })
                        // TODO: convert this into an error
                        .unwrap(),
                    if convolutions == 2 {
                        grid.kinematics()
                            .iter()
                            .zip(subgrid.node_values())
                            .find_map(|(kin, node_values)| {
                                matches!(kin, &Kinematics::X(idx) if idx == 1)
                                    .then_some(node_values)
                            })
                            // TODO: convert this into an error
                            .unwrap()
                    } else {
                        Vec::new()
                    },
                ];

                let appl_idx: Vec<Vec<_>> = izip!(&grids, &appl_grids, ["factorization scale muf2", "momentum fraction x1", "momentum fraction x2"])
                    .map(|(grid, appl_grid, label)| {
                        grid
                        .iter()
                        .map(|&value| {
                            appl_grid
                                .iter()
                                .position(|&appl_value| subgrid::node_value_eq(appl_value, value))
                                .map_or_else(
                                    || {
                                        if discard_non_matching_values {
                                            Ok(-1)
                                        } else {
                                            bail!("{label} = {value} not found in APPLgrid; try exporting with `--discard-non-matching-values`")
                                        }
                                    },
                                    |idx| Ok(idx.try_into().unwrap()),
                                )
                        })
                        .collect::<Result<_>>()
                    }
                ).collect::<Result<_>>()?;

                let mut weightgrid = ffi::igrid_weightgrid(igrid.pin_mut(), channel);

                'looop: for (indices, value) in subgrid.indexed_iter() {
                    // TODO: here we assume that all X are consecutive starting from the second
                    // element and are in ascending order

                    let appl_indices = [
                        appl_idx[0][indices[0]],
                        appl_idx[1][indices[1]],
                        if convolutions == 2 {
                            appl_idx[2][indices[2]]
                        } else {
                            0
                        },
                    ];

                    for (&appl_index, grid, &index, label) in izip!(
                        &appl_indices,
                        &grids,
                        &indices,
                        ["scale muf2", "momentum fraction x1", "momentum fraction x2"]
                    ) {
                        if appl_index == -1 {
                            if value != 0.0 {
                                println!(
                                    "WARNING: discarding non-matching {label} = {} in subgrid {:?}",
                                    grid[index],
                                    (order, bin, channel)
                                );
                            }

                            continue 'looop;
                        }
                    }

                    ffi::sparse_matrix_set(
                        weightgrid.as_mut(),
                        appl_indices[0],
                        appl_indices[1],
                        appl_indices[2],
                        factor * value,
                    );
                }

                // TODO: is this call needed?
                weightgrid.trim();
            }

            igrid.pin_mut().setlimits();

            unsafe {
                applgrid.pin_mut().add_igrid(
                    bin.try_into().unwrap(),
                    appl_order.into(),
                    igrid.into_raw(),
                );
            }
        }
    }

    applgrid.pin_mut().include_photon(true);

    let_cxx_string!(filename = output.to_str().unwrap());
    let_cxx_string!(empty = "");

    applgrid.pin_mut().Write(&filename, &empty, &empty);

    Ok((applgrid, order_mask))
}

// TODO: deduplicate this function from import
pub fn convolve_applgrid(grid: Pin<&mut grid>, conv_funs: &mut [Pdf]) -> Vec<f64> {
    let nloops = grid.nloops();

    // TODO: add support for convolving an APPLgrid with two functions
    assert_eq!(conv_funs.len(), 1);

    pineappl_applgrid::grid_convolve_with_one(grid, &mut conv_funs[0], nloops, 1.0, 1.0, 1.0)
}