#!/usr/bin/perl -w

use strict;
use Math::Libm ':all';

#############################################################################
############################# VARIABLES #####################################
#############################################################################

my $fitxerfasta;                         #nom de l'arxiu que conte les sequencies en format fasta
my $lmotiu;                              #la longitud del motiu que lusuari introdueix
my @input;                               #per treure 2 vectors referenciats de la funcio llegirseq
my @ident;                               #emmagatzema les identificacions de les sequencies
my @seq;                                 #emmagatzema els nucleotids de cada sequencia
my %mat_pesos;                           #freqs absoluts de cada nucleotid en el motiu-bckg
my %mat_relatives;                       #freqs relatives de cada nucleotid en el motiu-bckg
my $counts=0;                            #comptar les voltes que fa la iteracio
my %nucleotids;                          #freqs absolutes x nucleotid, sns importar la posicio
my @pos_candidats;                       # millors posicions d'inici del motiu per cada seq
my @pesos;                               #pesos tots candidats de cada sequencia
my @out;                                 #per treure una hash referenciat i un valor de la funció M_step
my $loglkh= -99999999;                   #la versemblança del model que estem calculant
my $current_loglkh= -999999999999;       #per comparar la versemb del model anterior-actual
my %mat_provisional;                     #per provar 10 matrius random avans de comensar
my $dif=1;


#############################################################################
########################### COS DEL PROGRAMA ################################
#############################################################################

if (@ARGV < 2) {

    print "EM.pl fitxerfasta longitudmotiu\n";  #informa que falten arguments si no hi hem posat els 2
    exit;
}

$fitxerfasta = $ARGV[0];                      #argument amb el fitxer fasta per processar
$lmotiu = $ARGV[1];                           #argument amb la longitud del motiu q volem buscar
                                              #veure funcions per seguir el programa
@input = llegirseq ($fitxerfasta);

@ident = @{$input[0]};
@seq = @{$input[1]};

%nucleotids = contarnucleotids (\@seq);      

while ($counts<10) {                                      #prova 10 vegades de:
   
    %mat_provisional = randommatrix (\@seq);              #fer una matriu random
    %mat_relatives = frec_relatives(\%mat_provisional);   #
    @pesos = E_step (\@seq, \%mat_relatives);             #
    @out = M_step (\@seq,\@pesos,\%nucleotids);           # Evaluar la versemblansa i :
    $loglkh = $out[1];
   
    if ($loglkh > $current_loglkh) {                      # si es millor que l'anterior

	$current_loglkh = $loglkh;                        # guardar el loglkh  
	%mat_pesos=%mat_provisional;                      # guardar la matriu que es millor que l'anterior

    }
    $counts = $counts +1;
}

%mat_relatives = frec_relatives (\%mat_pesos);           #fa les frequencies relatives de la millor 
                                                         # matriu random que ha trovat

$counts =0;

$loglkh = -99999999;                                    # reinicia els valors de loglkh
$current_loglkh = -99999999999999;                      #



while (($dif)>(1/10000) && $counts < 1000) { #  mentre el loglkh millori mes de 0,0001
                                                                #  i no hagi fet mes de mil iteracions

    $current_loglkh = $loglkh;                                  
 
    @pesos = E_step (\@seq, \%mat_relatives);                    # fes E i M step, i avalua el model

    @out = M_step (\@seq, \@pesos,\%nucleotids);

    %mat_relatives=%{$out[0]};

    $loglkh= $out[1];

    $counts = $counts +1;

    $dif = diferencia ($loglkh, $current_loglkh);    

}


@pesos = E_step ( \@seq, \%mat_relatives);                     #la matriu de pesos bona la pases per totes les seq

@pos_candidats = posiciocandidats (\@pesos);                   #et quedes amb la posició de cada seq que dona mes probabilitat

print "\n\nLa versemblansa del model es: $loglkh\n\n\n";       # Comensa l'output

output (\@seq, \@ident, \@pos_candidats, \%mat_relatives);





##############################################################################
#################################### FUNCIONS ################################
##############################################################################



sub llegirseq {                               #  posar les ident de la sequencia en el vector @ident
                                              #  i les sequencies concatenades en el vector @seq
    my $seq1 = "";
    my @ident;
    my @seq;
    my $i=0;
    my @input;

    if (!open(SEQUENCIES,"<$_[0]")) {         #si no es pot obrir mostra error i tanca el programa

	print "EM.pl impossible obrir $fitxerfasta\n";
	exit;

    }


    while () {

	if (m/(>.+)/) {

	    $ident[$i]=$1;
	    $seq1="";
	    $i=$i+1;

	} else {

	    m/([actgACTG]+)/;
	    $seq1=$seq1.$1;

	}
	$seq[@ident-1]=uc($seq1);
    }

    $input[0] = \@ident;                      #referenciem els vectors per posarlos al vector input
    $input[1] = \@seq;
    return @input;                            #vector amb les ident de les seqs i les seqs concatenades
}
                                              #  generacio d'un primer model a l'atzar,
sub randommatrix {                            #  generant a l'atzar els llocs d'inici
                                              #  del motiu per a cada sequencia
    my @seq =@{$_[0]};
    my $i=0;
    my %mat_pesos = iniciarmatriu ($lmotiu);  #  hash per anotar les freqs absolutes dels nucleotids
                                              #  en el background i en les pos del motiu

    while ($i<@seq) {

	my $rand;
	my $j=0;
	my @linia = ( $seq[$i] =~ m/[ACTG]/g );
	$rand = int(rand(@linia-$lmotiu));

	while ($j<@linia) {

	    if ($j == $rand) {                #  si estem dins el motiu, suma 1 al
                                              #  nucleotid que correspongui en cada posicio
		my $k = 0;

		while ( $k < $lmotiu ) {

		    $mat_pesos{$linia[$j+$k]}[$k+1] = $mat_pesos{$linia[$j+$k]}[$k+1] + 1;
		    $k = $k + 1;

		}

		$j = $j + $lmotiu;

	    }                                 #   quan no estem dins el motiu, suma 1 al nucleotid que correspongui
                                              #   en la posicio 0 del hash (background)
	    $mat_pesos{$linia[$j]}[0] = $mat_pesos{$linia[$j]}[0] + 1;
	    $j = $j +1;
	}
	$i = $i + 1;
     }

    return %mat_pesos;                        #  hash amb les freqs absolutes de cada
}                                             #  nucleotid en el background
                                              #  i suposant el lloc d'inici del motiu

sub contarnucleotids {                        #  calculem els nucleotids de cada que hi ha
                                              #  en totes les sequencies
    my @seq = @{$_[0]};
    my $i =0;
    my %nucleotids = ( "A" => "0" ,           #comencem inicialitzant tots els valors a 0
		       "C" => "0" ,
		       "T" => "0" ,
		       "G" => "0" );

    while ($i<@seq) {                         #  per cada nucleotid que trobem de cada,
                                              #  sumarem 1 al valor total d'aquell nucleotid
	my $j=0;
	my @linia = ($seq[$i] =~ m/[ACTG]/g);

	while ($j < @linia) {

	    $nucleotids{$linia[$j]}= $nucleotids{$linia[$j]} + 1;
	    $j =$j+1;
	}
	$i = $i +1;
    }
    return %nucleotids;                       #hash amb els valors absoluts de cada nucleotid
}


sub iniciarmatriu {                           #  inicialitzem la matriu amb el background i les
                                              #  posicions del motiu amb 0 a totes les posicions
    my $lmotiu = $_[0];
    my $i=0;
    my %mat_pesos;

    while ($i<$lmotiu+1) {

	$mat_pesos{A}[$i]=0;
	$mat_pesos{C}[$i]=0;
	$mat_pesos{T}[$i]=0;
	$mat_pesos{G}[$i]=0;
	$i = $i +1;

    }
    return %mat_pesos;
}


                                              #  calcul de les freqs relatives dels nucleotids
sub frec_relatives {                          #  en el background i les posicions del motiu

    my %mat_pesos = %{$_[0]};                 #hash amb les frequencies absolutes
    my %mat_relatives;                        #hash que anem a calcular amb les freqs relatives
    my $i=0;

    while ($i<$lmotiu+1) {                    #  recorrem totes les posicions del %mat_pesos i
                                              #  dividim cada valor perl a suma de valors de la columna

	my $suma = $mat_pesos{A}[$i]+$mat_pesos{C}[$i]+$mat_pesos{T}[$i]+$mat_pesos{G}[$i]+1;
	$mat_relatives{A}[$i] = ($mat_pesos{A}[$i]+(1/4))/$suma;
	$mat_relatives{C}[$i] = ($mat_pesos{C}[$i]+(1/4))/$suma;
	$mat_relatives{T}[$i] = ($mat_pesos{T}[$i]+(1/4))/$suma;
	$mat_relatives{G}[$i] = ($mat_pesos{G}[$i]+(1/4))/$suma;
	$i = $i +1;
   }
    return %mat_relatives;                    #hash de les frequencies relatives
}
                                              #  pas en que puntuem cada candidat en cada
sub E_step {                                  #  sequencia segons el model que estem usant
                                              #  el model el representa el %mat_relatives
    my @seq = @{$_[0]};
    my %mat_relatives = %{$_[1]};
    my @pesos;
    my $i=0;


    while ($i < @seq) {                       # per cada sequencia:

	my @linia = ($seq[$i] =~ m/[ACTG]/g); #  cada sequencia la fragmentem en un @linia
	my $inici = 0;                        #  que conte un nucleotid en cada posicio
	my $suma = 0;
        my @sum;

	while ($inici <= @linia-$lmotiu) {   #per cada possible lloc d'inici del motiu:
                                             
	    $pesos[$i][$inici] = 0;

	    my $j=0;

	    while ($j < @linia) {             #per cada nucleotid:


		if ($j == $inici) {           #  quan entrem dins el motiu, puntuem el candidat
                                              #  segons la posicio q ocupen els nucleotids amb els valors de %mat_relatives
		    my $k=0;

		    while ($k < $lmotiu) {

			$pesos[$i][$inici] = $pesos[$i][$inici]+log($mat_relatives{$linia[$j+$k]}[$k+1]);
			$k = $k+1;
		    }

		    $j =$j + $lmotiu;
		}

		if ($j != @linia) {           #  fora del motiu puntuem el candidat segons
                                              #  la freq del background, sns importar la posicio
		    
		    
		    $pesos[$i][$inici] = $pesos[$i][$inici]+log($mat_relatives{$linia[$j]}[0]);
		    $j = $j+1;
		}
	    }
	    
	    $sum[$inici]=$pesos[$i][$inici];
	    
	    $inici = $inici +1;
	}

	$suma=logsum(@sum);
	$inici=0;                             #  reinicialitzem altre cop per normalitzar
                                              #  els pesos de cada candidat en cada sequencia
	while ($inici<@{$pesos[$i]}) {

	    $pesos[$i][$inici] = exp($pesos[$i][$inici] - $suma);
	    $inici = $inici + 1;
	}

	$i=$i+1;
    }

    return @pesos;                            #  retorna una matriu amb les probabilitats 
                                              #  per cada posicio que, donada una  sequencia, 
}                                             #  el motiu comenci en aquella posicio


sub M_step {                                  #  updating de la matriu bckg-motiu per millorar
                                              #  el model segons les puntuacions dels candidats
    my @seq= @{$_[0]};
    my @pesos = @{$_[1]};
    my %nucleotids = %{$_[2]};
    my $i=0;
    my %mat_relatives;
    my %mat_pesos = iniciarmatriu ($lmotiu);  #reinicialitzem el hash amb 0 a totes posicions
    my %suma = ( "A" => "0",
		 "C" => "0",                  #
		 "T" => "0",                  #
		 "G" => "0" );                #  inicialitzem els sumatoris
                                              #  posant 0 a tots els valors
    my %sumalog = ( "A" => "0",               #
		    "C" => "0",               #
		    "T" => "0",               #
		    "G" => "0" );
    my $versem;                               #expressa la versemblansa del model


    while ($i<@seq) {                         #per cada sequencia:

	my $j=0;

	while ($j<@{$pesos[$i]}){           #per cada lloc d'inici possible:

	    my @motiu = (substr($seq[$i],$j,$lmotiu) =~ m/[ACTG]/g); #agafem el motiu en un vector
	    my $k=0;


	    while ($k<@motiu){               #recorrem el motiu i  suma el pes d'aquell candidat al nucleotid i la posicó que troba

		$mat_pesos{$motiu[$k]}[$k+1] = $mat_pesos{$motiu[$k]}[$k+1] + $pesos[$i][$j]; 
		$k = $k +1;

	    }
	    $j = $j +1;
	}

	$i=$i+1;
    }

    $i = 1;

    while ($i<$lmotiu+1){

	$suma{A} = $suma{A} + $mat_pesos{A}[$i];     #
	$suma{C} = $suma{C} + $mat_pesos{C}[$i];     #  nucleotids de cada dins del motiu
	$suma{T} = $suma{T} + $mat_pesos{T}[$i];     #  
	$suma{G} = $suma{G} + $mat_pesos{G}[$i];     #
	$i = $i +1;
    }

    $mat_pesos{A}[0] = $nucleotids{A}-$suma{A};      #
    $mat_pesos{C}[0] = $nucleotids{C}-$suma{C};      #  nucleotids de cada en el bckg, igual
    $mat_pesos{T}[0] = $nucleotids{T}-$suma{T};      #  al total de nucleotids 
    $mat_pesos{G}[0] = $nucleotids{G}-$suma{G};      #  menys els del motiu

    %mat_relatives = frec_relatives (\%mat_pesos);   #  recalculem les freqs relatives tenint
                                                     #  en compte les freq absolutes actualitzades
    $i=0;

    while ($i < $lmotiu+1) {                         #tornem a recorrer la matriu per calcular la versemblansa

	$sumalog{A} = $sumalog{A} + ($mat_pesos{A}[$i] * log($mat_relatives{A}[$i]));
	$sumalog{C} = $sumalog{C} + ($mat_pesos{C}[$i] * log($mat_relatives{C}[$i]));
	$sumalog{T} = $sumalog{T} + ($mat_pesos{T}[$i] * log($mat_relatives{T}[$i]));
	$sumalog{G} = $sumalog{G} + ($mat_pesos{G}[$i] * log($mat_relatives{G}[$i]));
	$i = $i +1;

    }

    $versem = $sumalog{A}+$sumalog{C}+$sumalog{T}+$sumalog{G};

    my @out = ( \%mat_relatives, $versem );         #retorna la matriu ponderada i el loglkh de la matriu

    return @out;

}


sub logsum {
my $len = scalar(@_);
my $x = 0;

if ($len > 1) {

$x = $_[0];
for (my $i = 1; $i < $len; $i++) {
$x = $x + log1p(exp($_[$i]-$x));
}

}

return $x;
}


sub diferencia {

    my $loglkh = $_[0];
    my $current_loglkh = $_[1];
    my $dif;

    if ($loglkh=~ m/\-(\d+.\d+)/) {
	$loglkh = $1;
    }

    if ($current_loglkh=~ m/\-(\d+.\d+)/) {
	$current_loglkh = $1;
    }

    if ($loglkh<$current_loglkh) {

	$dif = ($current_loglkh-$loglkh);

    } else {

	$dif = ($loglkh-$current_loglkh);
    }
    return $dif;
}



sub posiciocandidats {                        #  un cop acceptat el millor model, selecciona el millor candidat
                                              #  de cada sequencia, guardant la posicio que ocupa i la probabilitat que te
    my @pesos = @{$_[0]};
    my $i=0;
    my @pos_candidats;

    while ($i<@pesos) {                       #per cada seq:

	my $j=0;
	my $max = -999999999;

	while ($j<@{$pesos[$i]}) {           #per cada lloc d'inici:

	    if ($pesos[$i][$j] > $max ) {     #si el valor que mirem es millor que el que tenim ens el quedem, i si no el descartem

		$max = $pesos[$i][$j];
		$pos_candidats[$i][0] = $j;
		$pos_candidats[$i][1] = $max;

	    }
	    $j= $j+1;
	}
	$i = $i +1;
    }
    return @pos_candidats;                    #vector que conte el millor candidat de cada sequencia
}


sub output {                                  #disseny de l'output, que mostrara el programa per la pantalla

    my @seq = @{$_[0]};
    my @ident = @{$_[1]};
    my @pos_candidats = @{$_[2]};
    my %mat_relatives = %{$_[3]};
    my $i=0;
    my %mat_loglkh;
    my %mat_info;
    my @bits_info;
 
    print "Identificador \t Motiu \t   Posicio  \t  Pes\n";

    while ($i< @ident) {

	my $motiu = substr($seq[$i],$pos_candidats[$i][0],$lmotiu);
	print "$ident[$i]  \t$motiu  \t",($pos_candidats[$i][0]+1),"  \t";
	printf "%.4f\n",$pos_candidats[$i][1]; 
	$i = $i +1;
    }
    $i = 1;

    while ($i < $lmotiu +1) {

	$mat_loglkh{A}[$i] =log ($mat_relatives{A}[$i]) - log ($mat_relatives{A}[0]);
	$mat_loglkh{C}[$i] =log ($mat_relatives{C}[$i]) - log ($mat_relatives{C}[0]);
	$mat_loglkh{T}[$i] =log ($mat_relatives{T}[$i]) - log ($mat_relatives{T}[0]);
	$mat_loglkh{G}[$i] =log ($mat_relatives{G}[$i]) - log ($mat_relatives{G}[0]);

	$mat_info{A}[$i] = ($mat_relatives{A}[$i]*log($mat_relatives{A}[$i])/log(2));
	$mat_info{C}[$i] = ($mat_relatives{C}[$i]*log($mat_relatives{C}[$i])/log(2));
	$mat_info{T}[$i] = ($mat_relatives{T}[$i]*log($mat_relatives{T}[$i])/log(2));
	$mat_info{G}[$i] = ($mat_relatives{G}[$i]*log($mat_relatives{G}[$i])/log(2));
	
	$bits_info[$i]= ($mat_info{A}[$i]+$mat_info{C}[$i]+$mat_info{T}[$i]+$mat_info{G}[$i])+2;

	$i =$i+1;
    }
    print "\n\n";

    print "Matriu de frequencies relatives del motiu\n\n";
    printarhash (\%mat_relatives);

    $i=1;    
    
    print "Bits Info\nB\t";
    
    while ($i<@bits_info){

	printf "%.2f\t",$bits_info[$i];
	$i=$i+1;
    }
    print "\n\n";
    print "Matriu de pesos del motiu\n\n";
    printarhash (\%mat_loglkh);

   
}

sub printarhash {

    my %hash=%{$_[0]};
    my $i =1;
    my @nucleotids = ("A", "C", "T", "G");
    print "Posicio\t";

    while ($i<$lmotiu+1) {

	print "pos$i\t";
	$i = $i+1;
    }

    $i=0;
    print "\n";

    while ($i<4) {
	my $j=1;
	print $nucleotids[$i],"\t";
	while ($j<$lmotiu+1){
	    printf "%.2f\t",$hash{$nucleotids[$i]}[$j];
	    $j=$j+1;
	}
	print "\n";
	$i =$i+1;
    }
    print "\n\n";
}





##############################################################################################
##############################################################################################
########################                                           ###########################
########################  implemented by: EVA MERCE & SERGI REGOT  ###########################
########################                                           ###########################
##############################################################################################
##############################################################################################