exception Exception;;

let get_str num =
  match num with
    | 0 -> 'A' | 1 -> 'V' | 2 -> 'L' | 3 -> 'I' 
    | 4 -> 'P' | 5 -> 'F' | 6 -> 'M' | 7 -> 'W' 
    | 8 -> 'G' | 9 -> 'C' | 10 -> 'D' | 11 -> 'E'
    | 12 -> 'N' | 13 -> 'Q' | 14 -> 'S' | 15 -> 'T'
    | 16 -> 'Y' | 17 -> 'K' | 18 -> 'R' | 19 -> 'H' 
    | _ -> raise Exception;;

let (++) str1 str2 =
  let a = String.create ((String.length str1)+(String.length str2)) in
    for i=1 to (String.length str1) do
      a.[i-1] <- str1.[i-1]
    done;
    for i=(String.length str1) to (String.length str1)+(String.length str2)-1 do
      a.[i] <- str2.[i-(String.length str1)]
    done;
    a;;



let gen_random_string length =
  let str = String.create length in
    for i=1 to length do
      let num = Random.int 20 in
	str.[i-1] <- (get_str num)
    done;
    str;;

let mutate str num =
  let mutate_one str =
    let new_str = String.copy str in 
    let a = Random.int (String.length str) in
      new_str.[a] <- (get_str (Random.int 20));
      new_str in
  let new_str = ref (String.copy str) in
    for i=1 to num do
      new_str := (mutate_one !new_str)
    done;
    !new_str;;

let insert str num =
  let insert_one str =
    let new_str = String.create ((String.length str)+1) in
    let a = Random.int (String.length new_str) in
      for i=0 to (a-1) do
	new_str.[i] <- str.[i]
      done;
      new_str.[a] <- (get_str (Random.int 20));
      for i=(a+1) to ((String.length new_str)-1) do
	new_str.[i] <- str.[i-1]
      done;
      new_str in
  let new_str = ref (String.copy str) in
    for i=1 to num do
      new_str := (insert_one !new_str)
    done;
    !new_str;;

let delete str num =
  let delete_one str =
    let new_str = String.create ((String.length str)-1) in
    let a = Random.int (String.length str) in
      for i=0 to (a-1) do
	new_str.[i] <- str.[i]
      done;
      for i=a to ((String.length new_str)-1) do
	new_str.[i] <- str.[i+1]
      done;
      new_str in
  let new_str = ref (String.copy str) in
    for i=1 to num do
      new_str := (delete_one !new_str) 
    done;
    !new_str;;




let change str num =
  let change_one str =
    let a = Random.int 3 in
      match a with
	  0 -> (mutate str 1)
	| 1 -> (insert str 1)
	| 2 -> (delete str 1) 
	| _ -> raise Exception in
  let new_str = ref (String.copy str) in
    for i=1 to num do
      new_str := (change_one !new_str)
    done;
    !new_str;;

(* parameter *)

let score = 2;;  (* agreement score *)
let cost = -1;;  (* disagreement cost *)
let gap_init = -2;; (* gap initiation cost *)
let gap_ext = -1;; (* gap extension cost *)


(* alignment algorithm *)

(* global alignment : Needleman-Wunsch *)

let global_alignment str1 str2 =
  let a = (String.length str1) in
  let b = (String.length str2) in
  let mat1 = Array.make_matrix a b 0 in
  let mat2 = Array.make_matrix a b 0 in 
    (* arg1: F(i-1,j-1) val1: x_i val2:y_i arg2:F(i-1,j) arg3:F(i,j-1) *)
  let get_max_value arg1 val1 val2 arg2 arg3 =
    let tmp1 = ref 0 in
    let v2 = arg2 + gap_init in
    let v3 = arg3 + gap_init in
      if val1=val2 then
	tmp1 := arg1 + score
      else 
	tmp1 := arg1 + cost;
      let v1 = !tmp1 in
      let ret = ref (0,0) in
	if (v1 >= v2) && (v1>=v3) then
	  ret := (v1,1)
	else if (v2 >=v1) && (v2>=v3) then
	  ret := (v2,2) 
	else if (v3>=v1) && (v3>=v2) then
	  ret := (v3,3);
	!ret in
    for i=1 to a do
      mat1.(i-1).(0) <- (i-1)*gap_init;
      mat2.(i-1).(0) <- 2;
    done;
    for i=1 to b do
      mat1.(0).(i-1) <- (i-1)*gap_init;
      mat2.(0).(i-1) <- 3;
    done;
    for i=1 to (a-1) do
      for j=1 to (b-1) do
	let k = (get_max_value mat1.(i-1).(j-1) str1.[i] str2.[j] mat1.(i-1).(j) mat1.(i).(j-1)) in
	  mat1.(i).(j) <- (fst k);
	  mat2.(i).(j) <- (snd k);
      done;
    done;
    (mat1,mat2);;


let get_max_score1 str1 str2 =
  let a = (String.length str1) in
  let b = (String.length str2) in
  let c = global_alignment str1 str2 in
    (fst c).(a-1).(b-1);;

let print_global_alignment str1 str2 =
  let a = (String.length str1) in
  let b = (String.length str2) in
  let c = global_alignment str1 str2 in
  let d = (fst c) in
  let e = (snd c) in
  let ret1 = ref [] in
  let ret2 = ref [] in
  let i = ref (a-1) in
  let j = ref (b-1) in
    while ( not (!i=0 && !j=0) ) do
      if e.(!i).(!j)=1 then
	begin
	  ret1 := str1.[!i]::!ret1;
	  ret2 := str2.[!j]::!ret2;
	  i := !i - 1;
	  j := !j - 1;
	end
      else if e.(!i).(!j)=2 then
	begin
	  ret1 := str1.[!i]::!ret1;
	  ret2 := '-'::!ret2;
	  i := !i -1;
	end
      else if e.(!i).(!j)=3 then
	begin
	  ret1 := '-'::!ret1;
	  ret2 := str2.[!j]::!ret2;
	  j := !j -1;
	end;
    done;
    print_endline "string1:";
    print_endline str1;
    print_endline "string2:";
    print_endline str2;
    print_newline ();
    print_endline "global alignment print:";
    ignore(List.map print_char !ret1);
    print_newline ();
    ignore(List.map print_char !ret2);
    print_newline ();;






(* local alignment : Smith-Watermn *)

let local_alignment str1 str2 =
  let a = (String.length str1) in
  let b = (String.length str2) in
  let mat1 = Array.make_matrix a b 0 in
  let mat2 = Array.make_matrix a b 0 in
    (* arg1: F(i-1,j-1) val1: x_i val2:y_i arg2:F(i-1,j) arg3:F(i,j-1) *)
  let get_max_value arg1 val1 val2 arg2 arg3 =
    let tmp1 = ref 0 in
    let v2 = arg2 + gap_init in
    let v3 = arg3 + gap_init in
      if val1=val2 then
	tmp1 := arg1 + score
      else 
	tmp1 := arg1 + cost;
      let v1 = !tmp1 in
      let ret = ref (0,0) in
	if (v1 >= v2) && (v1>=v3) && (v1>=0)then
	  ret := (v1,1)
	else if (v2 >=v1) && (v2>=v3) && (v2>=0) then
	  ret := (v2,2) 
	else if (v3>=v1) && (v3>=v2) && (v3>=0) then
	  ret := (v3,3)
	else 
	  ret := (0,4);
	!ret in
    for i=1 to a do
      mat1.(i-1).(0) <- 0
    done;
    for i=1 to b do
      mat1.(0).(i-1) <- 0
    done;
    for i=1 to a-1 do 
      mat2.(i).(0) <- 4
    done;
    for i=1 to b-1 do
      mat2.(0).(i) <- 4
    done;
    for i=1 to (a-1) do
      for j=1 to (b-1) do
	let k = (get_max_value mat1.(i-1).(j-1) str1.[i] str2.[j] mat1.(i-1).(j) mat1.(i).(j-1)) in
	  mat1.(i).(j) <- (fst k);
	  mat2.(i).(j) <- (snd k);
      done;
    done;
    (mat1,mat2);;


let get_max_score2 str1 str2 =
  let a = (String.length str1) in
  let b = (String.length str2) in
  let c = (local_alignment str1 str2) in
  let d = ref 0 in
  let e = (fst c) in
    for i=1 to a do
      for j=1 to b do
	if e.(i-1).(j-1) > !d then d := e.(i-1).(j-1);
      done;
    done;
    !d;;


let print_local_alignment str1 str2 =
  let find_max_index mat1 =
    let i = ref 0 in
    let j = ref 0 in 
    let max_score = ref min_int in
      for x=1 to Array.length mat1 do
	for y=1 to Array.length (mat1.(0)) do
	  if !max_score<mat1.(x-1).(y-1) then
	    begin
	      max_score := mat1.(x-1).(y-1);
	      i := x - 1;
	      j := y - 1;
	    end;
	done;
      done;
      (!i,!j) in
  let a = (String.length str1) in
  let b = (String.length str2) in
  let c = local_alignment str1 str2 in
  let d = (fst c) in
  let e = (snd c) in
  let ret1 = ref [] in
  let ret2 = ref [] in
  let i = ref (fst (find_max_index d)) in
  let j = ref (snd (find_max_index d)) in
    while ( not (!i=0 && !j=0) ) do
      if e.(!i).(!j)=1 then
	begin
	  ret1 := str1.[!i]::!ret1;
	  ret2 := str2.[!j]::!ret2;
	  i := !i - 1;
	  j := !j - 1;
	end
      else if e.(!i).(!j)=2 then
	begin
	  ret1 := str1.[!i]::!ret1;
	  ret2 := '-'::!ret2;
	  i := !i - 1;
	end
      else if e.(!i).(!j)=3 then
	begin
	  ret1 := '-'::!ret1;
	  ret2 := str2.[!j]::!ret2;
	  j := !j - 1;
	end
      else if e.(!i).(!j)=4 then
	begin
	  i := 0;
	  j := 0;
	end;
    done;
    print_endline "string1:";
    print_endline str1;
    print_endline "string2:";
    print_endline str2;
    print_newline ();
    print_endline "local alignment print:";
    ignore(List.map print_char !ret1);
    print_newline ();
    ignore(List.map print_char !ret2);
    print_newline ();;



(* phylogenetic tree algorithm : UPGMA *)

let get_distance str1 str2 = 
  let a = (String.length str1) in
  let b = (String.length str2) in
  let mat1 = Array.make_matrix a b 0.0 in
    (* arg1: F(i-1,j-1) val1: x_i val2:y_i arg2:F(i-1,j) arg3:F(i,j-1) *)
    (* gap : +1 , disagreement : +1 , agreement : 0 *)
  let get_max_value arg1 val1 val2 arg2 arg3 =
    let tmp1 = ref 0. in
    let v2 = arg2 +. 1. in
    let v3 = arg3 +. 1. in
      if val1=val2 then
	tmp1 := arg1
      else 
	tmp1 := arg1 +. 1.;
      let v1 = !tmp1 in
      let ret = ref 0.0 in
	if (v1 <= v2) && (v1<=v3) then
	  ret := v1
	else if (v2<=v1) && (v2<=v3) then
	  ret := v2 
	else if (v3<=v1) && (v3<=v2) then
	  ret := v3;
	!ret in 
    for i=1 to a do
      mat1.(i-1).(0) <- (float_of_int (i-1))
    done;
    for i=1 to b do
      mat1.(0).(i-1) <- (float_of_int (i-1))
    done;
    for i=1 to (a-1) do
      for j=1 to (b-1) do
	let k = (get_max_value mat1.(i-1).(j-1) str1.[i] str2.[j] mat1.(i-1).(j) mat1.(i).(j-1)) in
	  mat1.(i).(j) <- k;
      done;
    done;
    mat1;;



let get_min_distance str1 str2 =
  let a = (String.length str1) in
  let b = (String.length str2) in
  let c = (get_distance str1 str2) in
    c.(a-1).(b-1);;


let test num1 num2 =
  let a = gen_random_string num1 in
  let b = change a num2 in
    ignore (global_alignment a b);;


let find_group group i = 
  let a = ref i in
    while (group.(!a) != !a) do
      a := group.(!a);
    done;
    !a;;


let find_min mat1 group index =
  let ret1 = ref infinity in
  let ret2 = ref 0 in
  let ret3 = ref 0 in 
    for i=0 to index do
      for j=0 to index do
	if (i<=j && (find_group group i)!=(find_group group j)) then
	  if (find_group group i)=i && (find_group group j)=j && mat1.(i).(j) < !ret1 then
	    begin
	      ret1 := mat1.(i).(j);
	      ret2 := i;
	      ret3 := j;
	    end;
      done;
    done;
    (!ret2,!ret3);;


let set_distance mat1 group min_pair1 min_pair2 new_node_num = 
  let get_group_list_num group part =
    let a = ref 0 in
    let b = find_group group part in
      for i=1 to  Array.length group do
	if (find_group group (i-1)) = b then
	  a := !a + 1;
      done;
      (float_of_int !a) in
    for i=0 to new_node_num do
      if ( (find_group group group.(new_node_num)) != (find_group group i)) then
	begin
	  let a = get_group_list_num group min_pair1 in
	  let b = get_group_list_num group min_pair2 in
	  let c = (mat1.(min_pair1).(i) *. a +. mat1.(min_pair2).(i) *. b) /. (a +. b) in
	    mat1.(new_node_num).(i) <- c;
	    mat1.(i).(new_node_num) <- c;
	end
    done;;


(* building phylogenetic tree *)
let phylo_tree str_array =
  let len = Array.length str_array in
  let mat1 = Array.make_matrix (len*2-1) (len*2-1) 0.0 in
  let group = Array.make (len*2-1) 0 in
    for i=1 to (Array.length str_array) do
      for j=1 to (Array.length str_array) do
	if i<=j then
	  if i=j then
	    mat1.(i-1).(j-1) <- 0.0
	  else
	    begin
	      mat1.(i-1).(j-1) <- (get_min_distance str_array.(i-1) str_array.(j-1));
	      mat1.(j-1).(i-1) <- mat1.(i-1).(j-1);
	    end;
      done;
    done;
    for i=1 to (Array.length group) do
      group.(i-1) <- i-1
    done;
    for i=1 to len-1 do
      let a = (find_min mat1 group (i+len-2)) in
      let b = (fst a) in
      let c = (snd a) in
	(* b,c : pair of min distance, i+len-1 : new_node_number *)
	set_distance mat1 group b c (i+len-1);
	group.(b) <- (i+len-1);
	group.(c) <- (i+len-1);
	group.(i+len-1) <- (i+len-1);
    done;
    (mat1,group);;



let sample_data0 = gen_random_string 100;;
let sample_data1 = change sample_data0 1;;
let sample_data2 = change sample_data0 20;;
let sample_data3 = change sample_data1 5;;
let sample_data4 = change sample_data3 5;;

let sample_array = [| sample_data0;sample_data1; sample_data2; sample_data3; sample_data4  |];;