open Printf;;
  Log.set_level 0;;
open Log;;
  
let nonzero arr = 
  let zerofilter a b = if b = 0 then a else -1 in
  List.filter (fun a -> a > 0) (Array.to_list (Array.mapi zerofilter arr))
    
let rec print_list list = match list with 
  [] -> printf "\n";
| (a,b)::t -> printf "(%d,%d)" a b; print_list t;;
   
(*returns the first n elements of a list*)
let rec prefix n list = if n = 0 then [] else 
match list with
  [] -> []
| h::t -> h::(prefix (n-1) t);;
   
let randlist list = match list with
  [] -> []
| _ -> Array.to_list (
    let array = Array.of_list list in
    let len = Array.length array in
    for i=0 to len - 1 do
      let pos = Random.int len in
      let temp = array.(pos) in
      array.(pos) <- array.(i);
      array.(i) <- temp
    done;
    array)
      
(* generates a list with element frequency given in counts *)
let rec gen_freq_list counts i = match counts with 
  [] -> []
| 0 :: rest -> gen_freq_list rest (i+1)
| n :: rest -> i :: (gen_freq_list ((n-1)::rest) i)
		     
exception Broken_input of (string)
exception Hit_target
    
type check = int
type data = int
type block_id = Check of check | Data of data
type edge = data * check

let dlist_to_blist dlist = 
  let rec conv dlist blist = match dlist with 
    [] -> blist
  | x :: rest -> conv rest ((Data x)::blist) in
  conv dlist []

let clist_to_blist clist = 
  let rec conv clist blist = match clist with 
    [] -> blist
  | x :: rest -> conv rest ((Check x)::blist) in
  conv clist []


      
(* compares two edges using the lexicographical ordering *)
let compare edge1 edge2 = match edge1, edge2 with
  (x1,y1), (x2,y2) -> if x1 = x2 then y1-y2 else x1-x2
    
let repeat trials rand_fun = 
 let total = ref 0 in
 for i = 0 to trials - 1 do
   total := !total + rand_fun
 done;
 !total

class virtual graph = object (self : 'a)
(* INFORMATION RETRIEVAL *)
  method virtual is_empty : bool
  method virtual edges : int
  method virtual dnodes : data list
  method virtual cnodes : check list
  method dblocks = dlist_to_blist self#dnodes
  method cblocks = clist_to_blist self#cnodes
  method dcount = List.length self#dnodes
  method ccount = List.length self#cnodes
  method unbalanced = self#dcount > self#ccount
  method virtual exist : data -> check -> bool
  method virtual dnodes_of : check -> data list
  method virtual cnodes_of : data -> check list
  method dcount_of data = List.length (self#dnodes_of data)
  method ccount_of check = List.length (self#cnodes_of check)
  method private dump_uedges = 
    let checks d = List.map (fun c -> (d,c)) (self#cnodes_of d) in
    (List.flatten (List.map checks self#dnodes))
  method dump_edges = List.sort compare self#dump_uedges
  method virtual consistent : bool
      
(* GRAPH MODIFICATION *)
  method virtual add : data -> check -> unit
  method add_edge (d,c) = self#add d c
  method virtual del : data -> check -> unit
  method deldata data = List.iter (self#del data) (self#cnodes_of data)
  method delcheck check = 
    List.iter 
      (fun a -> self#del a check) 
      (self#dnodes_of check)
      
(* GRAPH CREATION FOR TESTING GRAPHS *)      
  method virtual newgraph : 'a
  method newdata datalist = 
    let edgel = List.filter (fun (d,c) -> List.mem d datalist) self#dump_uedges in
    self#generatepair edgel
  method newtest dataneed checkmissing = 
    let edgel = List.filter 
	(fun (d,c) -> (List.mem d dataneed) && not (List.mem c checkmissing)) 
	self#dump_uedges in
    self#generatepair edgel
  method randsub numdata numcheck =
    let dataneed = prefix numdata (randlist self#dnodes) in
    let checkmissing = prefix numcheck (randlist self#cnodes) in
    self#newtest dataneed checkmissing
  method newblocks_gone blocks =
    let rec split dlist clist blist = match blist with 
      [] -> (dlist, clist)
    | Data d :: rest -> split (d :: dlist) clist rest
    | Check c :: rest -> split dlist (c :: clist) rest in
    let (dataneed, checkneed) = split [] [] blocks in
    let checkmissing = List.filter (fun x -> not (List.mem x checkneed)) self#cnodes in (*Make More Efficient*)
    self#newtest dataneed checkmissing

      
(*  DEBUGGING INFORMATION                     *)
  method print_checks lev =
    let pcheck check =
      let datas = self#dnodes_of check in
      log lev "# %d: " check;
      List.iter (fun d -> log lev "%d " d) datas;
      log lev "%s" "\n" in
    List.iter pcheck self#cnodes
  method print_data lev =
    let pdata data =
      let checks = self#cnodes_of data in
      log lev "# %d: " data;
      List.iter (fun c -> log lev "%d " c) checks;
      log lev "%s" "\n" in
    List.iter pdata self#dnodes
  method equal (g2: 'a) =
    let pairs1 = self#dump_edges and pairs2 = g2#dump_edges in
    List.fold_left2 (fun v (a,b) (c,d) -> v && ((a=c) && (b=d))) true pairs1 pairs2;
    
    
(*  DECODING ROUTINES                         *)
  method decode =
    let dropdata c = 
      match self#dnodes_of c with
	[] -> ()
      |	d::[] -> self#deldata d
      |	_ -> () in
    List.iter dropdata self#cnodes;
    assert self#consistent
  method gje_simplify =
    let keep = Hashtbl.create self#dcount in  (* we might need one per data *)
    let optimize elem =
      log 4 "E%d " elem;
      let checks = self#cnodes_of elem in
      let (kept,notkept) = List.partition (Hashtbl.mem keep) checks in
      match kept, notkept with
	_ , [] -> ()
      |	[], piv :: [] -> Hashtbl.add keep piv 1
      | rest2, piv :: rest1 -> 
	  log 4 "P%d " piv;
	  Hashtbl.add keep piv 1;
	  let pattern = self#dnodes_of piv in
	  let toggle check data =
	    if self#exist data check
	    then self#del data check
	    else self#add data check in
	  let xor_row toggles row = List.iter (toggle row) toggles in
	  List.iter (xor_row pattern) (List.append rest1 rest2) in
    List.iter optimize self#dnodes;
    assert self#consistent
  method hdecode =
    let rec maxl list elem = match list with
      [] -> elem
    | h::t -> maxl t (max elem h) in
    let heapsize = maxl self#cnodes 2 in
    let heap = new Sizeheap.sizeheap heapsize in
    let heapadd c = heap#insert c (self#dcount_of c) in
    List.iter heapadd self#cnodes;
    while heap#lookminval = 1 do
      try 
	let datalist = self#dnodes_of heap#popmin_id in
	assert ((List.length datalist) = 1);
	let data = List.hd datalist in
	log 4 "D%d " data;
	List.iter heap#decrement (self#cnodes_of data);
	self#deldata data
      with Failure "hd" -> assert false
    done;
    assert self#consistent
      
(* TARGETED DECODING ROUTINES *)
  method decodet targ =
    let dropdata c = 
      match self#dnodes_of c with
	[] -> ()
      |	x::[] -> 
	  if x=targ then raise Hit_target 
	  else self#deldata x
      |	_ -> () in
    List.iter dropdata self#cnodes;
    assert self#consistent
  method hdecodet targ =
    let rec maxl list elem = match list with
      [] -> elem
    | h::t -> maxl t (max elem h) in
    let heapsize = maxl self#cnodes 2 in
    let heap = new Sizeheap.sizeheap heapsize in
    let heapadd c = heap#insert c (self#dcount_of c) in
    List.iter heapadd self#cnodes;
    let heapdel d = List.iter heap#decrement (self#cnodes_of d);
      self#deldata d in
    while heap#lookminval = 1 do
      try 
	let datalist = self#dnodes_of heap#popmin_id in
	assert ((List.length datalist) = 1);
	let data = List.hd datalist in
	log 4 "D%d " data; 
	if data = targ then raise Hit_target
	else heapdel data
      with Failure "hd" -> assert false
    done;
    assert self#consistent
      
      
(* DECODING HELPER FUNCTIONS *)
  method fulldecode =
    log 3 "%s" "#  pre:"; self#print_checks 3;
    self#decode;
    self#decode;
    log 3 "%s" "#  mid:"; self#print_checks 3;
    self#gje_simplify;
    log 3 "%s" "# mid2:"; self#print_checks 3;
(*    self#hdecode; *)
    log 3 "%s" "# post:"; self#print_checks 3;
  method candecodet target = match target with 
    Data t -> 
      log 3 "%s" "#  pre:"; self#print_checks 3;
      self#decodet t;
      self#hdecodet t;
      log 3 "%s" "#  mid:"; self#print_checks 3;
      self#gje_simplify;
      log 3 "%s" "# post:"; self#print_checks 3;
      not self#unbalanced
  | Check c -> 
      log 3 "%s" "#  pre:"; self#print_checks 3;
      self#decode;
      self#hdecode;
      log 3 "%s" "#  mid:"; self#print_checks 3;
      self#gje_simplify;
      log 3 "%s" "# post:"; self#print_checks 3;
      not self#unbalanced
      
  method ranksub erasures =
    let subgraph = self#randsub erasures 0 in
    subgraph#fulldecode;
    not subgraph#unbalanced
  method testrank trials erasures = 
    repeat trials 
      (if self#ranksub erasures then 1 else 0)
  method rankreal erasures = 
    let keepdata = ref self#dcount in
    let keepcheck = ref self#ccount in
    if !keepdata + !keepcheck < erasures then raise (Broken_input (sprintf "rankreal: %d, %d, %d" !keepdata !keepcheck erasures ));
    for i = 1 to erasures do 
      if Random.int (!keepdata + !keepcheck) < !keepdata 
      then keepdata := !keepdata - 1
      else keepcheck := !keepcheck - 1
    done;
    let numdata = self#dcount - !keepdata in
    let numcheck = self#ccount - !keepcheck in
    assert (numdata + numcheck = erasures);
    let subgraph = self#randsub numdata numcheck in
    subgraph#fulldecode;
    not subgraph#unbalanced
  method testreal trials erasures =
    try 
      repeat trials 
	(if self#rankreal erasures then 1 else 0)
    with Broken_input x -> 
      eprintf "Testreal (%d)\n" erasures;
      eprintf "self: ";  self#print_checks (-1);
      raise (Broken_input x)
  method private stress blocks_gone block_left = 
    match block_left with 
      deletion :: remaining -> 
	let new_gone = deletion :: blocks_gone in 
	let test_graph = self#newblocks_gone new_gone in
	if test_graph#candecodet deletion
	then self#stress new_gone remaining
	else new_gone
    | [] -> raise (Broken_input "Ran out of blocks to add")
  method stress_data trials = 
    repeat trials 
      (List.length (self#stress [] ((randlist self#dblocks) @ self#cblocks)))
  method stress_real trials = 
    repeat trials
      (List.length (self#stress [] (randlist (self#dblocks @ self#cblocks))))
(*  GRAPH GENERATION                          *)
  method generate_rand data check =
    let llist = randlist(gen_freq_list data 0) in
    let rlist = gen_freq_list check 0 in
    let outgraph = self#newgraph in
    List.iter2 outgraph#add data check;
    outgraph
  method generateedg data check =
    let outgraph = self#newgraph in
    List.iter2 outgraph#add data check;
    outgraph
  method generatepair pair_list =
    let outgraph = self#newgraph in
    List.iter (fun (d,c) -> outgraph#add d c) pair_list;
    outgraph
  method generate_regular (num_edges: int) (num_data: int) =
    let outgraph = self#newgraph in 
    for i = 0 to num_edges-1 do 
      outgraph#add (i mod num_data) (i / num_data) 
    done;
    outgraph
  method generate_MDS (n: int) (irr: int) (checks: int) =
    let outgraph = self#newgraph in 
    let field = new Bffield.bffield n irr in
    let array = field#mds_array (checks/n) in
    let rows = Array.length array and cols = Array.length array.(0) in
    for c=0 to (rows*n)-1 do
      let fromrow = c / n and frombit = c mod n in
      for d = 0 to cols-1 do
	let elem = field#out array.(fromrow).(d) in
	let getbit int pos = ((int lsr pos) land 1) in
	if getbit elem frombit = 1 then
	  outgraph#add d c;
      done;
    done;
    outgraph
  method generate_cid (data: int) (checks: int) =
    let outgraph = self#newgraph in 
    let pattern_size = data / checks in
    for datanum = 0 to data-1 do
      let patnum = datanum / checks in
      let shiftnum = datanum mod checks in
      outgraph#add datanum shiftnum;
      outgraph#add datanum ((shiftnum + patnum) mod checks);
    done;
    outgraph
end;;

let rec uniq1 singles list = match list with
  [] -> singles
| h :: t -> 
    if List.mem h singles 
    then uniq1 singles t 
    else uniq1 (h::singles) t;;
let uniq list = uniq1 [] list;;

module OrderedPair =
  struct
    type t = int * int
    let compare (x1,y1) (x2,y2) = if x1 = x2 then y1-y2 else x1-x2
  end;;
module OPSet = Set.Make(OrderedPair);;
    
class sparseGraph1 = object (self) (* using a set to hold edge data *)
  val mutable edge = OPSet.empty  (* elements are (Data,Check) *)
  inherit graph

  method is_empty = OPSet.is_empty edge
  method edges = OPSet.cardinal edge
  method dnodes = uniq (List.map fst (OPSet.elements edge))
  method cnodes = uniq (List.map snd (OPSet.elements edge))
  method exist data check = OPSet.mem (data,check) edge
  method dnodes_of check = List.map fst (OPSet.elements (OPSet.filter (fun (a,b) -> b = check) edge))
  method cnodes_of data = List.map snd (OPSet.elements (OPSet.filter (fun (a,b) -> a = data) edge))
  method dump_edges = 
    let compare (x1,y1) (x2,y2) = if x1 = x2 then y1-y2 else x1-x2 in
    List.sort compare (OPSet.elements edge)
  method consistent = true

  method add data check = edge <- OPSet.add (data,check) edge
  method del data check = edge <- OPSet.remove (data,check) edge
  method deldata data = edge <- OPSet.filter (fun (a,b) -> a != data) edge
  method delcheck check = edge <- OPSet.filter (fun (a,b) -> b != check) edge
      
  method newdata data = {< edge=OPSet.filter (fun (a,b) -> List.mem a data) edge >}
  method newgraph = {< edge = OPSet.empty >}
end;;
    
class sparseGraph2 = object (self) (* Using a list to hold edge data *)
  val mutable edge = []  (* elements are (Data,Check) *)
  inherit graph
      
  method is_empty = edge = []
  method edges = List.length edge
  method dnodes = uniq (List.map fst edge)
  method cnodes = uniq (List.map snd edge)
  method exist data check = List.mem (data,check) edge
  method dnodes_of check = List.map fst (List.filter (fun (a,b) -> b = check) edge)
  method cnodes_of data = List.map snd (List.filter (fun (a,b) -> a = data) edge)
  method dump_edges = List.sort compare edge
  method consistent = true
      
  method add data check = edge <- (data, check)::edge
  method del data check = edge <- List.filter (fun (a,b) -> (a,b) != (data, check)) edge
  method deldata data = edge <- List.filter (fun (a,b) -> a != data) edge
  method delcheck check = edge <- List.filter (fun (a,b) -> b != check) edge
      
  method newdata data = {< edge=List.filter (fun (a,b) -> List.mem a data) edge >}
  method newgraph = {< edge = [] >}
end;;
    
    
let rec equal_plists l1 l2 = match l1 with 
  (a1,b1)::t1 -> begin match l2 with 
    (a2,b2)::t2 -> (a1==a2)&&(b1==b2)&&(equal_plists t1 t2)
  | [] -> false
  end
|  [] -> begin match l2 with
    (a2,b2)::t2 -> false
  | [] -> true
end

type connection = {from: int; mutable to_list: int list };;
class sparseGraph3 = 
  let find_node node list = List.find (fun a -> a.from = node) !list in
  let get_list node list = 
    try (find_node node list).to_list 
    with Not_found -> [] in
  let replace_list f new_to blist = 
    if List.length new_to = 0 then
      blist := List.filter (fun a -> a.from != f) !blist
    else try 
      (fun a -> a.to_list <- new_to) (find_node f blist)
    with Not_found -> blist := {from = f; to_list = new_to}::!blist in
  let remove_connection from t list = 
    let new_to = List.filter (fun a -> a != t) (get_list from list) in
    replace_list from new_to list in
  let add_connection from t list = 
    let new_to =  t :: (get_list from list) in
    replace_list from new_to list in
  object (self)
    val mutable edgecount = 0
    val dedges = ref []  (* elements are connections (Data,[C1;C2;...])*)
    val cedges = ref []  (* elements are connections (Check,[D1;D2;...]) *)
    inherit graph
    method is_empty = edgecount = 0
    method edges = edgecount
    method dnodes = List.map (fun a -> a.from) !dedges
    method cnodes = List.map (fun a -> a.from) !cedges
    method private existd data check = List.mem check (get_list data dedges)
    method private existc data check = List.mem data (get_list check cedges)
    method exist = self#existd
    method dnodes_of check = get_list check cedges
    method cnodes_of data = get_list data dedges
    method consistent = true;
    method add data check = 
      if not (self#exist data check) then begin
	add_connection data check dedges;
	add_connection check data cedges;
	edgecount <- edgecount + 1;
      end
    method del data check = 
      if self#exist data check then begin
	remove_connection data check dedges;
	remove_connection check data cedges;
	edgecount <- edgecount - 1;
      end
(*    method deldata data = 
      replace_list data [] dedges;
      List.iter (fun c -> remove_connection c data cedges) (self#cnodes_of data)  *)
    method delcheck check = 
      replace_list check [] cedges;
      List.iter (fun d -> remove_connection check d dedges) (self#dnodes_of check)
    method newgraph = 
      {< dedges = ref [];
	cedges = ref [];
	edgecount = 0; >}
  end;; 


class bitmatrixgraph1 datasize checksize = 
  let nonzero_values array = List.filter (fun a -> a != 0) (Array.to_list array) in
  let nonzero_indexes array = nonzero_values 
      (Array.mapi (fun i a -> if a=0 then i else 0) array) in
  object (self)
    val edge = Array.create_matrix datasize checksize false
    val datacount = Array.make datasize 0;
    val checkcount = Array.make checksize 0;
    val mutable edgecount = 0
    inherit graph
    method is_empty = edgecount = 0
    method edges = edgecount
    method dnodes = nonzero_indexes datacount
    method cnodes = nonzero_indexes checkcount
    method exist data check = edge.(data).(check)
    method dnodes_of check = 
      assert (check < checksize);
      let ret = ref [] in
      for i = 0 to datasize-1 do
	if edge.(i).(check) then ret := i :: !ret
      done;
      !ret
    method cnodes_of data = 
      assert (data < datasize);
      let ret = ref [] in
      for i = 0 to checksize-1 do
	if edge.(data).(i) then ret := i :: !ret
      done;
      !ret
    method add data check = 
      assert (data < datasize); assert (check < checksize);
      if not edge.(data).(check) then begin
	edge.(data).(check) <- true;
	datacount.(data) <- datacount.(data) + 1; 
	checkcount.(check) <- checkcount.(check) + 1;
	edgecount <- edgecount + 1;
      end
    method del data check = 
      assert (data < datasize); assert (check < checksize);
      if edge.(data).(check) then begin
	edge.(data).(check) <- false;
	datacount.(data) <- datacount.(data) - 1; 
	assert (datacount.(data) >= 0);
	checkcount.(check) <- checkcount.(check) - 1; 
	assert (checkcount.(check) >= 0);
	edgecount <- edgecount - 1;
	assert (edgecount >= 0);
      end
    method consistent = false
    method newdata data = 
      let newedge = Array.create datasize (Array.create checksize false) in
      let newdata = Array.make datasize 0 in
      let newcheck = Array.make checksize 0 in 
      let newedgecount = ref 0 in 
      let updatestuff d = 
	newedge.(d) <- Array.copy edge.(d);
	newdata.(d) <- datacount.(d);
	for i = 0 to checksize-1 do 
	  if edge.(d).(i) then begin 
	    newcheck.(i) <- newcheck.(i) + 1; 
	    newedgecount := !newedgecount + 1;
	  end;
	done; in 
      List.iter updatestuff data;
      {< edge=newedge;
	datacount = newdata;
	checkcount = newcheck; 
	edgecount = !newedgecount >}
    method newgraph = 
      {<
	edge = Array.create_matrix datasize checksize false;
	datacount = Array.make datasize 0;
	checkcount = Array.make checksize 0;
	edgecount = 0 >}

  end;;

class bitmatrixgraph2 datasize checksize = object (self)
  val edge = Array.create_matrix datasize checksize false
  val datacount = Array.make datasize 0;
  val checkcount = Array.make checksize 0;
  val mutable edgecount = 0
  inherit graph
  method new_graph = {< >}
  method is_empty = edgecount = 0
  method edges = edgecount
  method dnodes = List.filter (fun a -> a != 0) (Array.to_list (Array.mapi (fun a i -> if a=0 then i else 0) datacount))
  method cnodes = List.filter (fun a -> a != 0) (Array.to_list (Array.mapi (fun a i -> if a=0 then i else 0) checkcount))
  method exist data check = edge.(data).(check)
  method dnodes_of check = 
    let ret = ref [] in
    for i = 0 to datasize-1 do
      if edge.(i).(check) then ret := i :: !ret
    done;
    !ret
  method cnodes_of data = 
    let ret = ref [] in
    for i = 0 to checksize-1 do
      if edge.(data).(i) then ret := i :: !ret
    done;
    !ret
  method add data check = 
    if not edge.(data).(check) then begin
      edge.(data).(check) <- true;
      datacount.(data) <- datacount.(data) + 1; 
    end
  method del data check = 
    if edge.(data).(check) then begin
      edge.(data).(check) <- false;
      datacount.(data) <- datacount.(data) - 1;
      checkcount.(check) <- checkcount.(check) - 1;
      edgecount <- edgecount - 1;
    end
  method consistent = false
  method newdata data = 
    let newedge = Array.create datasize (Array.create checksize false) in
    let newdata = Array.make datasize 0 in
    let newcheck = Array.make checksize 0 in 
    let newedgecount = ref 0 in 
    let updatestuff d = 
      newedge.(d) <- Array.copy edge.(d);
      newdata.(d) <- datacount.(d);
      for i = 0 to checksize-1 do 
	if edge.(d).(i) then begin 
	  newcheck.(i) <- newcheck.(i) + 1; 
	  newedgecount := !newedgecount + 1;
	end;
      done; in 
    List.iter updatestuff data;
    {< edge=newedge; 
      datacount = newdata; 
      checkcount = newcheck; 
      edgecount = !newedgecount >}
  method newgraph = 
    {<edge = Array.create_matrix datasize checksize false;
      datacount = Array.make datasize 0;
      checkcount = Array.make checksize 0;
      edgecount = 0 >}

end;;