(ns myapp.sm
  (:require [clojure.core.matrix :as matr])
  (:use [clojure.test]))

(def m-zero {1 [[0]]
             2 (matr/matrix [[0 0] [0 0]])
             3 (matr/matrix [[0 0 0] [0 0 0] [0 0 0]])})

(defn empty-m [dim]
  "return a dim x dim matrix of zeroes"
  (if (<= 3 dim)
    (get m-zero dim)
    (matr/matrix (vec (take dim (cycle [(take dim (cycle [0]))]))))))


(defn m-sqrt [[[A B] [C D]]]
  "Calculates one square root of a 2x2 matrix"
  (let [[a b c d] (map #(Math/sqrt %) [A B C D])]
;    (println "Matrix:\n" [A B] \newline [C D] \newline a b c d)
        (cond
          (and (zero? B) (zero? C) (not (zero? A)) (not (zero? D)))
          (matr/matrix [[a 0] [0 d]])

          (and (zero? B) (every? #(not (zero? %)) [A C D]))
          (matr/matrix [[a 0] [(/ C (+ a d)) c]])

          (and (zero? C) (every? #(not (zero? %)) [A B D]))
          (matr/matrix [[a (/ B (+ a d))] [0 d]])

          :default
          (let [tau (+ A D)
                sigma (- (* A D) (* B C))]
            (cond
              (and (zero? sigma) (zero? tau))
              (empty-m 2)

              :default
              (let [s (Math/sqrt sigma)
                    t (Math/sqrt (+ tau (* 2 s)))]
                (println "s" s "t" t)
                (matr/div (matr/matrix [[(+ A s) B] [C (+ D s)]]) t)))))))


(defn- to-matrix [v & col]
  "Transform a vector to a matrix for required format"
  (let [n (count v)]
    (if col
      (matr/set-column (empty-m n) 0 v)
      (matr/set-row (empty-m n) 0 v))))

(defn com
  "Calculate center of mass"
  ([[x m]] (com x m))
  ([x m]
   {;^:doc "Calculate center of mass"
    :pre [(obj-valid? x m)]}
    (matr/mmul (reduce matr/add (map matr/mmul x m))
               (/ 1 (reduce + m)))))


(defn body [x m a]
  {:x0 x
   :x x
   :m m
   :t0 (com x m)
   :v (take (count m) (cycle [0]))
   :alpha a})


(defn- de-singularise [M]
  "Neccessary do avoid ugly singularities"
  (if-not (zero? (matr/det M))
    M
    (matr/matrix
     (for [i (range (count M))]
       (vec (for [j (range (count M))]
              (let [mij (get-in M [i j])]
                (if (and (= i j) (zero? mij)) 1e16
                    mij))))))))


(defn goal [{:keys [x x0 m t0]}]
  "calculate goal posititions gi for points x of body b"
  (let [t (com x m)
        p (map matr/sub x t)
        q (map matr/sub x0 t0)
        Apq (reduce matr/add
                    (map matr/mul
                         (map to-matrix p)
                         (map #(to-matrix % :transpose) q)))
        S (m-sqrt (matr/mmul (matr/transpose Apq) Apq))
        S (de-singularise S)
        R (matr/mmul Apq (matr/inverse S))
        t0 (matr/mmul t0 (matr/identity-matrix (count (first x))))
        delta (map #(matr/sub % t0) x)]
;    (println "R:"  R "delta:" delta)
    (map #(matr/add % t)
         (map #(matr/mmul R %) delta))))

(defn advance [{:keys [x v alpha] :as b} h]
  (let [g (goal b)
        delta (map matr/sub g x)
        delta (map #(matr/mmul (/ alpha h) %) delta)
        v (map matr/add v delta)
        v (map #(matr/mmul % h) v)
        x (map matr/add x v)]
    (assoc b :x x :v v)))

(deftest shape-matching
  (let [x0 [[0 0] [0 5] [5 5] [5 0]]
        x1 (matr/add 1 x0)
        m [1 1 1 1]
        b (assoc (body x0 m 0.5) :x x1)]
    (println "Body:\n" (clojure.string/join \newline b) \newline)
    (println (com x0 m))
    (println "goal:" (goal b))
    (loop [the-body b
           n 0]
      (when (< n 6)
        (println "step" n)
        (println (clojure.string/join \newline the-body))
        (recur (advance the-body 0.01)
               (inc n))))))