(module concise mzscheme
  (provide (rename my-app #%app))
  (require-for-syntax (lib "stx.ss" "syntax"))
  (require-for-template (only (lib "1.ss" "srfi") iota))
  (require-for-template (only (lib "26.ss" "srfi") cut))
  (require (only (lib "26.ss" "srfi") cut))
  (require (only (lib "1.ss" "srfi") iota))
  
  (define (interval n m)
    (if (<= n m)
        (interval/step n m 1)
        (interval/step n m -1)))
  
  (define (inclusive-interval n m)
    (if (<= n m)
        (inclusive-interval/step n m 1)
        (inclusive-interval/step n m -1)))
  
  (define (same-sign? x y)
    (positive? (* x y)))
  
  (define (interval/step n m s)
    (if (= n m)
        '()
        (begin
          (unless (same-sign? (- m n) s)
            (if (< n m)
                (error "When <from> is less than <to> it is required that <step> is positive")
                (error "When <to> is less than <from> it is required that <step> is negative")))
          (if (< n m)
              (let loop ([c 0] [xs '()])
                (let ([x (+ n (* c s))])
                  (if (< x m)
                      (loop (+ c 1) (cons x xs))
                      (reverse! xs))))
              (let loop ([c 0] [xs '()])
                (let ([x (+ n (* c s))])
                  (if (> x m)
                      (loop (+ c 1) (cons x xs))
                      (reverse! xs))))))))
  
  (define (inclusive-interval/step n m s)
    (if (= n m)
        (list m)
        (begin
          (unless (same-sign? (- m n) s)
            (if (< n m)
                (error "When <from> is less than <to> it is required that <step> is positive")
                (error "When <to> is less than <from> it is required that <step> is negative")))
          (if (< n m)
              (let loop ([c 0] [xs '()])
                (let ([x (+ n (* c s))])
                  (if (<= x m)
                      (loop (+ c 1) (cons x xs))
                      (reverse! xs))))
              (let loop ([c 0] [xs '()])
                (let ([x (+ n (* c s))])
                  (if (>= x m)
                      (loop (+ c 1) (cons x xs))
                      (reverse! xs))))))))
  
  (define-syntax (my-app stx)
    (define (bracket-shaped? stx)
      (equal? (syntax-property stx 'paren-shape) #\[))
    (define (set-shaped? stx)
      (equal? (syntax-property stx 'paren-shape) #\{))
    (define (dots? stx)
      (and (identifier? stx)
           (module-or-top-identifier=? 
            stx (datum->syntax-object stx '...))))
    (define (range? stx)
      (syntax-case* stx (..) module-or-top-identifier=?
        [[_ a .. b]           (bracket-shaped? stx) ]
        [[_ a .. b s]         (bracket-shaped? stx)]
        [[_ a dots b]    (and (bracket-shaped? stx) (dots? #'dots)) #t]
        [[_ a dots b s]  (and (bracket-shaped? stx) (dots? #'dots)) #t]
        [[_ a .. b s]         (bracket-shaped? stx)]
        [[_ count]            (bracket-shaped? stx)]
        [[_ count start]      (bracket-shaped? stx)]
        [[_ count start step] (bracket-shaped? stx)]
        [_            #f]))
    (define (handle-range stx)
      (syntax-case* stx (..) module-or-top-identifier=?
        [[_ from .. to]
         (syntax/loc stx
           (interval from to))]
        [[_ from .. to step]
         (syntax/loc stx
           (interval/step from to step))]
        [[_ from dots to]
         (dots? #'dots)
         (syntax/loc stx
           (inclusive-interval from to))]
        
        [[_ from dots to step]
         (dots? #'dots)
         (syntax/loc stx
           (inclusive-interval/step from to step))]
        [[_ count]
         (syntax/loc stx
           (iota count))]
        [[_ count start]
         (syntax/loc stx
           (iota count start))]
        [[_ count start step]
         (syntax/loc stx
           (iota count start step))]
        [_else
         (error stx)]))
    
    (define (handle-cut stx)
      (define (underscore-to-slot stx)
        (syntax-case* stx (_) module-or-top-identifier=?
          [_ (syntax/loc stx <>)]
          [_else stx]))
      (with-syntax ([(se ...)
                     (map underscore-to-slot
                          (cdr (syntax->list stx)))])
        (syntax/loc stx
          (cut se ...))))
    
    (syntax-case* stx (.. : is ->) module-or-top-identifier=?      
      ; lambda shorthand
      [{_ n ... -> expr}
         (set-shaped? stx)
         (syntax/loc stx
           (lambda (n ...) expr))]
      [{expr-or-slot ...}
       (set-shaped? stx)
       (handle-cut stx)]
      ; "comprehension" syntax
      [[_ expr : n is while-pred-expr]
       (bracket-shaped? stx)
       (syntax/loc stx
         (let loop ([n 0])
           (if while-pred-expr
               (cons expr (loop (+ n 1)))
               '())))]
      [(_ expr : n is while-pred-expr filter-expr)
       (bracket-shaped? stx)
       (syntax/loc stx
         (let loop ([n 0])
           (if while-pred-expr
               (if filter-expr
                   (cons expr (loop (+ n 1)))
                   (loop (+ n 1)))
               '())))]
      [(_ n : while-pred-expr)
       (bracket-shaped? stx)
       (syntax/loc stx
         (let loop ([n 0])
           (if while-pred-expr
               (cons n (loop (+ n 1)))
               '())))]      
      [(_ n : while-pred-expr filter-expr)
       (bracket-shaped? stx)
       (syntax/loc stx
         (let loop ([n 0])
           (if while-pred-expr
               (if filter-expr
                   (cons n (loop (+ n 1)))
                   (loop (+ n 1)))
               '())))]
      ; ranges
      [range
       (range? stx)
       (handle-range stx)]
      [(_ . more)
       (syntax/loc stx 
         (#%app . more))])))


(require concise)


;;;
;;; RANGES
;;;

'normal-range
(equal? [1 .. 4]  '(1 2 3))
(equal? [1 ... 4] '(1 2 3 4))

'reverse-range
(equal? [4 .. 1]  '(4 3 2))
(equal? [4 ... 1] '(4 3 2 1))

'empty-or-one-element-range
(equal? [4 .. 4] '())
(equal? [4 ... 4] '(4))

'normal-range-with-step
(equal? [1 .. 7 2] '(1 3 5))
(equal? [1 .. 8 2] '(1 3 5 7))
(equal? [1 ... 7 2] '(1 3 5 7))
(equal? [1 ... 8 2] '(1 3 5 7))

'reverse-range
(equal? [4 .. 1]  '(4 3 2))
(equal? [4 ... 1] '(4 3 2 1))

'reverse-range-with-step
(equal? [7 .. 1 -2] '(7 5 3))     ; TODO: better error location for 2 instead of -2
(equal? [8 .. 1 -2] '(8 6 4 2))
(equal? [7 ... 1 -2] '(7 5 3 1))
(equal? [8 ... 1 -2] '(8 6 4 2))

;;;
;;; IOTA
;;;

(equal? [4]     '(0 1 2 3))
(equal? [4 3]   '(3 4 5 6))
(equal? [4 3 5] '(3 8 13 18))   


;;;
;;; "COMPREHENSION"
;;;

(equal? [n : (< n 5)]                       '(0 1 2 3 4))
(equal? [n : (< n 5) (even? n)]             '(0 2 4))
(equal? [(* n n) : n is (< n 5)]            '(0 1 4 9 16))
(equal? [(* n n) : n is (< n 5) (even? n)]  '(0 4 16))

;;;
;;; "CURRY"
;;;

(equal? ({- _} 3) -3)
(equal? ({list 1 _ 3 _ 5 } 2 4)  '(1 2 3 4 5))

;;;
;;; LAMBDA
;;;

(equal? (map {n -> (+ n 1)} [0 ... 3])  '(1 2 3 4))

;;;
;;; EXAMPLES
;;;

(define (square? n)
  (= (integer-sqrt n) (sqrt n)))

(equal? (map {+ _ 3} [1 ... 3])  '(4 5 6))
(equal? (filter {square? _} [1 ... 100])  '(1 4 9 16 25 36 49 64 81 100))
(equal? (filter {n -> (and (odd? n) (square? n))}  [1 ... 100])
        '(1 9 25 49 81))