open Phylogenetics
open Codepitk.Tdg09
open Core_bench

let () =
  let stationary_distribution = Amino_acid.random_profile 0.5 in
  let exchangeability_matrix = Rate_matrix.Amino_acid.make (fun aa_i aa_j->
      let i, j = Amino_acid.(to_int aa_i, to_int aa_j) in
      if i <= j then float_of_int i *. 0.01 +. float_of_int j *. 0.001
      else float_of_int j *. 0.01 +. float_of_int i *. 0.001
    ) in
  let p = {
    stationary_distribution;
    Evolution_model.exchangeability_matrix;
    scale=1.;
  } in
  let t = 100. in
  let transition_matrix_diag = Evolution_model.transition_probability_matrix p in
  let transition_matrix_exp =
    let m = Evolution_model.rate_matrix p in
    fun t -> Amino_acid.Matrix.(expm (scal_mul t m))
  in
  Core.Command.run (Bench.make_command [
      Bench.Test.create ~name:"diag" (fun () -> transition_matrix_diag t) ;
      Bench.Test.create ~name:"exp" (fun () -> transition_matrix_exp t) ;
    ])