use std::{collections::HashMap, iter::zip};

use crate::pyo3_prelude::*;
use crate::{
    cached_lambda,
    circuit::{Circuit, HashBytes},
};

use super::{CircuitNode, CircuitRc, NamedAxes};

pub fn set_named_axes<T: CircuitNode + Clone>(node: &T, named_axes: NamedAxes) -> T {
    assert!(named_axes.len() == node.info().rank());
    let mut result: T = node.clone();
    result.info_mut().named_axes = named_axes.clone();
    let mut result = result.init_info().unwrap();
    result.info_mut().named_axes = named_axes;
    result
}

#[pyfunction]
#[pyo3(name = "set_named_axes")]
pub fn set_named_axes_py(circuit: CircuitRc, named_axes: NamedAxes) -> CircuitRc {
    set_named_axes(&**circuit, named_axes).rc()
}
pub fn merge_named_axes(a: &NamedAxes, b: &NamedAxes) -> NamedAxes {
    assert!(a.len() == b.len());
    zip(a, b)
        .map(|(ai, bi)| bi.clone().or(ai.clone()))
        .collect()
}

pub fn named_axes_backward<T: CircuitNode + Clone>(
    circuit: &T,
    named_axes: &NamedAxes,
) -> Vec<NamedAxes> {
    let child_axis_map = circuit.child_axis_map();
    child_axis_map
        .iter()
        .map(|z| {
            z.iter()
                .map(|x| (*x).and_then(|i| named_axes[i].clone()))
                .collect()
        })
        .collect()
}

#[pyfunction]
pub fn propagate_named_axes(circuit: CircuitRc, named_axes: NamedAxes) -> CircuitRc {
    cached_lambda! {
        let mut recurse = |circ: Circuit,axes:NamedAxes| -> CircuitRc = {
            let new_out_axes = merge_named_axes(&circ.info().named_axes,&axes);
            let new_child_axes = named_axes_backward(&circ, &new_out_axes);
            let child_axis_names:HashMap<HashBytes,NamedAxes>=zip(circ.children(),new_child_axes).map(|(c,a)|(c.info().hash,a)).collect();

            set_named_axes(&circ,new_out_axes).map_children_unwrap(&mut |childy|recurse(childy.clone(),child_axis_names[&childy.info().hash].clone())).rc()
        };
    }
    recurse((**circuit).clone(), named_axes)
}
