module Stdlib.Data.List.Base;

import Juvix.Builtin.V1.List open public;
import Stdlib.Data.Fixity open;
import Stdlib.Data.Bool.Base open;
import Stdlib.Function open;
import Stdlib.Data.Nat.Base open;
import Stdlib.Data.Maybe.Base open;
import Stdlib.Trait.Ord open;
import Stdlib.Data.Pair.Base open;

--- 𝒪(𝓃). Returns ;true; if the given
--- object is in the ;List;.
{-# specialize: [1] #-}
elem {A} (eq : A → A → Bool) (s : A) : List A → Bool
  | nil := false
  | (x :: xs) := eq s x || elem eq s xs;

--- Right-associative fold.
{-# specialize: [1] #-}
foldr {A B} (f : A → B → B) (z : B) : List A → B
  | nil := z
  | (h :: hs) := f h (foldr f z hs);

syntax iterator rfor {init := 1; range := 1};

{-# specialize: [1] #-}
rfor {A B} (f : B → A → B) (acc : B) : List A → B
  | nil := acc
  | (x :: xs) := f (rfor f acc xs) x;

--- Left-associative and tail-recursive fold.
{-# specialize: [1] #-}
foldl {A B} (f : B → A → B) (z : B) : List A → B
  | nil := z
  | (h :: hs) := foldl f (f z h) hs;

syntax iterator for {init := 1; range := 1};

{-# inline: 0 #-}
for : {A B : Type} → (B → A → B) → B → List A → B := foldl;

syntax iterator map {init := 0; range := 1};

--- 𝒪(𝓃). Maps a function over each element of a ;List;.
{-# specialize: [1] #-}
map {A B} (f : A → B) : List A → List B
  | nil := nil
  | (h :: hs) := f h :: map f hs;

syntax iterator filter {init := 0; range := 1};

--- 𝒪(𝓃). Filters a ;List; according to a given predicate, i.e.,
--- keeps the elements for which the given predicate returns ;true;.
{-# specialize: [1] #-}
filter {A} (f : A → Bool) : List A → List A
  | nil := nil
  | (h :: hs) := ite (f h) (h :: filter f hs) (filter f hs);

--- 𝒪(𝓃). Returns the length of the ;List;.
length {A} (l : List A) : Nat := for (acc := zero) (_ in l) {suc acc};

--- 𝒪(𝓃). Returns the given ;List; in reverse order.
reverse {A} (l : List A) : List A := for (acc := nil) (x in l) {x :: acc};

--- Returns a ;List; of length n where every element is the given value.
replicate {A} : Nat → A → List A
  | zero _ := nil
  | (suc n) x := x :: replicate n x;

--- Takes the first n elements of a ;List;.
take {A} : Nat → List A → List A
  | (suc n) (x :: xs) := x :: take n xs
  | _ _ := nil;

--- Drops the first n elements of a ;List;.
drop {A} : Nat → List A → List A
  | (suc n) (x :: xs) := drop n xs
  | _ xs := xs;

--- 𝒪(𝓃). splitAt n xs returns a tuple where first element is xs
--- prefix of length n and second element is the remainder of the ;List;.
splitAt {A} : Nat → List A → Pair (List A) (List A)
  | _ nil := nil, nil
  | zero xs := nil, xs
  | (suc zero) (x :: xs) := x :: nil, xs
  | (suc m) (x :: xs) := case splitAt m xs of l1, l2 := x :: l1, l2;

--- 𝒪(𝓃 + 𝓂). Merges two lists according the given ordering.
terminating
merge {A} {{Ord A}} : List A → List A → List A
  | xs@(x :: xs') ys@(y :: ys') :=
    if 
      | isLT (Ord.cmp x y) := x :: merge xs' ys
      | else := y :: merge xs ys'
  | nil ys := ys
  | xs nil := xs;

--- 𝒪(𝓃). Returns a tuple where the first component has the items that
--- satisfy the predicate and the second component has the elements that don't.
partition {A} (f : A → Bool) : List A → Pair (List A) (List A)
  | nil := nil, nil
  | (x :: xs) :=
    case partition f xs of l1, l2 := ite (f x) (x :: l1, l2) (l1, x :: l2);

syntax operator ++ cons;

--- Concatenates two ;List;s.
++ : {A : Type} → List A → List A → List A
  | nil ys := ys
  | (x :: xs) ys := x :: xs ++ ys;

--- 𝒪(𝓃). Append an element.
snoc {A} (xs : List A) (x : A) : List A := xs ++ x :: nil;

--- Concatenates a ;List; of ;List;s.
flatten : {A : Type} → List (List A) → List A := foldl (++) nil;

--- 𝒪(𝓃). Inserts the given element before every element in the given ;List;.
prependToAll {A} (sep : A) : List A → List A
  | nil := nil
  | (x :: xs) := sep :: x :: prependToAll sep xs;

--- 𝒪(𝓃). Inserts the given element inbetween every two elements in the given ;List;.
intersperse {A} (sep : A) : List A → List A
  | nil := nil
  | (x :: xs) := x :: prependToAll sep xs;

--- 𝒪(1). Drops the first element of a ;List;.
tail {A} : List A → List A
  | (_ :: xs) := xs
  | nil := nil;

head {A} (a : A) : List A -> A
  | (x :: _) := x
  | nil := a;

syntax iterator any {init := 0; range := 1};

--- 𝒪(𝓃). Returns ;true; if at least one element of the ;List; satisfies the predicate.
{-# specialize: [1] #-}
any {A} (f : A → Bool) : List A → Bool
  | nil := false
  | (x :: xs) := ite (f x) true (any f xs);

syntax iterator all {init := 0; range := 1};

--- 𝒪(𝓃). Returns ;true; if all elements of the ;List; satisfy the predicate.
{-# specialize: [1] #-}
all {A} (f : A -> Bool) : List A -> Bool
  | nil := true
  | (x :: xs) := ite (f x) (all f xs) false;

--- 𝒪(1). Returns ;true; if the ;List; is empty.
null {A} : List A → Bool
  | nil := true
  | _ := false;

--- 𝒪(min(𝓂, 𝓃)). Returns a list containing the results of applying a function
--- to each pair of elements from the input lists.
{-# specialize: [1] #-}
zipWith {A B C} (f : A -> B -> C) : List A -> List B -> List C
  | nil _ := nil
  | _ nil := nil
  | (x :: xs) (y :: ys) := f x y :: zipWith f xs ys;

--- 𝒪(min(𝓂, 𝓃)). Returns a list of pairs formed from the input lists.
zip {A B} : List A -> List B -> List (Pair A B)
  | nil _ := nil
  | _ nil := nil
  | (x :: xs) (y :: ys) := (x, y) :: zip xs ys;

--- 𝒪(𝓃 log 𝓃). Sorts a list of elements in ascending order using the MergeSort
--- algorithm.
mergeSort {A} {{Ord A}} (l : List A) : List A :=
  let
    terminating
    go : Nat -> List A -> List A
      | zero _ := nil
      | (suc zero) xs := xs
      | len xs :=
        let
          len' : Nat := div len (suc (suc zero));
          splitXs : Pair (List A) (List A) := splitAt len' xs;
          left : List A := fst splitXs;
          right : List A := snd splitXs;
        in merge (go len' left) (go (sub len len') right);
  in go (length l) l;

--- On average 𝒪(𝓃 log 𝓃), worst case 𝒪(𝓃²). Sorts a list of elements in
--- ascending order using the QuickSort algorithm.
terminating
quickSort {A} {{Ord A}} : List A → List A :=
  let
    terminating
    go : List A → List A
      | nil := nil
      | xs@(_ :: nil) := xs
      | (x :: xs) :=
        case partition (Ord.cmp x >> isGT) xs of l1, l2 := go l1 ++ x :: go l2;
  in go;

--- 𝒪(𝓃) Filters out every ;nothing; from a ;List;.
catMaybes {A} : List (Maybe A) -> List A
  | nil := nil
  | (just h :: hs) := h :: catMaybes hs
  | (nothing :: hs) := catMaybes hs;

syntax iterator concatMap {init := 0; range := 1};

--- Applies a function to every item on a ;List; and concatenates the result.
--- 𝒪(𝓃), where 𝓃 is the number of items in the resulting list.
concatMap {A B} (f : A -> List B) : List A -> List B := map f >> flatten;

--- 𝒪(𝓃 * 𝓂). Transposes a ;List; of ;List;s interpreted as a matrix.
transpose {A} : List (List A) -> List (List A)
  | nil := nil
  | (xs :: nil) := map λ {x := x :: nil} xs
  | (xs :: xss) := zipWith (::) xs (transpose xss);