open Lacaml.D

module De = Determinant
module Ex = Excitation
module Sp = Spindeterminant

type t = float list


let non_zero integrals degree_a degree_b ki kj =
  let kia = De.alfa ki and kib = De.beta ki
  and kja = De.alfa kj and kjb = De.beta kj
  in
  let single h p spin same opposite = 
    let same_spin_mo_list = 
      Sp.to_list same
    and opposite_spin_mo_list = 
      Sp.to_list opposite
    in
    fun one_e two_e ->
      let same_spin = 
        List.fold_left (fun accu i -> accu +. two_e h i p i spin spin) 0. same_spin_mo_list
      and opposite_spin = 
        List.fold_left (fun accu i -> accu +. two_e h i p i spin (Spin.other spin) ) 0. opposite_spin_mo_list
      in (one_e h p spin) +. same_spin +. opposite_spin
  in
  let diag_element =
    let mo_a = Sp.to_list kia
    and mo_b = Sp.to_list kib
    in
    fun one_e two_e ->
      let one =
        (List.fold_left (fun accu i -> accu +. one_e i i Spin.Alfa) 0. mo_a)
        +.
        (List.fold_left (fun accu i -> accu +. one_e i i Spin.Beta) 0. mo_b)
      in
      let two =
        let rec aux_same spin accu = function
          | [] -> accu
          | i :: rest -> 
            let new_accu = 
              List.fold_left (fun accu j -> accu +. two_e i j i j spin spin) accu rest
            in
            aux_same spin new_accu rest
        in
        let rec aux_opposite accu other = function
          | [] -> accu
          | i :: rest ->
            let new_accu =
              List.fold_left (fun accu j -> accu +. two_e i j i j Spin.Alfa Spin.Beta) accu other
            in
            aux_opposite new_accu other rest
        in
        (aux_same Spin.Alfa 0. mo_a)  +.  (aux_same Spin.Beta 0. mo_b)  +.
        (aux_opposite 0. mo_a mo_b)
      in
      one +. two
  in
  let result =
    match degree_a, degree_b with
    | 1, 1 -> (* alpha-beta double *)
      begin
        let ha, pa, phase_a = Ex.single_of_spindet kia kja in
        let hb, pb, phase_b = Ex.single_of_spindet kib kjb in
        match phase_a, phase_b with
        | Phase.Pos, Phase.Pos
        | Phase.Neg, Phase.Neg -> fun _ two_e ->    two_e ha hb pa pb Spin.Alfa Spin.Beta
        | Phase.Neg, Phase.Pos                       
        | Phase.Pos, Phase.Neg -> fun _ two_e -> -. two_e ha hb pa pb Spin.Alfa Spin.Beta
      end

    | 2, 0 -> (* alpha double *)
      begin
        let h1, p1, h2, p2, phase = Ex.double_of_spindet kia kja in
        match phase with 
        | Phase.Pos -> fun _ two_e ->    two_e h1 h2 p1 p2 Spin.Alfa Spin.Alfa
        | Phase.Neg -> fun _ two_e -> -. two_e h1 h2 p1 p2 Spin.Alfa Spin.Alfa
      end

    | 0, 2 -> (* beta double *)
      begin
        let h1, p1, h2, p2, phase = Ex.double_of_spindet kib kjb in
        match phase with 
        | Phase.Pos -> fun _ two_e ->    two_e h1 h2 p1 p2 Spin.Beta Spin.Beta
        | Phase.Neg -> fun _ two_e -> -. two_e h1 h2 p1 p2 Spin.Beta Spin.Beta
      end

    | 1, 0 -> (* alpha single *)
      begin
        let h, p, phase = Ex.single_of_spindet kia kja in
        match phase with
        | Phase.Pos -> fun one_e two_e ->    single h p Spin.Alfa kia kib one_e two_e
        | Phase.Neg -> fun one_e two_e -> -. single h p Spin.Alfa kia kib one_e two_e 
      end

    | 0, 1 -> (* beta single *)
      begin
        let h, p, phase = Ex.single_of_spindet kib kjb in
        match phase with
        | Phase.Pos -> fun one_e two_e ->    single h p Spin.Beta kib kia one_e two_e
        | Phase.Neg -> fun one_e two_e -> -. single h p Spin.Beta kib kia one_e two_e
      end

    | 0, 0 -> (* diagonal element *)
      diag_element

    | _ -> assert false
  in
  List.map (fun (one_e, two_e) -> result one_e two_e) integrals


let make integrals ki kj =
  let degree_a, degree_b =
    De.degrees ki kj 
  in
  if degree_a+degree_b > 2 then
    List.map (fun _ -> 0.) integrals
  else
    non_zero integrals degree_a degree_b ki kj
    



let make_s2 ki kj =
  let degree_a = De.degree_alfa ki kj in
  let kia = De.alfa ki in
  let kja = De.alfa kj in
  if degree_a > 1 then 0.
  else
    let degree_b = De.degree_beta ki kj in
    let kib = De.beta ki in
    let kjb = De.beta kj in
    match degree_a, degree_b with
    | 1, 1 -> (* alpha-beta double *)
      let ha, pa, phase_a = Ex.single_of_spindet kia kja in
      let hb, pb, phase_b = Ex.single_of_spindet kib kjb in
      if ha = pb && hb = pa then
        begin
          match phase_a, phase_b with
          | Phase.Pos, Phase.Pos
          | Phase.Neg, Phase.Neg -> -1.
          | Phase.Neg, Phase.Pos                       
          | Phase.Pos, Phase.Neg -> 1.
        end
      else 0.
    | 0, 0 -> 
      let ba = Sp.bitstring kia and bb = Sp.bitstring kib in
      let tmp = Z.(logxor ba bb) in
      let popcount x =
        if x = Z.zero then 0 else
          Z.popcount x
      in
      let n_a = Z.(logand ba tmp) |> popcount in
      let n_b = Z.(logand bb tmp) |> popcount in
      let s_z = 0.5 *. float_of_int (n_a - n_b) in
      float_of_int n_a +. s_z *. (s_z -. 1.)
    | _ -> 0.