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


module PolyMap = Map.Make(Z)

type t = Z.t PolyMap.t



let zero =
  PolyMap.empty

let one =
  PolyMap.singleton Z.zero Z.one

let singleton ~coeff ~expnt =
  if coeff = Z.zero then
    zero
  else
    PolyMap.singleton expnt coeff

let ( ~$ ) (a, b) =
  singleton ~coeff:a ~expnt:b



let of_int a =
  PolyMap.singleton Z.zero (Z.of_int a)

let of_zint a =
  PolyMap.singleton Z.zero a



let of_arraylist_helper p (coeff,expnt) =
  if coeff = Z.zero then
    p
  else
    if PolyMap.mem expnt p then
      PolyMap.add expnt Z.((PolyMap.find expnt p) + coeff) p
    else
      PolyMap.add expnt coeff p

let to_list a =
  List.rev_map (fun (expnt,coeff) -> (coeff,expnt)) (PolyMap.bindings a)

let of_list a =
  List.fold_left of_arraylist_helper zero a

let to_array a =
  Array.of_list (to_list a)

let of_array a =
  Array.fold_left of_arraylist_helper zero a



let degree p =
  match (PolyMap.max_binding_opt p) with
      None -> Z.zero
    | Some (k,v) -> k

let get_coeff ~expnt p =
  match (PolyMap.find_opt expnt p) with
      None -> Z.zero
    | Some v -> v



let stringifier (coeff, expnt) =
  let sign_str = if coeff < Z.zero then "-" else "+" in
  let coeff_str =
    if Z.abs coeff = Z.one && expnt <> Z.zero then ""
    else Z.to_string (Z.abs coeff) in
  let expnt_str =
    if expnt = Z.zero then ""
    else if expnt = Z.one then "x"
      else "x^" ^ (Z.to_string expnt) in
  (sign_str, coeff_str ^ expnt_str)

let to_string a =
  let pieces = List.map stringifier (to_list a) in
  if List.length pieces = 0 then "0"
  else
    let first_term_sign = if fst (List.hd pieces) = "-" then "-" else "" in
    let first_term = first_term_sign ^ (snd (List.hd pieces)) in
    let term_concat s (x,y) = s ^ " " ^ x ^ " " ^ y in
    let remaining_terms = List.fold_left term_concat "" (List.tl pieces) in
    first_term ^ remaining_terms

let print a =
  print_string (to_string a)

let output channel a =
  output_string channel (to_string a)

let pp_print format a =
  let pieces = List.map stringifier (to_list a) in
  let do_first item =
    if fst item = "-" then Format.pp_print_string format "-";
    Format.pp_print_string format (snd item) in
  let do_rest item =
    Format.pp_print_string format (" " ^ (fst item));
    Format.pp_print_space format ();
    Format.pp_print_string format (snd item) in
  if List.length pieces = 0 then
    Format.pp_print_string format "0"
  else
   (do_first (List.hd pieces);
    List.iter do_rest (List.tl pieces))



let neg =
  PolyMap.map Z.neg

let add a b =
  let f key x y =
    if Z.(x + y) = Z.zero then None
    else Some Z.(x + y)
  in PolyMap.union f a b

let sub a b =
  let f key x y =
    match (x, y) with
        (None, None) -> None
      | (Some v, None) -> Some v
      | (None, Some v) -> Some Z.(~- v)
      | (Some v1, Some v2) ->
          if Z.(v1 - v2) = Z.zero then None else Some Z.(v1 - v2)
  in PolyMap.merge f a b

let single_mul expnt coeff p =
  let helper expnt2 coeff2 r =
    PolyMap.add Z.(expnt + expnt2) Z.(coeff * coeff2) r
  in PolyMap.fold helper p zero

let mul a b =
  PolyMap.fold (fun k v result -> add result (single_mul k v b)) a zero

let div_rem a b =
  let rec helper dividend quotient remainder =
    if dividend = zero then
      (quotient, remainder)
    else
      if degree b > degree dividend then
        (quotient, add dividend remainder)
      else
        let (dsor_expnt, dsor_coeff) = PolyMap.max_binding b in
        let (dend_expnt, dend_coeff) = PolyMap.max_binding dividend in
        let (coeff_div, coeff_rem) = Z.div_rem dend_coeff dsor_coeff in
        let expnt_diff = Z.(dend_expnt - dsor_expnt) in
        let divisor_adjusted = single_mul expnt_diff coeff_div b in
        let dividend_adjusted = sub dividend divisor_adjusted in
        let new_dividend = PolyMap.remove dend_expnt dividend_adjusted in
        let new_quotient =
          if coeff_div <> Z.zero then
            PolyMap.add expnt_diff coeff_div quotient
          else
            quotient in
        let new_remainder =
          if coeff_rem <> Z.zero then
            PolyMap.add dend_expnt coeff_rem remainder
          else
            remainder in
        helper new_dividend new_quotient new_remainder in
  if b = zero then
    raise Division_by_zero
  else
    helper a zero zero

let div a b =
  fst (div_rem a b)

let rem a b =
  snd (div_rem a b)



let ( + ) = add

let ( - ) = sub

let ( * ) = mul

let ( / ) = div