#lang plai-typed

(require [typed-in racket
                   (random : (number -> number))
                   (log : (number -> number))
                   (quotient : (number number -> number))
                   (inexact->exact : (number -> number))])

(define-type (Join-List 'a)
  (empty-join-list)
  (one (elt : 'a))
  (join-list (lst1 : (Join-List 'a)) (lst2 : (Join-List 'a)) (size : number)))

;; is-join-list?: any/c -> boolean
;; takes a datum and returns true if it is a join list
;; false otherwise
#|(define (is-join-list? (dat : 'a)) : boolean
  (or (empty-join-list? dat)
      (one? dat)
      (join-list? dat)))
|#
;; new-cons: any/c list -> list
;; redefines cons to throw an error when used with append lists
;; to be exported as "cons"
#|(define (new-cons (elt : 'a) (lst : (listof 'a)))
  (cond
    [(or (one? lst) (join? lst))
     (error 'cons "do not use cons with join lists!")]
    [(or (empty? lst) (cons? lst)) (cons elt lst)]
    [else (error 'cons "second argument must be of type <list>, given ~a and ~a" elt lst)]))|#

;; join?: any/c -> boolean
;; takes a datum and returns true if it is a join list with multiple elements
;;    (a join-list struct), returns false otherwise
#|(define (join? (lst : 'a)) : boolean
  (type-case lst
    [join-list (l1 l2 size) #t]
    [else #f]))

(define (one? (lst : 'a)) : boolean
  (type-case lst
    [one (elem) #t]
    [else #f]))

(define (empty? (lst : 'a)) : boolean
  (type-case lst
    [empty #t]
    [else #f]))
|#
;; size: join-list -> number
;; returns the size of the join-list
(define (size (lst : (Join-List 'a))) : number
  (type-case (Join-List 'a) lst
    [empty-join-list () 0]
    [one (elem) 1]
    [join-list (lst1 lst2 s) (join-list-size lst)]))
 
#|(cond
    [(empty-join-list? lst) 0]
    [(one? lst) 1]
    [(join-list? lst) (join-list-size lst)]))|#

;; get: one -> any/c
;; consumes a singleton join list and returns the data contained within
#|(define (get (lst : (Join-List 'a))) : 'a
  (if (one? lst)
      (one-elt lst)
      (error 'get "give get a one")))
|#

;; join: join-list join-list -> join-list
;; consumes two join lists and joins them together into a single list
(define (join (lst1 : (Join-List 'a)) (lst2 : (Join-List 'a))) : (Join-List 'a)
  (cond
    [(empty-join-list? lst1) lst2]
    [(empty-join-list? lst2) lst1]
    [else (join-list lst1 lst2 (+ (size lst1) (size lst2)))]))

;; get-moved: join-list join-list number number -> join-list join-list
;; consumes two join lists, a maximum number of elements to move, and a direction to move them
;; produces two join lists, the first is the left half after moving the specified number of
;;    elements in the appropriate direction
;;    the second is the right half
(define (get-moved (left : (Join-List 'a)) (right : (Join-List 'a)) (num : number) (the-dir : number)) : ((Join-List 'a)*(Join-List 'a))
  (local [;; get-sub-lists: join-list dir -> join-list join-list
          ;; consumes a join-list and a direction
          ;; returns two join-lists where the first either a prefix (the-dir = 0) or
	  ;;    a suffix (the-dir = 1) of the inputted list with length no greater than num
          ;;    the second is the other part of the inputted list
          (define (get-sub-lists (a-list : (Join-List 'a)) (dir : ((Join-List 'a) -> (Join-List 'a)))) : ((Join-List 'a)*(Join-List 'a))
            (if (<= (size a-list) num)
                (values a-list (empty-join-list))
                (local ((define-values (sub-list new-side) (get-sub-lists (dir a-list) dir)))
                  (values sub-list
                          (if (= the-dir 0)
                              (join new-side (join-list-lst2 a-list))
                              (join (join-list-lst1 a-list) new-side))))))]
    (local ((define-values (moved stayed) (if (= the-dir 0)
                                              (get-sub-lists right join-list-lst1)
                                              (get-sub-lists left join-list-lst2))))
      (if (= the-dir 0)
          (values (join left moved) stayed)
          (values stayed (join moved right))))))

;; split: join-list (join-list join-list -> a) -> a
;; consumes a join-list with multiple elements and a handler
;; the handler takes two halves of a list and does a computation
(define (split (lst : (Join-List 'a)) (proc : ((Join-List 'a) (Join-List 'a) -> 'b))) : 'b
  (type-case (Join-List 'a) lst
    [empty-join-list () (error 'join-list->list "expected argument of type <join-list>")]
    [one (elem) (error 'join-list->list "expected argument of type <join-list>")]
    [join-list (lst1 lst2 s) (local ((define left (join-list-lst1 lst))
                                     (define right (join-list-lst2 lst))
                                     (define left-size (size left))
                                     (define right-size (size right))
                                     (define-values (new-left new-right)
                                       (cond
                                         [(= (size lst) 2) (values left right)] ;; do nothing
                                         [(= left-size right-size) (get-moved left right 1 (random 2))]
                                         [(> left-size right-size) (get-moved left
                                                                              right
                                                                              (+ (random (- left-size right-size)) 1)
                                                                              1)]
                                         [(< left-size right-size) (get-moved left
                                                                              right
                                                                              (+ (random (- right-size left-size)) 1)
                                                                              0)])))
                               (proc new-left new-right))]))
      #|   ; (begin 
          ;  (set-join-list-lst1! lst new-left)
           ; (set-join-list-lst2! lst new-right)
          (proc new-left new-right))
  (if (not (join-list? lst))
      (error 'split "expected first argument of type <join-list> with multiple elements")
      (local ((define left (join-list-lst1 lst))
             (define right (join-list-lst2 lst))
             (define left-size (size left))
             (define right-size (size right))
             (define-values (new-left new-right)
                   (cond
                     [(= (size lst) 2) (values left right)] ;; do nothing
                     [(= left-size right-size) (get-moved left right 1 (random 2))]
                     [(> left-size right-size) (get-moved left
                                                          right
                                                          (+ (random (- left-size right-size)) 1)
                                                          1)]
                     [(< left-size right-size) (get-moved left
                                                          right
                                                          (+ (random (- right-size left-size)) 1)
                                                          0)])))
         ; (begin 
          ;  (set-join-list-lst1! lst new-left)
           ; (set-join-list-lst2! lst new-right)
          (proc new-left new-right))))|#

;; list->join-list: (listof any/c) -> join-list
;; consumes a scheme list and outputs a mostly-balanced join list representation
;;    of that same list
(define (list->join-list (lst : (listof  'a))) : (Join-List 'a)
  (local
    [(define (list->join-list-help (a-list : (listof 'a)) (len : number)) : ((Join-List 'a)*(listof 'a))
       (cond
         [(= len 0) (values (empty-join-list) a-list)]
         [(= len 1) (values (one (first a-list)) (rest a-list))]
         [(= len 2) (values (join (one (first a-list))
                                (one (second a-list)))
                          (rest (rest a-list)))]
         [else (let* ([range (inexact->exact (max 3 (floor (* 2 (log len)))))]
                      [rand (- (random range) (quotient range 2))]
                      [left-output (list->join-list-help a-list
                                                         (- (floor (/ len 2)) rand))]
                      [right-output (list->join-list-help (snd left-output)
                                                          (+ (ceiling (/ len 2)) rand))])
                 (values (join (fst left-output) (fst right-output))
                       (snd right-output)))]))
     (define constructed (list->join-list-help lst (length lst)))]
    (if (not (empty? (snd constructed)))
        (error 'list->join-list "something went wrong with join-list construction")
        (fst constructed))))

(define (fst (tuple : ('a * 'b))) : 'a
  (local ((define-values (x y) tuple))
    x))

(define (snd (tuple : ('a * 'b))) : 'a
  (local ((define-values (x y) tuple))
    y))

;; join-list->list: join-list -> list
;; consumes a join list and converts it into a scheme list representation
;; note: this does not call split
(define (join-list->list (lst : (Join-List 'a))) : (listof 'a)
  (type-case (Join-List 'a) lst
    ;[empty-join-list () (error 'join-list->list "expected argument of type <join-list>")]
    [empty-join-list () empty]
    [one (elem) (cons elem empty)]
    [join-list (lst1 lst2 s)
               (local [(define (join-list->list-help a-list acc)
                         (type-case (Join-List 'a) a-list
                           [empty-join-list () acc]
                           [one (elem) (cons (one-elt a-list) acc)]
                           [join-list (lst1 lst2 s) 
                                      (join-list->list-help lst1
                                                            (join-list->list-help lst2 acc))]))]
                 (join-list->list-help lst empty))]))
                           
      ;(error 'join-list->list "expected argument of type <join-list>")))

;[(empty-join-list? a-list) acc]
 ;                          [(one? a-list) (cons (get a-list) acc)]
  ;                         [else (join-list->list-help (join-list-lst1 a-list)
   ;                                                    (join-list->list-help (join-list-lst2 a-list)
    ;                                                                         acc))]))]

;; join-list=?: join-list join-list -> boolean
;; consumes two join-lists and returns true if they have the same elements in the same order
(define (join-list=? (lst1 : (Join-List 'a)) (lst2 : (Join-List 'a))) : boolean
   (equal? (join-list->list lst1) (join-list->list lst2)))

;; provide statement
#|(provide (rename-out (new-cons cons))
         (rename-out (make-one one))
         one?
         get
         join
         join?
         split
         (rename-out (is-join-list? join-list?))
         join-list=?
         list->join-list
         join-list->list)|#