(defun sage-copy-optional-annotation ()
  "In a 'sage: ' line of a docstring, copy '# optional' from a previous line and
advance to the end of the next 'sage: ' line or to the end of the current docstring.
If invoked elsewhere, just advance to the end of the next 'sage: ' line."
  (interactive)
  (if (save-excursion (beginning-of-line)
		      (looking-at " *sage:"))
      (multiple-value-bind (tab-stop-list text advance)
	  (save-excursion
	    (previous-line)
	    (end-of-line)
	    (condition-case nil
		(re-search-backward "# *optional.*$")
	      (error
	       (values '(88)  ; alignment point for modularization "# optional"s
		       "# optional - "
		       nil))
	      (:success
	       (values (list (- (match-beginning 0)
				(save-excursion
				  (beginning-of-line) (point))))
		       (match-string-no-properties 0)
		       t))))
	(end-of-line)
	(just-one-space)
	(tab-to-tab-stop)
	(insert text)
	(when advance
	  (re-search-forward "sage: \\|\"\"\"")
	  (end-of-line)))
    (re-search-forward "sage:")
    (end-of-line)))

(defun sage--align-comment (alignment)
  (let ((tab-stop-list alignment))
    (just-one-space 0)
    (tab-to-tab-stop)
    (re-search-forward "sage: \\|\"\"\"")
    (end-of-line)))
  

(defun sage-align-optional-annotation ()
  (interactive)
  (if (save-excursion (beginning-of-line)
		      (looking-at " *sage:"))
      (let ((eol (save-excursion (end-of-line)
				 (point))))
	(beginning-of-line)
	(condition-case nil
	    (re-search-forward "# optional - \\(sage\\|numpy\\|scipy\\|sympy\\|fpylll\\|mpmath\\|pplpy\\|sphinx\\|pexpect\\)" eol)  ; standard packages
	  (:success
	   (goto-char (match-beginning 0))
	   (sage--align-comment '(88 100 120)))  ; 88 = alignment point for modularization "# optional"s
					         ; in sphinx furo style
	  (error
	   (condition-case nil
	       (re-search-forward "# optional" eol)  ; optional packages
	     (:success
	      (goto-char (match-beginning 0))
	      (sage--align-comment '(64)))
	     (error
	      (end-of-line)
	      (re-search-forward "# optional")
	      (end-of-line))))))
    (re-search-forward "sage:")
    (end-of-line)))

(defun sage-newline-keep-comment ()
  (interactive)
  (let ((pt (point))
	(eol (save-excursion (end-of-line)
			     (point))))
    (condition-case nil
	(re-search-forward "\\([^#]*\\)\\(#.*\\)?$" eol)
      (error
       (newline))
      (:success
       (let ((text (match-string 1)))
	 (replace-match (make-string (length text) ? )
			nil nil nil 1)
	 (end-of-line)
	 (newline)
	 (py-indent-line)
	 (insert "....:     ")
	 (save-excursion (insert text)))))))

(with-eval-after-load "python-mode"
  (define-key python-mode-map (kbd "C-M-;") 'sage-copy-optional-annotation)
  (define-key python-mode-map (kbd "C-M-'") 'sage-align-optional-annotation)
  (define-key python-mode-map (kbd "C-M-<return>") 'sage-newline-keep-comment))

(font-lock-add-keywords 'python-mode
			'(("^ *\\(sage: \\|[.][.][.][.]: \\)" 1
			   'font-lock-warning-face prepend)))