turbo 0.4.2

use std::collections::HashMap;

use Code;
use daggy::{Dag, EdgeIndex, NodeIndex, Walker};
use daggy::petgraph::{EdgeDirection};
use rand::random;

use {b1, b3};

pub trait MLDecoder {
  /// Runs the 'forward pass' step of the Viterbi decoder. Returns the generated
  /// state tree and a vector of indices into the state tree corresponding to
  /// the surviving leaf states.
  fn estimate_states(&self, signal: &Vec<b1>) -> (Dag<DecoderState, usize>,
                                                  Vec<NodeIndex>);

  /// Reads back the maximum likelihood bits in reverse order from the pruned
  /// tree and the surviving leaf.
  fn read_tree(&self, pruned_tree: (Dag<DecoderState, usize>,
                                    Vec<NodeIndex>))
               -> Vec<b1>;

  /// Runs the forward and backward steps of the decoder to produce the fully
  /// decoded string.
  fn decode(&self, signal: &Vec<b1>) -> Vec<b1>;
}

pub trait MAPDecoder {
  fn estimate_states(signal: Vec<b1>, code: Code) -> (Dag<DecoderState, usize>,
                                                      Vec<NodeIndex>);
  fn decode(signal: Vec<b1>, code: Code) -> Vec<b1>;
}

#[derive(Clone, Copy, Debug)]
pub struct DecoderState {
  /// Index of the decoder state
  pub state_index: usize,
  /// (Hamming) distance from the observed signal
  pub dist: usize
}

impl MLDecoder for Code {
  /// Hard decision decoding, using the Hamming distance.
  // TODO write some tests here
  fn estimate_states(&self, signal: &Vec<b1>) -> (Dag<DecoderState, usize>,
                                                  Vec<NodeIndex>) {
    // TODO some check here for required minimum signal length

    // We only need to consider a certain number of samples before storing the
    // metrics and pruning down to the code dictionary set.
    let chunk_width = self.polys.len() + 1;
    
    // TODO make (lazy) static
    let start_index: usize = self.start_state[0].into();

    let mut decoder_tree = Dag::new();
    let root = decoder_tree.add_node(
      DecoderState { state_index: start_index, dist: 0 });
    
    let mut previous_nodes = vec![root];

    // Three iterators need to be at work here:
    //
    // 1. An iterator over the signal elements
    // 2. An iterator over the previous candidate states
    // 3. An iterator over the possible next states for each candidate
    
    // We can get (1) just by iterating. (2) has to be gotten by retaining a
    // reference to the computed candidates as they are appended to the
    // state. (3) can be computed from the node under iteration.

    // Dictionary of pruned states, used once every chunk_width.
    // Stores the DAG node index along with the decoder state for later use.
    let mut pruned_states: HashMap<usize, (NodeIndex, DecoderState)> =
      HashMap::new();
    
    for next_bit_chunk in signal.chunks(chunk_width) {

      // TODO make variable word length
      for next_bit in next_bit_chunk {

        let mut descendant_nodes = vec![];

        for previous_node in previous_nodes {

          // Compute the distances for all of the reachable states
          let next_state_inds =
            self.next_states(decoder_tree[previous_node].state_index);

          for (iter_ind, state_ind) in next_state_inds.iter().enumerate() {
            // TODO remove the binary magic number/redundant computation here
            let next_bit_guess: usize = iter_ind % (next_bit.max() + 1);

            // TODO use parameterized distance metric here
            let dist = next_bit_guess ^ usize::from(*next_bit);
            let prev_dist = decoder_tree[previous_node].dist;
            
            descendant_nodes.push(
              decoder_tree.add_child(
                previous_node,
                // storing local distance on the edges also
                dist,
                DecoderState {
                  state_index: *state_ind,
                  // accumulated distance
                  dist: prev_dist + dist
                }
              ).1);
          }
        }

        previous_nodes = descendant_nodes;
      }

      // Once we've finished a chunk, we know that we're at a point where we
      // have arrived at all of the possible states. At this point, we want to
      // iterate over the previous nodes and store only the nodes that have the
      // smallest accumulated distance metric along the paths they represent.

      // We clear the map since the distances are stored accumulated, so we
      // don't need to know immediately about the past.
      pruned_states.clear();

      // This means updating the pruned_states hash map...
      for (node_ind, survivor_node) in previous_nodes
        .iter()
        .map(|&node_ind| (node_ind, decoder_tree[node_ind])) {
          if pruned_states.contains_key(&survivor_node.state_index) {
            let cur_node = pruned_states.get_mut(&survivor_node.state_index)
                                        .unwrap();

            if cur_node.1.dist > survivor_node.dist {
              *cur_node = (node_ind, survivor_node);
            }
          } else {
            pruned_states.insert(survivor_node.state_index,
                                 (node_ind, survivor_node));
          }
        }

      // ...and rewriting the previous_states vec so that the next loop starts
      // from that new state.
      previous_nodes = pruned_states.keys().filter_map(|key| {
        if let Some(state_tuple) = pruned_states.get(key) {
          Some(state_tuple.0)
        } else {
          None
        }
      }).collect();

      prune_state_tree(&mut decoder_tree, &previous_nodes);      
      
      // NOTE Technically, the 1..(n-1)th chunks should use all of the possible
      // states as start states, but there is always the chance of the signal
      // being cut off.
    }

    // Return the generated tree and the node indices of the survivor states.
    (decoder_tree, previous_nodes)
  }

  
  fn read_tree(&self, pruned_tree: (Dag<DecoderState, usize>,
                                    Vec<NodeIndex>))
               -> Vec<b1> {
    
    // The list of survivors, if we've waited long enough, should only have one
    // node index in it. However, there's always the chance that two branches
    // could both have the same metric and therefore arrive. So the first thing
    // we have to do is choose a random one (under uniform prior assumption).
    let survivor_node = pruned_tree.1[random::<usize>() % pruned_tree.1.len()];
    let tree = pruned_tree.0;

    tree.recursive_walk(survivor_node, |tree, node_ind| {
      // Each node should only have one parent.
      tree.parents(node_ind).iter(tree).nth(1)
    }).iter(&tree)
      .map(|(_, node_ind)| {
        // Will panic if state_index is out of bounds, but that's expected to be
        // enforced by the Code struct
        let code_word: b3 = tree[node_ind].state_index.into();
        code_word.bits().to_vec()
      })
      .flat_map(|x| x)
      .collect()
  }

  fn decode(&self, signal: &Vec<b1>) -> Vec<b1> {
    // TODO add number of bits argument
    self.read_tree(self.estimate_states(signal))
        .iter().rev().cloned().collect()
  }
}


// impl MAPDecoder for Code {}

// TODO make a trait out of this and make the implementation more
// efficient/clean?
// fn hamming_dist(word_1: b3, word_2: b3) -> usize {
//   (word_1 ^ word_2)
//     .bits()
//     .iter()
//     .fold(0, |accum, &bit| (accum + usize::from(bit)))
// }

// fn pluck<'a, T>(inds: Vec<usize>, target: T) -> T
//   where <T as Index<usize>>::Output: Sized + Clone,
//         T: Index<usize> + IntoIterator +
//            iter::FromIterator<<T as Index<usize>>::Output> {
//   inds.iter().map(|&ind| target[ind].clone()).collect()
// }

// Prunes the state tree so that only lineages leading to survivor nodes are
// kept.
fn prune_state_tree(tree: &mut Dag<DecoderState, usize>,
                    survivor_nodes: &Vec<NodeIndex>) {
  // Notes on this horrible implementation follow.
  //
  // Algorithm I am thinking of using to prune:
  //
  // 1. Get the leaf node indices disjoint from the survivor ones
  //
  // 2. For each leaf, remove the leaf and the edge leading from it to its
  // immediate parent.
  //
  // 3. If the parent has a survivor node in its descendants, keep it (and
  // memoize the index?).
  //
  // 4. If the parent does not have a survivor node, recurse with parent as
  // leaf.
  
  // NOTE the Dag structure is not stable under removals, so we also need an
  // iterator that can reliably reach the leaf nodes after pruning. (If
  // nothing else works filtering on tree.raw_nodes() for .next[0] ==
  // EdgeIndex(End) works.
  
  let end_index = EdgeIndex::<u32>::end();

  // Gets the leaf indices that aren't in the survivor node indices
  let dead_leaf_inds: Vec<NodeIndex> = tree.raw_nodes().iter().filter(|&node| {
    node.next_edge(EdgeDirection::Outgoing) == end_index
  }).filter_map(|leaf| {
    // Complex rigmarole to get a connected node's own ID
    let twig = leaf.next_edge(EdgeDirection::Incoming);
    let node_id = tree.edge_endpoints(twig).unwrap().1;

    if !survivor_nodes.contains(&node_id) {
      Some(node_id)
    } else {
      None
    }      
  }).collect();

  let mut good_parents: Vec<NodeIndex> = vec![];
  let mut nodes_to_remove: Vec<NodeIndex> = vec![];
  let mut edges_to_remove: Vec<EdgeIndex> = vec![];

  for leaf_ind in dead_leaf_inds {
    tree.recursive_walk(leaf_ind, |tree, node_ind| {
      let mut walk_res = None;
      
      for (parent_edge, parent_ind) in tree.parents(node_ind).iter(tree) {
        if good_parents.contains(&parent_ind) { continue; }
        else if tree.children(parent_ind).any(tree, |_, _, child_ind| {
          survivor_nodes.contains(&child_ind)
        }) {
          good_parents.push(parent_ind);
        } else {
          nodes_to_remove.push(parent_ind);
          edges_to_remove.push(parent_edge);
          walk_res = Some((parent_edge, parent_ind));
        }
      }
      walk_res
    }).last(&tree);
    // HACK calling last() just so the recursive walker will run, since this
    // 'almost-iterator' situation is too gnarly for words
  }
  
  // TODO check this for index stability
  for node_ind in nodes_to_remove.iter() {
    tree.remove_node(*node_ind);
  }
  for edge_ind in edges_to_remove.iter() {
    tree.remove_edge(*edge_ind);
  }
}