#!/usr/local/bin/perl -w

use lib '/ad/eng/support/software/linux/all/all/lib/perl'; # path to BioPerl on cluster

use strict;
use Bio::SeqIO;  # BioPerl module to handle fasta files




# construct a usage statement to print
my $usage = <<USAGE;

    USAGE: perl 4mer_motifs.pl <fasta_file> <number_of_randomizations>

    options:
    1. a fasta file with any number of sequences
    2. number of randomizations to generate null distribution


    NOTE
    calculate 4mer over/under representation based on observed 4mer frequencies
    relative to a null distribution.

    will print out 5 columns as follows:
    4mer  observed_count  random_mean  random_sd  z-score

  
USAGE


my $fasta_file               = shift @ARGV;
my $number_of_randomizations = shift @ARGV;

# check command line arguments and display help if needed
unless ($fasta_file && $number_of_randomizations) { 
    die "$usage"; 
}

# initialize hash to keep track of the observed 4mer counts
# the hash will have the following organization
# key = 4mer
# val = count
# this subroutine is number 1
my $four_mer_hash_ref = &initialize_4mer_hash_counts_to_zero();

# get hash from memory location
my %four_mer_hash     = %{$four_mer_hash_ref};


# check to make sure this subroutine was coded properly, and if it was not
# exit the program and print an error message
&check_subroutine1(\%four_mer_hash);


# initialize a hash of hashes to hold the random counts
# the hash will have the following organization
# key1 = 4mer
# key2 = randomization count
# val  = count for that 4mer for that randomization
# this subroutine is number 2
my $rand_hash = &initialize_random_4mer_hash_of_hashes($number_of_randomizations);

# get hash from memory location
my %random_hash = %{$rand_hash};


# check to make sure this subroutine was coded properly, and if it was not
# exit the program and print an error message
&check_subroutine2(\%random_hash);


# open BioPerl SeqIO object
my $seq_object = &open_bioperl_seqio_object($fasta_file);



# go through each sequence in the fasta file
while (my $seq_obj = $seq_object->next_seq) {

    # get the sequence header from the object
    # we don't need this, we're just showing you how to do it FYI
    my $name = $seq_obj->id;


    # get the nucleotide sequence from the object
    my $sequence = $seq_obj->seq;


    # scan through the current sequence in the fasta file 4 nucleotides at a time
    # update %four_mer_hash with the observed counts
    for (my $i = 0; $i <= length($sequence) - 4; $i++) {
	my $four_mer = uc(substr $sequence, $i, 4);
	$four_mer_hash{$four_mer}++;
    }


    # now do the specified number of randomizations for the current sequence
    for (my $n = 0; $n < $number_of_randomizations; $n++) {

	# generate a random sequence that is the same length as the current sequence
	# this subroutine is number 3
	my $rand_seq = &create_random_sequence_given_length(length($sequence));

	# check to make sure subroutine 3 is working
	# you can comment out this checkpoint once the subroutine works
	&check_subroutine3(length($sequence), $rand_seq);
	

	# scan through the random sequence and tally the 4mer counts for each iteration
	for (my $j = 0; $j <= length($rand_seq) - 4; $j++) {
	    my $rand_four_mer = uc(substr $rand_seq, $j, 4);
	    $random_hash{$rand_four_mer}{$n}++;
	}
    }
}

# now go through each 4mer and calculate the mean and sd from the randomizations
# then calculate a z-score for the observed frequency
# declare a hash to hold each 4mer as the key and the mean, sd, or z-score as a value
my %means;
my %sds;
my %z_scores;

for my $key1 (keys %four_mer_hash) {
    
    # declare an array to hold the randomized data so we can calculate a mean and sd
    my @data;

    # for the current 4mer, put all the randomization data into an array
    for my $key2 (keys %{$random_hash{$key1}}) {
	push @data, $random_hash{$key1}{$key2};
    }

    # now calculate the mean and sd for the given array of data
    # this subroutine is number 4
    my ($mean, $sd) = &calculate_mean_and_sd_from_array_ref(\@data);

    # check to make sure subroutine 4 is working
    # you can comment out this checkpoint once the subroutine works
    &check_subroutine4($mean, $sd);

    # now calculate a z-score based on the observed counts, and the mean and sd from the randomizations
    my $z = sprintf("%0.5f", ($four_mer_hash{$key1} - $mean) / $sd);

    # format all the numbers for pretty printing
    $means{$key1}    = sprintf("%0.1f", $mean);
    $sds{$key1}      = sprintf("%0.2f", $sd);
    $z_scores{$key1} = sprintf("%0.2f", $z);
}

# now sort by the 4mers and print out
for my $key (sort keys %z_scores) {
    my $observed = $four_mer_hash{$key};
    print "$key\t$observed\t$means{$key}\t$sds{$key}\t$z_scores{$key}\n";
}


#####################
#### SUBROUTINES ####
#####################



# SUBROUTINE 1 #
sub initialize_4mer_hash_counts_to_zero {
    
    # given: nothing
    # returns: 1 scalar (a reference to a hash)



    # initialize a hash to hold the observed counts
    # the hash will have the following organization
    # key = 4mer
    # val = count, initialized to zero, for that 4mer
    my %hash;


    # declare an array to hold all possible nucleotides
    my @bases = qw(A C G T);

    ####---> INSERT CODE HERE <---####
    foreach my $line (@seqs){
              chomp $line; 
	    
      

	
  
    # HINTS
    # loop through all possible nucleotides multiple times to generate all possible 4mers


		    # inside your loops, concatenate all single nucleotides together to create a 4mer
                    # uncomment the line below when you are ready to use it
		    #my $seq = "$nt1" . "$nt2" . "$nt3" . "$nt4";

		    # initialize the count for this 4mer to zero
                    # uncomment the line below when you are ready to use it
		    #$hash{$seq} = 0;



    # return a reference to the hash
    return \%hash;
}



# SUBROUTINE 2 #
sub initialize_random_4mer_hash_of_hashes {
    # given: 1 scalar (number of randomizations to do)
    # returns: 1 scalar (a reference to a hash of hashes)


    # get the number of randomizations passed to this subroutine
    my $num_of_rands = shift;


    # initialize a hash of hashes to hold the random counts
    # the hash will have the following organization
    # key1 = 4mer
    # key2 = randomization count
    # val  = count, initialized to zero, for that 4mer for that randomization
    my %hash;


    # declare an array to hold all possible nucleotides
    my @bases = qw(A C G T);

    ####---> INSERT CODE HERE <---####
  
    # HINTS
    # loop through all possible nucleotides multiple times to generate all possible 4mers
    

		    # inside your loops, concatenate all single nucleotides together to create a 4mer
                    # uncomment the line below when you are ready to use it
		    #my $seq = "$nt1" . "$nt2" . "$nt3" . "$nt4";

		    # loop through all possible randomization counts

			
			# initialize the count for this 4mer and randomization number to zero
                        # uncomment the line below when you are ready to use it
			#$hash{$seq}{$n} = 0;



    # return a reference to the hash of hashes
    return \%hash;
}



# SUBROUTINE 3 #
sub create_random_sequence_given_length {
    # given: 1 scalar (the length of a random DNA sequence to generate)
    # returns: 1 scalar (a randomly generated DNA sequence)

    # get the length of the random sequence to produce
    my $seq_length = shift;


    # declare an array to hold each base of the randomly generated sequence
    my @seq;


    ####---> INSERT CODE HERE <---####
    # ASSUME ALL NUCLEOTIDES ARE EQUALLY PROBABLE! #





    # assuming you created a random sequence in an array, join the array elements into
    # a scalar variable and return that variable.
    my $sequence = join '', @seq;

    return $sequence;
}



# SUBROUTINE 4 #
sub calculate_mean_and_sd_from_array_ref {
    # given: 1 scalar (a reference to an array of numbers)
    # returns: 2 scalars (the mean and sd of all the values in the array)

    # get the array reference that was passed to this subroutine
    my $array_ref = shift;


    # declare meand and sd variables to report back
    my $mean = 0;
    my $sd   = 0;

    ####---> INSERT CODE HERE <---####



    # return the mean and sd that was calculated
    return $mean, $sd;
}














sub open_bioperl_seqio_object {
    
    # NO CODE NEEDED FOR THIS SUBROUTINE #

    my $file = shift;

    my $seqio_obj;

    if ($file =~ m/\.gz$/) {
	$seqio_obj = Bio::SeqIO->new(-file   => "gunzip -c $file |",
				     -format => "fasta" );
    }
    else {
	$seqio_obj = Bio::SeqIO->new(-file   => "$file",
				     -format => "fasta" );
    }
    return $seqio_obj;
}



###############################################
#### ADDITIONAL ERROR-CHECKING SUBROUTINES ####
####                                       ####
####          IGNORE BELOW HERE            ####
###############################################


sub check_subroutine1 {
    my $ref = shift;
    my %hash = %{$ref};
    unless (keys %hash > 0) {
	die "\n\nsubroutine 1: \"initialize_4mer_hash_counts_to_zero\" was not coded properly\n\nquitting...\n\n";
    }
    for my $key (keys %hash) {
	unless ($hash{$key} == 0) {
	    die "\n\nsubroutine 1: \"initialize_4mer_hash_counts_to_zero\" was not coded properly\n\nquitting...\n\n";
	}
    }
}

sub check_subroutine2 {
    my $ref = shift;
    my %hash = %{$ref};
    unless (keys %hash > 0) {
	die "\n\nsubroutine 2: \"initialize_random_4mer_hash_of_hashes\" was not coded properly\n\nquitting...\n\n";
    }
    for my $key1 (keys %hash) {
	for my $key2 (keys %{$hash{$key1}}) {
	    unless ($hash{$key1}{$key2} == 0) {
		die "\n\nsubroutine 2: \"initialize_random_4mer_hash_of_hashes\" was not coded properly\n\nquitting...\n\n";
	    }
	}
    }
}

sub check_subroutine3 {
    my $length   = shift;
    my $rand_seq = shift;

    unless ($rand_seq =~ m/[ACGT]/ && length($rand_seq) == $length) {
	die "\n\nsubroutine 3: \"create_random_sequence_given_length\" was not coded properly\n\nquitting...\n\n";
    }
}


sub check_subroutine4 {
    my $mean = shift;
    my $sd   = shift;

    unless ($mean > 0 && $sd > 0) {
	die "\n\nsubroutine 4: \"calculate_mean_and_sd_from_array_ref\" was not coded properly\n\nquitting...\n\n";
    }
}