(in-package :lisplab) ;; should not be part of lisplab package

(defun 0elm (m)
  (if (matrix? m)
      (coerce 0 (element-type m))
      (coerce 0 (type-of m))))

(defclass template ()
  ((symbol 
    :initarg :symbol 
    :accessor template-symbol
    :documentation "The variable")
   (dynamic-symbol ;;; TODO use this rather than a dynamic variant of the other
    :initarg :dynamic-symbol 
    :accessor template-dynamic-symbol
    :documentation "The dynamic variable")
   (type 
    :initarg :type 
    :accessor template-type
    :documentation "The actual run-time type")
   ))

(defmethod print-object ((tl template) stream)
  (print-unreadable-object (tl stream :type t :identity t)
    (prin1 (template-symbol tl) stream)))

(defgeneric create-template (type symbol &rest rest))

(defgeneric handle (what template code)) 

(defgeneric extra-let*s (template))

(defgeneric extra-declares (template))


;;; Defaults

(defmethod handle (what template code) code)

(defmethod extra-let*s (template) nil)

(defmethod extra-declares (template) nil)

;;; Blas real templates

(defclass template-blas-real (template)
  ((store-symbol 
    :initform  (gensym "store")
    :accessor template-store-symbol
    :documentation "Temp variable for store")
   (rows-symbol 
    :initform  (gensym "rows")
    :accessor template-rows-symbol
    :documentation "Temp variable for rows")
   (cols-symbol 
    :initform  (gensym "cols")
    :accessor template-cols-symbol
    :documentation "Temp variable for columns")
   ))

(defmethod create-template ((type (eql 'blas-real)) 
			    symbol &rest rest)
  (make-instance 'template-blas-real :symbol symbol))

(defmethod handle ((what (eql 'mref)) (tl template-blas-real) code )
  (destructuring-bind (ref a i j) code
    (if (eql a (template-symbol tl))
	(list 'ref-blas-real-store (template-store-symbol tl) i j (template-rows-symbol tl))
	code)))

(defmethod handle ((what (eql 'rows)) (tl template-blas-real) code )
  (destructuring-bind (rows a) code
    (if (eql a (template-symbol tl))
	(template-rows-symbol tl)
	code)))

(defmethod handle ((what (eql 'cols)) (tl template-blas-real) code)
  (destructuring-bind (cols a) code
    (if (eql a (template-symbol tl))
	(template-cols-symbol tl)
	code)))

(defmethod extra-let*s ((tl template-blas-real))
  (list `(,(template-store-symbol tl) (store ,(template-symbol tl)))
	`(,(template-rows-symbol tl) (rows ,(template-symbol tl)))
	`(,(template-cols-symbol tl) (cols ,(template-symbol tl)))))

(defmethod extra-declares ((tl template-blas-real))
  (list `(blas-real ,(template-symbol tl))
	`((simple-array double-float (*)) ,(template-store-symbol tl))	
	`(type-blas-idx ,(template-rows-symbol tl))
	`(type-blas-idx ,(template-cols-symbol tl))))

;;; Double floats

(defclass template-double-float (template)())

(defmethod create-template ((type (eql 'double-float)) 
			    symbol &rest rest)
  (make-instance 'template-double-float :symbol symbol))

(defmethod extra-declares ((tl template-double-float))
  (list `(double-float ,(template-symbol tl))))


;;;; The actual optimizations

(defun handle-tree (tl code)
  (if (consp code)
      (let ((code2 (mapcar (lambda (code) (handle-tree tl code)) code)))
	(handle (car code2) tl code2))
      code))

(defun handle-all (templates code)
  (if templates 
      (handle-all (cdr templates)
		  (handle-tree (car templates) code))
      code))

(defun generate-code (syms vals code)
  (let* ((templates (mapcar #'create-template 
			   (mapcar #'type-of vals)
			   syms ))
	 (let*s (mapcan #'extra-let*s templates))
	 (declares (mapcan #' extra-declares templates))
	 (code2 (handle-all templates code)))
    `(let* ,let*s
       (declare ,@declares)
       ,@code2)))
	       
(defmacro w/dynamic (args &body body)
  "Optimized code, but without any structure information and anything 
that should be otimized must be an argument"
  (let ((run (gensym "run")))
    `(progv ',args (list ,@args)
       (let ((,run  (generate-code
		     ',args
		     (list ,@args)
		     ',body)))
	 (eval ,run)))))


(defun test-m* (A B)
  (let* ((M (rows a))
	 (N (cols b))
	 (S (cols a))
	 (c (create a 0 (list M N)))
	 (tmp 0d0))
    (w/dynamic (a b c tmp)
      (dotimes (i (rows A))
	(dotimes (j (cols B))
	  (setf tmp 0d0)
	  (dotimes (k (cols A))
	    (incf tmp (* (mref a i k) (mref b k j))))
	  (setf (mref c i j) tmp)))
      c)))

#+nil (defun test-m* (a b)
  (let* ((M (rows a))
	 (N (cols b))
	 (S (cols a))
	 (c (create a 0 (list M N)))
	 (tmp 0))
    (w/dynamic (a b c M N S)
      (dotimes (i M)
	(dotimes (j N)
	  #+nil (setf tmp 0)
	  (dotimes (k S)
	    (incf (mref C i j) (* (mref a i k) (mref b k j))))))
      c)))












#|


;; The parsing context

(defclass context ()
  ((templates
    :initarg :templates
    :initform nil
    :accessor context-templates
    :documentation "The context")
   #+nil (code
    :initarg :code
    :initform nil
    :accessor context-code
    :documentation "The code")))

(defmethod print-object ((c context) stream)
  (print-unreadable-object (c stream :type t :identity t)
    (dolist (tl (context-templates c))
      (format stream "~&~A" tl))))

(defgeneric push-template (context template))

(defmethod push-template ((c context) tl)
  (setf (context-templates c)
	(cons tl (context-templates c)))
  tl) 

(defgeneric pop-template (context))

(defmethod pop-template ((c context))
  (let ((x (car (context-templates c))))
    (setf (context-templates c)
	  (cdr (context-templates c)))
    x)) 
  

#+nil (defgeneric optimize-context (template context))

;;; defaults

(defmethod handle (what template code) nil)

(defmethod handle-all (context code)
  (if (consp code)
      (progn 
	;; TODO let a new kind of declare to update context
	(let ((code2 (mapcar (lambda (c) 
			       (handle-all context c)) 
			     code)))
	  (dolist (tl (context-templates context))
	    (let ((x (handle (car code2) tl code2)))
	      (if x (return-from handle-all x))))
	  code2))
      code))

(defgeneric extra-declares (template))
  











(defgeneric apply-template (template code))
  
(defgeneric make-template-let*-forms (template))

(defgeneric make-template-declare-forms (template))

;;;; Blas real templates

(defmethod make-template-declare-forms ((tl template-blas-real))
  `((type type-blas-store ,(template-store-symbol tl) )
    (type type-blas-idx ,(template-rows-symbol tl) )))

(defmethod make-template-let*-forms ((tl template-blas-real ))
  `((,(template-store-symbol tl) (store ,(template-symbol tl) ))
    (,(template-rows-symbol tl) (rows ,(template-symbol tl)))))

(defmethod apply-template ((tl template-blas-real) code)
  (if (consp code)
      (let* ((code (sublis '((.+ . +) (.* . *) (.- . -) (./ . /) (.^ . expt)) code))  
	     (f (car code))) 
	(cond ((and (eql f 'mref)
		    (eql (cadr code) (template-symbol tl)))
	       `(ref-blas-real-store 
		 ,(template-store-symbol tl) 
		 ,(third code) 
		 ,(fourth code) 
		 ,(template-rows-symbol tl)))
	      ((and (eql f 'vref)
		    (eql (cadr code) (template-symbol tl)))
	       `(aref ,(template-store-symbol tl) 
		      ,(third code)))
	      (t code)))
      code))

(defun apply-templates (templates code)  
  (if (consp code)
      (if (eql (car code) 'declare)
	  (cons 'delcare 
		(mapcar (lambda (decl) 
			  (if (and (consp decl)
				   (eql (car decl) 'ttype))
			      ;; TODO real code here
			      decl
			      decl))
			(cdr code)))
	  (progn
	    (dolist (tl templates)
	      (setf code (apply-template tl code)))
	    (mapcar (lambda (c) (apply-templates templates c)) code)))
      code))




(defun template-generate-optimized-code (templates code)
  (let ((let*s (mapcan #'make-template-let*-forms templates))
	(declares (mapcan #'make-template-declare-forms templates))
	(code2 (apply-templates templates code)))
    `(let* ,let*s
       (declare ,@declares)
       ,@code2)))
	       
#+nil (defun generate-optimized-code-trial (templates code)
  (cons 'progn code))

(defun template-assemble-code (syms vals code)
  (let ((templates  (mapcar #'create-template 
			    (mapcar #'type-of vals)
			    syms )))
    (template-generate-optimized-code templates code)))

(defmacro template-lambda (args &body body)
  "Optimized code, but without any structure information and anything 
that should be otimized must be an argument"
  (let ((run (gensym "run")))
    `(lambda ,args 
       (progv ',args (list ,@args)
	 (let ((,run  (template-assemble-code
		       ',args
		       (list ,@args)
		       ',body)))
	   (eval ,run))))))

(defun test-template-lambda ()
  (template-lambda (a b)
    (dotimes (i (rows a))
      (dotimes (j (cols b))
	(setf (mref a i j) (.* 4 (mref b i j)))))
    a))

(defun test-mmax (x)
  (funcall 
   (template-lambda (a)
     (let ((max (vref a 0)))
       (dotimes (i (size a))
	 (when (> (vref a i) max)
	   (setf max (vref a i))))
       max))
     x))





































#+nil (defmacro template-lambda (args &body body)
  "Optimized code, but without any structure information and anything 
that should be otimized must be an argument"
  (let ((body2 (gensym "body2")))
    `(lambda ,args 
       (macrolet ((,body2 () `(generate-optimized-code 
			      (mapcar #'create-template 
				      (mapcar #'type-of (list ,,@args))
				      ',,args )
			      ,,body)))
	 (,body2)))))

#+nil (defmacro template-lambda (args &body body)
  "Optimized code, but without any structure information and anything 
that should be otimized must be an argument"
  (let ((body2 (gensym "body2")))
    (list 'lambda args 
	  (list 'macrolet (list 
			   (list body2 '() `(generate-optimized-code 
					    (mapcar #'create-template 
						    (mapcar #'type-of (list ,@args))
						    ',args )
					    ',body)))
	 (list body2)))))


(defun defmat-generate-inline-expansions (syms type element-type)
  "Inlining of particular functions. First the subsitutions. Second the definitions. Third declarations. "
  ;; TODO generalize
  '(
    ((.+ . +) 
     (.* . *) 
     (./ . /) 
     (.- . -) 
     (.^ . expt)
     ((mref a i j) . (ref-blas-real-store storea i j rowsa))
     ((mref b i j) . (ref-blas-real-store storeb i j rowsb))
     ((mref c i j) . (ref-blas-real-store storec i j rowsc)))
    ;; The symbols must be generated here. Since the symbols and declarations must be seen in 
    ;; connection
    ((storea (store a))
     (storeb (store b))
     (storec (store c))
     (rowsa (rows a))
     (rowsb (rows b))
     (rowsc (rows c))
     )
    ((type type-blas-store storea)
     (type type-blas-store storeb)
     (type type-blas-store storeb)
     (type type-blas-idx rowsa)
     (type type-blas-idx rowsb)
     (type type-blas-idx rowsc))))

(defun defmat-make-symbol-hash (operations)
  (let ((ht (make-hash-table )))
    (dolist (op operations)
      (setf (gethash op ht) (gensym (symbol-name op))))
    ht))

(defun defmat-generate-let-form (template sym helper-syms type)
  "TODO make dynamic"
  (case template
    (:ge 
     (case type 
       (blas-real 
	(let ((store (gensym "store")))
	  (list 
	   (list `(,store (store ,sym)))
	   (list `(type type-blas-store ,store))
	   (list `((mref ,sym i j) . (ref-blas-real-store store i j rowsa))))))
       (blas-complex 
	(let ((store (gensym "store")))
	  (list 
	   (list `(,store (store ,sym)))
	   (list `(type type-blas-store ,store))
	   (list `((mref ,sym i j) . (ref-blas-complex-store store i j rowsa)))))
	)))
    ;; TODO element and index
     
     ))
		 
     
#+todo-fix (defun defmat-optimize-code (template-sym-type-list code) ;; TODO types
  (destructuring-bind (lets declares subtitutions)
      (defmat-generate-inline-expansions syms t t)
    `(let ,lets
       (declare ,@declares)
       (progn . ,(defmat-apply-matches syms patsubs code)))))

(defun defmat-apply-matches (syms matches code)
  "Applies the matches to the code"
  (if matches 
      (defmat-apply-matches 
	  syms 
	  (cdr matches)
	  (defmat-apply-single-match-tree syms (caar matches) (cdar matches) code))
      code))
    
(defun defmat-apply-single-match-tree (syms pat sub code)
  (defmat-apply-single-match
      syms pat sub
      (if (consp code)
	  (mapcar (lambda (code) (defmat-apply-single-match-tree syms pat sub code))
		  code)
	  code)))
      
(defun defmat-apply-single-match (syms pat sub code)
  "Aplies one pattern pattern matche to the code"
  (if (consp code)
      (if (consp pat)
	  (let ((matches (defmat-find-matches syms pat code)))
	    (if matches
		(sublis matches sub)
		code))
	  (if (eql pat (car code))
	      (cons sub (cdr code))
	      code))
      (if (eql pat code)
	  sub
	  code)))

(defun defmat-find-matches (syms pat code)
  "Binds the mathces in the pattern to the code"
  (when (eql (car pat) (car code))
    (remove nil
	    (mapcar (lambda (p s)
		      (if (member p syms)
			  (when (not (eql p s))
			    (return-from defmat-find-matches nil))
			(cons p s)))
		    (cdr pat) (cdr  code)))))
		
(defun defmat-apply-substitutions (subs body)
  (if subs
      (defmat-apply-substitutions (cdr subs) (subst (cdar subs) (caar subs) body))
      body))

(defun defmat-parse-args (args)
  (mapcar #'car args))

(defun defmat-make-optimized-body (matforms body)
  (let ((syms (defmat-parse-args matforms)))
    (defmat-optimize-code syms body)))

#+nil (defmacro defmat (name return-type args defs &body body)
  (let ((args2 (defmat-parse-args args))
	(subst (gensym))
	(body2 (gensym)))
    `(defun ,name ,args2 
       (let ((,subst (mapcar #'cons ',args2 (list ,@args2)))
	     (,body2 (defmat-make-optimized-body ',args ',(cons 'progn body))))
	 (eval (sublis ,subst ,body2))))))

#+nil (defmacro defmat-lambda (return-type args &body body)
  (let ((args2 (defmat-parse-args args))
	(subst (gensym))
	(body2 (gensym)))
    `(lambda ,args2 
       (let ((,subst (mapcar #'cons ',args2 (list ,@args2)))
	     (,body2 #+nil (cons 'progn ',body) (defmat-make-optimized-body ',args ',body)))
	 (eval (sublis ,subst ,body2))))))

(defmacro defmat-lambda (args &body body)
  "Optimized code, but without any structure information and anything 
that should be otimized must be an argument"
  (let ((body2 (gensym "body2")))
    `(lambda ,args 
       (macrolet ((,body2 ()  (defmat-optimize-code ',args ',body)))
	(,body2)))))




;; TODO maby macrolet is enough, and eval is not needed
;; TODO make the body the hard way using (list 'progn .. etc)

#+nil (defmat m*5 :ge ((A :ge) (B :ge))
  ((C (create A (rows A) (cols B))))
  (dotimes (i (rows A))
    (dotimes (j (cols B))
      (let ((cij 0))	
	(dotimes (k (cols A))
	  (setf cij (.+ cij (.* (mref A i k ) (mref B k j)))))
	(setf (mref C i j) cij)))) 
  C)


(defmacro deftest (name args &body body)
  (let ((body2 (gensym)))
    `(defun ,name ,args
       (macrolet ((tmp () ',@body))
	 (tmp)))))


#+nil (defmacro deftest (name args &body body)
  (let ((subst (gensym))
	(body2 (gensym)))
    `(defun ,name ,args 
       (let ((,subst (mapcar #'cons ',args (list ,@args)))
	     (,body2 (copy-tree ',@body)))
	 (eval (sublis ,subst ,body2))))))

#+nil (defmacro deftest (name args &body body)
  (let ((body1 (gensym))
	(body2 (gensym)))
    `(defun ,name ,args 
       (let ((,body1 (copy-tree ',body)))
	 (macrolet ((,body2 () `(progn ,@,body1)))
	   (,body2))))))






;;;; TRASH


#+nil (defun make-defmat-form (matrix-type sym store-sym)
  ;; TODO make form based on th type
  (list :ge sym store-sym 'ref-blas-real-store 'blas-real))  

#+nil (defun make-defmat-macrolet-form (type sym) 
  ;; TODO make based on type
  (let ((i (gensym "i"))
	(j (gensym "j")))   
    `(,sym (,i ,j) 
	   (ref-blas-real-store (syms-store ',sym)  ,i ,j (syms-rows ',sym)))))

#+nil (defun make-store-form (matform)
  (destructuring-bind (tp sym store-sym ref store-type) matform
    `(,store-sym (store ,sym))))  

#+nil (defun make-store-type-declare (matform)
  (destructuring-bind (tp sym store-sym ref store-type) matform
    `(type ,store-type ,store-sym)))  


#+nil (defmacro defmat (name return-type args defs &body body)
  (let ((body2 (gensym)))
    `(defmethod ,name ,(defmat-parse-args args)
       (let ,defs
	 (let ((body ',body))
	   (macrolet ((,body2 () `(progn ,@body)))
	     (,body2)))))) )

|#