Rust MCTS eval_and_expand

pub fn eval_board(
    bs: &BoardStack,
    net: &Net,
    tree: &mut Tree,
    selected_node_idx: &usize,
) -> Vec<usize> {
    let contents = get_contents();
    let b = convert_board(bs);

    let output = eval_state(b, &net).expect("Error");

    let (board_eval, policy) = output; // check policy, eval ordering!

    let board_eval = board_eval.squeeze();

    let board_eval: Vec<f32> = Vec::try_from(board_eval).expect("Error");

    let board_eval = Tensor::from_slice(&vec![board_eval[0]]);

    let value = Tensor::tanh(&board_eval);

    let policy = policy.squeeze();
    let policy: Vec<f32> = Vec::try_from(policy).expect("Error");
    let value = f32::try_from(value).expect("Error");

    let value = match bs.board().side_to_move() {
            Color::Black => -value,
            Color::White => value,
    };

    // step 1 - get the corresponding idx for legal moves

    let mut legal_moves: Vec<Move> = Vec::new();
    bs.board().generate_moves(|moves| {
        // Unpack dense move set into move list
        legal_moves.extend(moves);
        false
    });

    let mut fm: Vec<Move> = Vec::new();
    if bs.board().side_to_move() == Color::Black {
        // flip move
        for mv in &legal_moves {
            fm.push(Move {
                from: mv.from.flip_rank(),
                to: mv.to.flip_rank(),
                promotion: mv.promotion,
            })
        }
    } else {
        fm = legal_moves.clone();
    }

    legal_moves = fm;

    let mut idx_li: Vec<usize> = Vec::new();

    for mov in &legal_moves {
        // let mov = format!("{}", mov);
        if let Some(idx) = contents.iter().position(|x| mov == x) {
            idx_li.push(idx as usize);
        }
    }

    // step 2 - using the idx in step 1, index all the policies involved
    let mut pol_list: Vec<f32> = Vec::new();
    for id in &idx_li {
        pol_list.push(policy[*id]);
    }

    // println!("{:?}", pol_list);

    // step 3 - softmax

    let sm = Tensor::from_slice(&pol_list);

    let sm = Tensor::softmax(&sm, 0, Kind::Float);

    let pol_list: Vec<f32> = Vec::try_from(sm).expect("Error");

    // println!("{:?}", pol_list);

    // println!("        V={}", &value);

    // step 4 - iteratively append nodes into class
    let mut counter = 0;
    let ct = tree.nodes.len();
    for (mv, pol) in legal_moves.iter().zip(pol_list.iter()) {
        tree.nodes[*selected_node_idx].eval_score = value;
        // tree.nodes[*selected_node_idx].eval_score = 0.0;
        let fm: Move;
        if bs.board().side_to_move() == Color::Black {
            // flip move
            fm = Move {
                from: mv.from.flip_rank(),
                to: mv.to.flip_rank(),
                promotion: mv.promotion,
            };
        } else {
            fm = *mv;
        }
        let mut child = Node::new(0.0, Some(*selected_node_idx), Some(fm));
        // let mut child = Node::new(*pol, Some(*selected_node_idx), Some(fm));
        // println!("{:?}, {:?}, {:?}", mv, child.policy, child.eval_score);
        tree.nodes.push(child); // push child to the tree Vec<Node>
        tree.nodes[*selected_node_idx].children.push(counter + ct); // push numbers
        counter += 1
    }
    // println!("{:?}", tree.nodes.len());
    idx_li
}