open Lacaml.D

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

type t =
  {
    det_space   : Determinant_space.t ;
    h_matrix    : Mat.t lazy_t ;
    eigensystem : (Mat.t * Vec.t) lazy_t;
  }

let det_space    t = t.det_space

let h_matrix     t = Lazy.force t.h_matrix

let eigensystem  t = Lazy.force t.eigensystem 

let eigenvectors t = 
  let (x,_) = eigensystem t in x

let eigenvalues t = 
  let (_,x) = eigensystem t in x


let h_ij mo_basis ki kj =
  let one_e_ints  = MOBasis.one_e_ints mo_basis
  and two_e_ints  = MOBasis.two_e_ints mo_basis
  and degree_a, degree_b = De.degrees ki kj 
  in
  if degree_a+degree_b > 2 then 
    0.
  else
    begin
      let kia = De.alfa ki and kib = De.beta ki
      and kja = De.alfa kj and kjb = De.beta kj
      in
      let integral =
        ERI.get_phys two_e_ints
      in
      let anti_integral i j k l =
        integral i j k l -. integral i j l k
      in
      let single h p same opposite = 
        let same_spin = 
          Sp.to_list same
          |> List.fold_left (fun accu i -> accu +. anti_integral h i p i) 0.
        and opposite_spin = 
          Sp.to_list opposite
          |> List.fold_left (fun accu i -> accu +. integral h i p i) 0.
        and one_e =
          one_e_ints.{h,p}
        in one_e +. same_spin +. opposite_spin
      in
      let diag_element () =
        let mo_a = Sp.to_list kia
        and mo_b = Sp.to_list kib
        in
        let one_e =
          List.concat [mo_a ; mo_b]
          |> List.fold_left (fun accu i -> accu +. one_e_ints.{i,i}) 0. 
        and two_e =
          let      two_index i j =      integral i j i j
          and anti_two_index i j = anti_integral i j i j
          in
          let rec aux_same accu = function
            | [] -> accu
            | i :: rest -> 
              let new_accu = 
                List.fold_left (fun accu j -> accu +. anti_two_index i j) accu rest
              in
              aux_same 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_index i j) accu other
              in
              aux_opposite new_accu other rest
          in
          (aux_same 0. mo_a)  +.  (aux_same 0. mo_b)  +.  (aux_opposite 0. mo_a mo_b)
        in
        one_e +. two_e
      in
        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 -> +. integral ha hb pa pb
            | Phase.Neg, Phase.Pos
            | Phase.Pos, Phase.Neg -> -. integral ha hb pa pb
          end

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

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

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

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

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

        | _ -> assert false
    end


let make det_space = 
  let ndet = Determinant_space.size det_space in
  let det  = Determinant_space.determinants det_space in
  let mo_basis = Determinant_space.mo_basis det_space in
  let h_matrix = lazy (
    Array.init ndet (fun i ->
        let ki = det.(i) in
        Array.init ndet (fun j ->
            let kj = det.(j) in
            h_ij mo_basis ki kj
          )
      )
   |> Mat.of_array
  )
  in
  let eigensystem = lazy (
    Lazy.force h_matrix
    |> Util.diagonalize_symm 
  )
  in
  { det_space ; h_matrix ; eigensystem }