(* Programmed by Jedidiah Barber *)
(* Licensed under the Sunset License v1.0 *)


let usage_msg = "Usage: wallsolve [options]"
let input_seq = ref (Array.make 0 (Z.of_int 0))
let depth_limit = ref 20



let anon_fun item =
  prerr_endline "Error: Don't know what to do with anonymous argument";
  exit 2



let input_immediate items =
  if Array.length !input_seq > 0 then
   (prerr_endline "Error: Multiple input sequences specified";
    exit 2)
  else
    try input_seq := Util.read_integer_list items
    with Util.Not_an_integer s ->
      prerr_endline ("Error: The argument '" ^ s ^ "' is not an integer");
      exit 2

let input_from_file filename =
  if Array.length !input_seq > 0 then
   (prerr_endline "Error: Multiple input sequences specified";
    exit 2)
  else
    try input_seq := Util.read_sequence_file filename
    with Util.Not_an_integer s ->
      prerr_endline ("Error: The item '" ^ s ^ "' from file '" ^ filename ^ "' is not an integer");
      exit 2



let speclist =
  [("-f", Arg.String input_from_file, "Read input integer sequence from file");
   ("-i", Arg.String input_immediate, "Integer input sequence with values separated by commas");
   ("-l", Arg.Set_int depth_limit, "Sets maximum levels to generate, default 20")]



module My_Algebra =
  struct
    type t = Poly.t
    let zero = Poly.zero
    let one = Poly.one
    let add = Poly.add
    let sub = Poly.sub
    let mul = Poly.mul
    let div = Poly.div
    let to_string = Poly.to_string
  end

module My_Wall = Wall.Make (My_Algebra)



let _ =
  if Array.length Sys.argv <= 1 then
   (Arg.usage speclist usage_msg;
    exit 1);

  Arg.parse speclist anon_fun usage_msg;

  if Array.length !input_seq = 0 then
   (prerr_endline "Error: No input sequence provided";
    exit 2);

  if Array.length !input_seq < 2 then
   (prerr_endline "Error: Input sequence not long enough to solve";
    exit 2);

  (* Convert input integer sequence into the right sort of input polynomial sequence. *)
  let poly_offset n =
    Poly.of_list [(Z.neg !input_seq.(n), Z.one); (!input_seq.(n+1), Z.zero)] in
  let poly_inputs =
    Array.init (Array.length !input_seq - 1) poly_offset in

  (* Generate the number wall and look at the bottom most non-zero row. *)
  let scribbles =
    My_Wall.create
      ~dimx:(Array.length poly_inputs)
      ~dimy:!depth_limit
      ~init:poly_inputs in
  let row_of_interest =
    My_Wall.last_nonzero_row scribbles in

  (* Test whether there is actually a row of zeros under the row of interest. *)
  if row_of_interest = My_Wall.y_length scribbles - 1 then
   (prerr_endline "Error: Depth limit too low to solve sequence";
    exit 2);

  (* Test whether the row of interest has enough entries to be sure of them all being the same. *)
  let exists_pred x =
    if My_Wall.has_value scribbles ~x ~y:row_of_interest then 1 else 0 in
  let array_sum =
    Array.fold_left ( + ) 0 in
  let naturals_to n =
    Array.init n (fun n -> n) in
  if array_sum (Array.map exists_pred (naturals_to (My_Wall.x_length scribbles))) < 3 then
   (prerr_endline "Error: Input sequence not long enough to solve";
    exit 2);

  (* Need to normalize the polynomials in the row as if they were set equal to zero. *)
  let rec gcd a b =
    Z.(if b = zero then a else gcd b (a mod b)) in
  let normalize p =
    let sign_correct x =
      if Z.(Poly.get_coeff ~expnt:(Poly.degree x) x < zero) then Poly.neg x else x in
    let coefficients =
      List.map fst (Poly.to_list p) in
    let divisor =
      List.fold_left gcd Z.zero coefficients in
    sign_correct (Poly.div p (Poly.of_zint divisor)) in

  (* Now let's look at the polynomial right in the middle. *)
  let key_poly =
    normalize (My_Wall.get scribbles
      ~x:((Array.length poly_inputs) / 2 + 1)
      ~y:row_of_interest) in

  (* Check for degeneracy. *)
  if Poly.degree key_poly = Z.zero then
   (prerr_endline "Error: Result is degree zero, not sure how that happened";
    exit 2);

  (* Check that all the polynomials in the row of interest are equal once normalized. *)
  let equals_pred x =
    let value = My_Wall.get_opt scribbles ~x ~y:row_of_interest in
    if Option.is_some value then
      normalize (Option.get value) = key_poly
    else
      true in
  if not (Array.for_all equals_pred (naturals_to (My_Wall.x_length scribbles))) then
   (prerr_endline "Error: Polynomials differ, maybe try a longer sequence?";
    exit 2);

  (* Print out the polynomial as if it was a recurrence relation with the highest degree term *)
  (* being on the left hand side and the remaining terms being on the right hand side. *)
  let print_relation p =
    let lhs = List.hd p in
    let rhs = List.map (fun (x,y) -> (Z.neg x, y)) (List.tl p) in
    if fst lhs <> Z.one then print_string (Z.to_string (fst lhs)); print_string "S(n) =";
    let print_rhs_item minus plus (c,e) =
      print_string " ";
      if c < Z.zero then print_string minus else print_string plus;
      if (Z.abs c) <> Z.one then print_string (Z.to_string (Z.abs c));
      print_string ("S(n-" ^ Z.to_string Z.(snd lhs - e) ^ ")") in
    print_rhs_item "-" "" (List.hd rhs);
    List.iter (print_rhs_item "- " "+ ") (List.tl rhs) in

  print_relation (Poly.to_list key_poly);
  print_newline ();

  (* This... should probably never trigger, but might as well check anyway. *)
  if Poly.degree key_poly > Z.of_int (Array.length !input_seq) then
   (prerr_endline "Error: Recurrence relation longer than input sequence... wtfhow?";
    exit 2);

  (* Calculate and print the predicted next number of the input sequence. *)
  let calc_next p inp =
    let lhs = List.hd p in
    let rhs = List.map (fun (x,y) -> (Z.neg x, y)) (List.tl p) in
    let rec process result terms =
      if terms = [] then
        result
      else
        let current = List.hd terms in
        let index = Array.length inp - Z.(to_int (snd lhs - snd current)) in
        process Z.(result + fst current * inp.(index)) (List.tl terms) in
    Z.div (process Z.zero rhs) (fst lhs) in

  print_string ("Next: " ^ Z.to_string (calc_next (Poly.to_list key_poly) !input_seq));
  print_newline ()