open Common

type single_exc =
  {
    hole     : int    ;
    particle : int    ;
    spin     : Spin.t ;
  }

type t =
| Identity  of Phase.t
| Single    of Phase.t * single_exc
| Double    of Phase.t * single_exc * single_exc
| Triple    of Phase.t * single_exc * single_exc * single_exc
| Multiple  of Phase.t * single_exc list


let single_of_spindet t t' =
  assert (Spindeterminant.excitation_level t t' = 1);
  let d  = Spindeterminant.bitstring t
  and d' = Spindeterminant.bitstring t'
  in
  let tmp = Bitstring.logxor d d' in
  let shift_left_one = Bitstring.(shift_left_one (numbits tmp)) in
  let hole_z     = Bitstring.logand (Spindeterminant.bitstring t ) tmp
  and particle_z = Bitstring.logand (Spindeterminant.bitstring t') tmp
  in
  let hole     = 1 + Bitstring.trailing_zeros hole_z
  and particle = 1 + Bitstring.trailing_zeros particle_z
  in
  (* Phase calculation *)
  let low, high =
    if particle > hole then hole, particle
    else particle, hole
  in
  let mask =
    let h = high-1 in
    let l = low in
    let mask_up = shift_left_one h |> Bitstring.minus_one
    and mask_dn = Bitstring.plus_one @@ Bitstring.lognot (shift_left_one l)
    in Bitstring.logand mask_up mask_dn
  in
  let phase =
    Phase.multiply (Phase.multiply (Spindeterminant.phase t) (Spindeterminant.phase t'))
      (Phase.of_nperm (Bitstring.popcount @@ Bitstring.logand d mask ))
  in
  (hole, particle, phase)


let single_of_det t t' =
  assert Determinant.(beta t = beta t' || alfa t = alfa t');
  if Determinant.(beta t = beta t') then
    let hole, particle, phase =
      single_of_spindet (Determinant.alfa t) (Determinant.alfa t')
    in
    Single (phase, { hole ; particle ; spin=Spin.Alfa })
  else
    let hole, particle, phase =
      single_of_spindet (Determinant.beta t) (Determinant.beta t')
    in
    Single (phase, { hole ; particle ; spin=Spin.Beta })


let multiple_of_spindet t t' =
  let holes     = Spindeterminant.holes_of t t'
  and particles = Spindeterminant.particles_of t t'
  in
  let t'' =
    List.fold_left (fun accu h -> Spindeterminant.annihilation h accu) t holes
  in
  let t'' =
    List.fold_left (fun accu p -> Spindeterminant.creation p accu) t'' particles
  in
  assert (t' = t'' || t' = Spindeterminant.negate_phase t'');
  let phase =
    if Spindeterminant.phase t' = Spindeterminant.phase t'' then
      Phase.Pos
    else
      Phase.Neg
  in
  (phase, List.rev @@ List.rev_map2 (fun hole particle -> (hole, particle)) holes (List.rev particles) )


let double_of_spindet t t' =
  match multiple_of_spindet t t' with
  | (phase, (h1,p1)::(h2,p2)::[]) -> (h1, p1, h2, p2, phase)
  | _ -> invalid_arg "t and t' are not doubly excited"


let triple_of_spindet t t' =
  match multiple_of_spindet t t' with
  | (phase, (h1,p1)::(h2,p2)::(h3,p3)::[]) -> (h1, p1, h2, p2, h3, p3, phase)
  | _ -> invalid_arg "t and t' are not doubly excited"


let multiple_of_det t t' =
  let pa, a =
    multiple_of_spindet (Determinant.alfa t) (Determinant.alfa t')
  and pb, b =
    multiple_of_spindet (Determinant.beta t) (Determinant.beta t')
  in
  let phase = Phase.multiply pa pb in
  Multiple (phase, List.concat [
    List.rev @@ List.rev_map (fun (hole, particle) -> { hole ; particle ; spin=Spin.Alfa }) a ;
    List.rev @@ List.rev_map (fun (hole, particle) -> { hole ; particle ; spin=Spin.Beta }) b ])


let double_of_det t t' =
  match multiple_of_det t t' with
  | Multiple (phase, [e1 ; e2]) -> Double (phase, e1, e2)
  | _ -> assert false


let triple_of_det t t' =
  match multiple_of_det t t' with
  | Multiple (phase, [e1 ; e2 ; e3]) -> Triple (phase, e1, e2, e3)
  | _ -> assert false


let of_det t t' =
  match Determinant.excitation_level t t' with
  | 0 -> if Determinant.phase t = Determinant.phase t' then
          Identity Phase.Pos
        else
          Identity Phase.Neg
  | 1 -> single_of_det t t'
  | 2 -> double_of_det t t'
  | 3 -> triple_of_det t t'
  | _ -> multiple_of_det t t'

let pp_s_exc ppf t =
  Format.fprintf ppf "@[T^{%s}_{%d->%d}@]"
    (match t.spin with
      | Spin.Alfa -> "\\alpha"
      | Spin.Beta -> "\\beta " )
    t.hole t.particle

let pp ppf t =
  let phase, l =
    match t with
    | Identity p -> p, []
    | Single (p,x) -> p, x::[]
    | Double (p,x,y) -> p, x::y::[]
    | Triple (p,x,y,z) -> p, x::y::z::[]
    | Multiple (p,l) -> p, l
  in
  Format.fprintf ppf "@[%c"
    (match phase with
      | Phase.Pos -> '+'
      | Phase.Neg -> '-' );
  List.iter (fun x -> Format.fprintf ppf "@[T^{%s}_{%d->%d}@]"
                      (match x.spin with
                        | Spin.Alfa -> "\\alpha"
                        | Spin.Beta -> "\\beta " )
                      x.hole x.particle) l;
  Format.fprintf ppf "@]"