type t = 
  | Bohr of float array
  | Angstrom of float array

(** Bohr radius, taken from https://physics.nist.gov/cgi-bin/cuu/Value?bohrrada0 *)
let a0 = 
  0.529_177_210_67

let zero = Bohr [| 0. ; 0. ; 0. |]

let of_float_triplet (x,y,z) = function
| `Bohr     -> Bohr     [|x;y;z|]
| `Angstrom -> Angstrom [|x;y;z|]

let of_3_floats x y z =
  of_float_triplet (x,y,z)

let to_string y =
  let result x = 
    (string_of_float x.(0))^" "^(string_of_float x.(1))^" "^(string_of_float x.(2))
  in
  match y with
  | Bohr     x -> (result x) ^ " Bohr"
  | Angstrom x -> (result x) ^ " Angstrom"


let to_float_array = function
  | Bohr     a
  | Angstrom a -> a

let x a = (to_float_array a).(0)
let y a = (to_float_array a).(1)
let z a = (to_float_array a).(2)


let coord a = function
  | 0 -> (to_float_array a).(0)
  | 1 -> (to_float_array a).(1)
  | 2 -> (to_float_array a).(2)
  | _ -> raise (Invalid_argument "Coordinate")


(** Linear algebra *)
let (|.) s a =
  match a with
  | Bohr     [|x;y;z|] -> Bohr     [| s*.x; s*.y; s*.z |]
  | Angstrom [|x;y;z|] -> Angstrom [| s*.x; s*.y; s*.z |]
  | _ -> assert false

let to_Angstrom = function
  | Angstrom a -> Angstrom a
  | Bohr     a -> Angstrom (a0 |. Bohr a |> to_float_array)

let to_Bohr = function
  | Angstrom a -> Bohr (1./.a0 |. Angstrom a |> to_float_array)
  | Bohr     a -> Bohr a 

let (|-), (|+) =
  let rec op f p q =
    match (p, q) with
    | (Angstrom a, Angstrom b) -> Angstrom (f a b)
    | (Bohr     a, Bohr     b) -> Bohr     (f a b)
    | (Angstrom a, Bohr     b) -> op f (to_Bohr p) q
    | (Bohr     a, Angstrom b) -> op f p (to_Bohr q)
  in
  (op (fun a b -> 
       match a,b with
       | [|x;y;z|], [|x';y';z'|] -> [| x-.x'; y-.y'; z-.z' |]
       | _ -> assert false
     ) , 
   op (fun a b -> 
       match a,b with
       | [|x;y;z|], [|x';y';z'|] -> [| x+.x'; y+.y'; z+.z' |]
       | _ -> assert false
     )
  )


let dot =
  let rec op f p q =
    match (p, q) with
    | (Angstrom a, Angstrom b) 
    | (Bohr     a, Bohr     b) -> f a b
    | (Angstrom a, Bohr     b) -> op f (to_Bohr p) q
    | (Bohr     a, Angstrom b) -> op f p (to_Bohr q)
  in
  op (fun a b ->
      match a,b with
      | [|x;y;z|], [|x';y';z'|] ->  x*.x' +. y*.y' +. z*.z' 
      | _ -> assert false
    )


let norm u =
  sqrt @@ dot u u