open Expr
open Simplify

type poly_term = {
  coeff: float;
  powers: (string * int) list;
}

type polynomial = poly_term list

let rec expr_to_poly expr vars =
  match expr with
  | Const c -> [{coeff = c; powers = []}]
  | Var v when List.mem v vars -> [{coeff = 1.0; powers = [(v, 1)]}]
  | Var _ -> [{coeff = 1.0; powers = []}]
  | Mul (Const c, e) | Mul (e, Const c) ->
      List.map (fun t -> {t with coeff = t.coeff *. c}) (expr_to_poly e vars)
  | Add (e1, e2) ->
      (expr_to_poly e1 vars) @ (expr_to_poly e2 vars)
  | Mul (e1, e2) ->
      let p1 = expr_to_poly e1 vars in
      let p2 = expr_to_poly e2 vars in
      List.concat_map (fun t1 ->
        List.map (fun t2 ->
          let new_coeff = t1.coeff *. t2.coeff in
          let new_powers = combine_powers t1.powers t2.powers in
          {coeff = new_coeff; powers = new_powers}
        ) p2
      ) p1
  | Pow (Var v, Const n) when List.mem v vars && Float.is_integer n ->
      [{coeff = 1.0; powers = [(v, int_of_float n)]}]
  | _ -> [{coeff = 1.0; powers = []}]

and combine_powers p1 p2 =
  let rec merge acc = function
    | [] -> List.rev acc
    | (v, n) :: rest ->
        match List.assoc_opt v acc with
        | Some m -> merge ((v, m + n) :: List.remove_assoc v acc) rest
        | None -> merge ((v, n) :: acc) rest
  in
  merge p1 p2

let poly_to_expr poly =
  let term_to_expr t =
    let coeff_expr = if t.coeff = 1.0 && t.powers <> [] then None
                     else Some (Const t.coeff) in
    let vars_expr = List.fold_left (fun acc (v, n) ->
      let var_power = if n = 1 then Var v
                      else Pow (Var v, Const (float_of_int n)) in
      match acc with
      | None -> Some var_power
      | Some e -> Some (Mul (e, var_power))
    ) None t.powers in
    match (coeff_expr, vars_expr) with
    | None, None -> Const 1.0
    | Some c, None -> c
    | None, Some v -> v
    | Some c, Some v -> Mul (c, v)
  in
  match poly with
  | [] -> Const 0.0
  | [t] -> term_to_expr t
  | t :: rest ->
      List.fold_left (fun acc term ->
        Add (acc, term_to_expr term)
      ) (term_to_expr t) rest

let expand expr =
  let rec expand_expr = function
    | Add (e1, e2) -> simplify (Add (expand_expr e1, expand_expr e2))
    | Sub (e1, e2) -> simplify (Sub (expand_expr e1, expand_expr e2))
    | Mul (Add (a, b), c) | Mul (c, Add (a, b)) ->
        simplify (Add (Mul (expand_expr a, expand_expr c),
                      Mul (expand_expr b, expand_expr c)))
    | Mul (Sub (a, b), c) | Mul (c, Sub (a, b)) ->
        simplify (Sub (Mul (expand_expr a, expand_expr c),
                      Mul (expand_expr b, expand_expr c)))
    | Mul (e1, e2) -> simplify (Mul (expand_expr e1, expand_expr e2))
    | Pow (Add (a, b), Const n) when Float.is_integer n && n >= 0.0 && n <= 10.0 ->
        let rec binomial_expand base pow =
          if pow = 0.0 then Const 1.0
          else Mul (expand_expr base, binomial_expand base (pow -. 1.0))
        in
        expand_expr (binomial_expand (Add (a, b)) n)
    | Pow (e, n) -> Pow (expand_expr e, n)
    | e -> e
  in
  let expanded = expand_expr expr in
  simplify expanded

let rec poly_gcd p1 p2 =
  if List.length p2 = 0 then p1
  else
    let remainder = poly_remainder p1 p2 in
    poly_gcd p2 remainder

and poly_remainder _p1 _p2 =
  []

let degree expr var =
  let rec find_degree = function
    | Var v when v = var -> 1
    | Pow (Var v, Const n) when v = var && Float.is_integer n -> int_of_float n
    | Mul (e1, e2) -> find_degree e1 + find_degree e2
    | Add (e1, e2) -> max (find_degree e1) (find_degree e2)
    | Sub (e1, e2) -> max (find_degree e1) (find_degree e2)
    | _ -> 0
  in
  find_degree expr

let collect expr _var =
  let expanded = expand expr in
  simplify expanded

let factor expr =
  [expr]