open Core
open Bistro
open Bistro.Shell_dsl
open File_formats

let assign k v =
  seq ~sep:"=" [ string k ; v ]

let conf_file_bppml ~tree ~faa ~out ~config =
  seq ~sep:"\n" (
    [
      assign "OUT" (out) ;
      assign "input.sequence.file" (dep faa) ;
      assign "alphabet" (string "Protein") ;
      assign "input.sequence.remove_stop_codons" (string "no") ;
      assign "input.sequence.sites_to_use" (string "all") ;
      assign "input.sequence.remove_saturated_sites" (string "yes") ;


      assign "input.tree.file"  (dep tree) ;
      assign "init.tree" (string "user") ;
      assign "input.tree.format" (string "Nhx") ;

      assign "optimization.topology" (string "false") ;

      assign "output.tree.file" (string "$(OUT)/tree.nhx") ;
      assign "output.tree.format" (string "Nhx") ;


      assign "output.infos" (string "$(OUT)/infos.tsv") ;
      assign "output.estimates" (string "$(OUT)/estimates.tsv") ;
    ]
    @ config
  )

let bppml ?(descr="") ~faa ~tree ~config : _ workflow =
  let config_f = dest // "config_bppml.bpp" in
  let out = ident dest in
  Workflow.shell ~descr:("bppsuite.bppml" ^ descr) [
    within_container Env.env_bppsuite (
      and_list [
        mkdir_p dest;
        cmd "cat" ~stdout:config_f [(file_dump (conf_file_bppml ~tree ~faa ~out ~config ))];
        cmd "bppml" [
          assign "param"  config_f;
        ]
      ]
    )
  ]

let topological ?(descr="") ~(tree:_ pworkflow) ~(tree_conv:_ pworkflow) ~(faa:aminoacid_fasta pworkflow) ~prot_model () : [`topological] dworkflow =
  let config = [assign "model" (string prot_model)] in
  let run_bppml = bppml ~descr:"" ~tree ~config ~faa in
  let run_bppml_conv = bppml ~descr:".conv" ~tree:tree_conv ~config ~faa in
  let bppml_config = Workflow.select run_bppml ["config_bppml.bpp"] in
  let bppml_config_conv = Workflow.select run_bppml_conv ["config_bppml.bpp"] in
  let bppml_out = Workflow.select run_bppml ["infos.tsv"] in
  let bppml_out_conv = Workflow.select run_bppml_conv ["infos.tsv"] in
  let out = dest // "out.tsv" in
  Workflow.shell ~descr:("topological.parse_"^prot_model^"."^descr) [
    mkdir dest ;
    cmd "cp" [dep bppml_out ; dest // "estimates.bppml_out.tsv" ];
    cmd "cp" [dep bppml_out_conv ; dest // "estimates.bppml_out_conv.tsv" ];
    cmd "cp" [dep bppml_config ; dest // "estimates.bppml_config" ];
    cmd "cp" [dep bppml_config_conv ; dest // "estimates.bppml_config_conv" ];
    cmd "python" ~img:Env.env_pcoc [
      file_dump (string Scripts.calc_topological) ;
      opt "-bppml_non_conv" dep bppml_out ;
      opt "-bppml_conv" dep bppml_out_conv ;
      opt "-o"  ident out ;
    ]
  ]

let results run_topological : text_file pworkflow =
  Workflow.select run_topological ["out.tsv"]