open Powers

exception AngularMomentumError of string

type t =
  | S | P | D | F | G | H | I | J | K | L | M | N | O

let of_char = function
  | 's' | 'S' -> S      | 'p' | 'P' -> P
  | 'd' | 'D' -> D      | 'f' | 'F' -> F
  | 'g' | 'G' -> G      | 'h' | 'H' -> H
  | 'i' | 'I' -> I      | 'j' | 'J' -> J
  | 'k' | 'K' -> K      | 'l' | 'L' -> L
  | 'm' | 'M' -> M      | 'n' | 'N' -> N
  | 'o' | 'O' -> O
  | c -> raise (AngularMomentumError (String.make 1 c))

let to_string = function
  | S -> "S"       | P -> "P"
  | D -> "D"       | F -> "F"
  | G -> "G"       | H -> "H"
  | I -> "I"       | J -> "J"
  | K -> "K"       | L -> "L"
  | M -> "M"       | N -> "N"
  | O -> "O"

let to_char = function
  | S -> 'S'       | P -> 'P'
  | D -> 'D'       | F -> 'F'
  | G -> 'G'       | H -> 'H'
  | I -> 'I'       | J -> 'J'
  | K -> 'K'       | L -> 'L'
  | M -> 'M'       | N -> 'N'
  | O -> 'O'

let to_int = function
  | S -> 0         | P -> 1
  | D -> 2         | F -> 3
  | G -> 4         | H -> 5
  | I -> 6         | J -> 7
  | K -> 8         | L -> 9
  | M -> 10        | N -> 11
  | O -> 12

let of_int = function
  | 0 -> S         | 1 -> P
  | 2 -> D         | 3 -> F
  | 4 -> G         | 5 -> H
  | 6 -> I         | 7 -> J
  | 8 -> K         | 9 -> L
  | 10 -> M        | 11 -> N
  | 12 -> O
  | c -> raise (AngularMomentumError (string_of_int c))



type kind =
| Singlet of  t
| Doublet of (t*t)
| Triplet of (t*t*t)
| Quartet of (t*t*t*t)


let n_functions a =
  let a = 
    to_int a
  in
  (a*a + 3*a + 2)/2


(** Returns an array of Zkeys corresponding to all possible angular momenta *)
let zkey_array a = 
  let keys_1d l =
    let create_z { x ; y ; _ } =
      Powers.of_int_tuple (x,y,l-(x+y))
    in
    let rec create_y accu xyz =
      let { x ; y ; z } = xyz in
      match y with
      | 0 -> (create_z xyz)::accu
      | i -> let ynew = y-1 in
        create_y ( (create_z xyz)::accu) (Powers.of_int_tuple (x,ynew,z))
    in
    let rec create_x accu xyz =
      let { x ; y ; z } = xyz in
      match x with
      | 0 -> (create_y [] xyz)@accu
      | i -> let xnew = x-1 in
        let ynew = l-xnew in
        create_x ((create_y [] xyz)@accu) (Powers.of_int_tuple (xnew, ynew, z))
    in
    create_x [] (Powers.of_int_tuple (l,0,0))
    |> List.rev
  in

  begin
    match a with
    | Singlet l1 -> 
      List.map (fun x -> Zkey.of_powers (Zkey.Three x)) (keys_1d @@ to_int l1)

    | Doublet (l1, l2) -> 
      List.map (fun a ->
          List.map (fun b ->
              Zkey.of_powers (Zkey.Six (a,b))) (keys_1d @@ to_int l2)
        ) (keys_1d @@ to_int l1)
      |> List.concat

    | Triplet (l1, l2, l3) ->

      List.map (fun a ->
          List.map (fun b ->
              List.map (fun c ->
                  Zkey.of_powers (Zkey.Nine (a,b,c))) (keys_1d @@ to_int l3)
            ) (keys_1d @@ to_int l2)
          |> List.concat
        ) (keys_1d @@ to_int l1)
      |> List.concat

    | Quartet (l1, l2, l3, l4) ->

      List.map (fun a ->
          List.map (fun b ->
              List.map (fun c ->
                  List.map (fun d ->
                      Zkey.of_powers (Zkey.Twelve (a,b,c,d))) (keys_1d @@ to_int l4)
                ) (keys_1d @@ to_int l3)
              |> List.concat
            ) (keys_1d @@ to_int l2)
          |> List.concat
        ) (keys_1d @@ to_int l1)
        |> List.concat
  end
  |> Array.of_list