open Core
open Bistro
open File_formats

let%pworkflow multinomial_ocaml_implementation ~meth ~(tree_sc:_ file) ~(faa:aminoacid_fasta file) (* : cpt file *) =
  let open Phylogenetics in
  let module MT = Phylogenetics_convergence.Multinomial_test in
  let open Codepitk in
  let meth = [%param meth] in
  let test = match meth with
    | `Asymptotic_LRT -> MT.LRT.asymptotic_test
    | `Simulation_LRT -> MT.LRT.simulation_test ~sample_size:10_000
    | `Asymptotic_sparse -> MT.Sparse.asymptotic_test
    | `Simulation_sparse -> MT.Sparse.simulation_test ~sample_size:10_000
  in
  let alignment =
    Alignment.from_fasta [%path faa]
    |> Rresult.R.get_ok
  in
  let tree =
    Convergence_tree.from_file [%path tree_sc]
    |> Rresult.R.failwith_error_msg
  in
  let site c0 c1 =
    let c0 = (c0 : int Amino_acid.table :> int array) in
    let c1 = (c1 : int Amino_acid.table  :> int array) in
    let d = MT.data ~x1:c0 ~x2:c1 in
    let r = test d in
    r._T_, r.pvalue
  in
  let res = Convergence_tree.alignment_counts_map tree alignment site in
  let res_lines =
    List.mapi res ~f:(fun i (lr, pval) ->
        sprintf "%d\t%f\t%f" (i + 1) (1. -. pval) lr
      )
  in
  let header = "Sites\tMultinomial_1mp\tMultinomial_LikelihoodRatio" in
  Out_channel.write_lines [%dest] (header :: res_lines)

let multinomial_asymptotic_lrt ~tree_sc ~faa = multinomial_ocaml_implementation ~meth:`Asymptotic_LRT ~tree_sc ~faa
let multinomial_simulation_lrt ~tree_sc ~faa = multinomial_ocaml_implementation ~meth:`Simulation_LRT ~tree_sc ~faa
let multinomial_asymptotic_sparse ~tree_sc ~faa = multinomial_ocaml_implementation ~meth:`Asymptotic_sparse ~tree_sc ~faa
let multinomial_simulation_sparse ~tree_sc ~faa = multinomial_ocaml_implementation ~meth:`Simulation_sparse ~tree_sc ~faa