pipeline.ml

open Core
open Bistro_utils
open Bistro.EDSL
open Bistro.Std
open File_formats

let ready_dataset_of_raw_dataset raw_dataset =
    let input_tree = raw_dataset.input_tree in
    let fna = raw_dataset.fna in
    let tree_dataset = Tree_dataset.prepare input_tree in
    let faa = Bppsuite.fna2faa ~fna in
    { input_tree; tree_dataset ; fna; faa}

let raw_dataset_of_ready_dataset ready_dataset =
    let input_tree = ready_dataset.input_tree in
    let fna = ready_dataset.fna in
    { input_tree; fna}

let repo_of_ready_dataset (rd : ready_dataset) =
  Repo.[
    item ["input_tree.nhx"] rd.input_tree ;
    item ["tree.H0.node_ids" ] (Tree_dataset.nodes rd.tree_dataset H0) ;
    item ["tree.Ha.node_ids" ] (Tree_dataset.nodes rd.tree_dataset Ha) ;
    item ["tree.only_convergent_tags.nhx" ] (Tree_dataset.tree rd.tree_dataset `Detection) ;
    item ["tree.only_node_ids.nhx" ] (Tree_dataset.tree rd.tree_dataset `Simulation) ;
    item ["tree.diffsel" ] (Tree_dataset.diffsel_tree rd.tree_dataset) ;
    item ["simulated_sequences.fna"] rd.fna ;
    item ["simulated_sequences.faa"] rd.faa ;
  ]
  |> Repo.shift "ready_dataset"

let repo_of_raw_dataset (raw_dataset:raw_dataset) =
  Repo.[
    item ["input_tree.nhx"] raw_dataset.input_tree ;
    item ["simulated_sequences.fna"] raw_dataset.fna ;
  ]
  |> Repo.shift "raw_dataset"


let repo_of_dataset_l ~preview dataset_l =
  List.map dataset_l ~f:(fun dataset ->
    let model_prefix = dataset.model_prefix in
    let tree_prefix = dataset.tree_prefix in
    let w_dataset =
      if preview then
        Ready_dataset dataset.ready_dataset
      else Raw_dataset (raw_dataset_of_ready_dataset dataset.ready_dataset) in
    let repo_d = match w_dataset with
      | Ready_dataset d -> repo_of_ready_dataset d
      | Raw_dataset d ->  repo_of_raw_dataset d
    in
      Repo.shift tree_prefix (Repo.shift model_prefix repo_d)
    )
    |> List.concat

let derive_from_model ~model ~tree ~tree_dataset ~tree_prefix ~preview =
  let nb_sites = if preview then 10 else 100 in
  let nodes = Tree_dataset.nodes tree_dataset model in
  let config = Convergence_hypothesis.bpp_config nodes model in
  let tree = Tree_dataset.tree tree_dataset `Simulation in
  let fna = Bppsuite.bppseqgen ~nb_sites ~tree ~config in
  let faa = Bppsuite.fna2faa ~fna in
  let ready_dataset = { input_tree = tree ; tree_dataset ; fna; faa} in
  let model_prefix = Convergence_hypothesis.string_of_model model in
    { model_prefix; tree_prefix; ready_dataset }

let derive_from_tree ~tree_dir ~tree ~preview =
  let tree_prefix = Filename.chop_extension tree in
  let tree = input (Filename.concat tree_dir tree) in
  let tree_dataset = Tree_dataset.prepare tree in
  let models = Convergence_hypothesis.[H0; Ha] in
  List.map models ~f:(fun model ->
      derive_from_model ~model ~tree ~tree_dataset ~tree_prefix ~preview
    )

let derive_sim ~tree_dir ~trees ~preview =
  List.map trees ~f:(fun tree ->
      derive_from_tree ~tree_dir ~tree ~preview)
  |> List.concat


let repo_of_det_results_l ~det_results_l =
  List.map det_results_l ~f:(fun det_results ->
    let model_prefix = det_results.dataset.model_prefix in
    let tree_prefix = det_results.dataset.tree_prefix in
    let det_meth_prefix = string_of_det_meth det_results.det_meth in
    let open Convergence_detection in
    let w = det_results.det_result
      in
    let repo_d = Repo.[
      item [det_meth_prefix] (w (*/ selector ["RUN*"]*))  ]
    in
    Repo.shift tree_prefix (Repo.shift model_prefix repo_d)
  )
  |> List.concat

let derive_from_det_meth ~det_meth ~dataset ~preview =
  let open Convergence_detection in
  let faa = dataset.ready_dataset.faa in
  let fna = dataset.ready_dataset.fna in
  let phy_n = Bppsuite.fa2phy ~fna in
  let pcoc_tree = Tree_dataset.tree dataset.ready_dataset.tree_dataset `Detection in
  let diffsel_tree = Tree_dataset.diffsel_tree dataset.ready_dataset.tree_dataset in
  let w_every = if preview then 1 else 10 in
  let n_cycles = if preview then 100 else 1000 in
  let det_result = match det_meth with
    | Pcoc -> Convergence_detection.pcoc ~gamma:false ~faa ~tree:pcoc_tree
    | Pcoc_gamma -> Convergence_detection.pcoc ~gamma:true ~faa ~tree:pcoc_tree
    | Diffsel -> Convergence_detection.diffsel ~phy_n ~tree:diffsel_tree ~w_every ~n_cycles
  in
  {det_meth; det_result; dataset}




let derive_from_dataset ~dataset ~preview=
  let det_meths = [Pcoc;Pcoc_gamma;Diffsel] in
  List.map det_meths ~f:(fun det_meth ->
    derive_from_det_meth ~det_meth ~dataset ~preview
    )

let derive_det ~dataset_l ~preview=
  List.map dataset_l ~f:(fun dataset ->
      derive_from_dataset ~preview ~dataset)
  |> List.concat

let main ~outdir ?(np = 2) ?(mem = 2) ~tree_dir ~preview () =
 let logger =
    Logger.tee [
      Console_logger.create () ;
      Dot_output.create "dag.dot" (*dot -Tpdf example/dag.dot -o dag.pdf*)
    ] in
  let trees = Array.to_list @@ Sys.readdir tree_dir in
  let dataset_l = derive_sim ~tree_dir ~trees ~preview in
  let det_results_l = derive_det ~dataset_l ~preview in
  let repo = [
    repo_of_dataset_l dataset_l ~preview ;
    repo_of_det_results_l det_results_l;
  ]
  |> List.concat
  in
  Repo.build ~outdir ~np ~mem:(`GB mem) ~logger repo

let command =
  let open Command.Let_syntax in
  Command.basic
    ~summary:"Run simulation pipeline"
    [%map_open
      let outdir =
        flag "--outdir" (required string) ~doc:"PATH Output directory"
      and preview =
        flag "--preview-mode" no_arg ~doc:" Preview mode"
      and np =
        flag "--np" (optional int) ~doc:"INT Number of available processors"
      and mem =
        flag "--mem" (optional int) ~doc:"INT Available memory (in GB)"
      and tree_dir =
        flag "--tree-dir" (required string) ~doc:"PATH Path to tree directory"
      in
      main ~outdir ?np ?mem ~tree_dir ~preview
    ]