#!/usr/bin/perl -w
use strict;
use warnings;
use Getopt::Long;
use Cwd;


# Read the coverage output generated by the script genome_methylation2bedGraph.pl which must be run in --counts mode
# The format needs to be: <chromosome>  <start position>  <end position>  <methylation percentage>   <count methylated>  <count unmethylated>

my $parent_dir = getcwd
my %chromosomes; # storing the genome of interest


########################################################################################################################
### OPTIONS
########################################################################################################################


my $help;
my $genome_folder;
my $CpG_only;
my $CX_context;
my $dinucleotide_context;
my $zero;

GetOptions("h|help" => \$help,
	   "genome_folder=s" => \$genome_folder,
	   "CX" => \$CX_context,
	   "dinucleotide_context" => \$dinucleotide_context,
	   "zero_based" => \$zero,
	  );

### HELP
if ($help){
  usage();
}

### NO INPUT FILE
if(scalar @ARGV < 1){
  warn "Missing input file\n";
  die usage();
}

### GENOME FOLDER
if ($genome_folder){
  unless ($genome_folder =~/\/$/){
    $genome_folder =~ s/$/\//;
  }
  warn "Genome folder was specified as $genome_folder\n";
}
else{
  $genome_folder = '/data/public/Genomes/Mouse/NCBIM37/';
  warn "Using the default genome folder /data/public/Genomes/Mouse/NCBIM37/\n";
}

### Cyosine context
if ($CX_context){
  warn "Reporting methylation for all cytosine contexts. Be aware that this will generate enormous files\n";
  sleep(2);
}
else{ # default
  $CpG_only = 1;
  warn "Reporting cytosine methylation in CpG context only\n";
  sleep(2);
}

### Additional Dinucleotide context
if ($dinucleotide_context){
  warn "Reporting an additional column with dinucleotide cytosine context (e.g. CG or CA)\n";
  sleep(2);
}

### Zero-based coordinates
if ($zero){
  warn "Using zero-based genomic coordinates (user-defined)\n";
  sleep(2);
}
else{ # default, 1-based coords
  warn "Using 1-based genomic coordinates (default)\n\n";
  sleep(2);
}


##########################################################################################################################

read_genome_into_memory();

foreach my $in (@ARGV){

  open (IN,$in) or die $!;

  my $last_chr;
  my %chr; # storing reads for one chromosome at a time

  my $count = 0;
  while (<IN>){
    chomp;
    ++$count;
    my ($chr,$start,$end,undef,$meth,$nonmeth) = (split /\t/);

    ### die if the input file format seems wrong
    if ($count == 1){
      die "The input file needs to have been generated with the genome_methylation2bedGraph script while using the '--counts' option. The input format needs to be:\n<chromosome>  <start position>  <end position>  <methylation percentage>   <count methylated>  <count unmethylated>\n\n" unless (defined $meth and defined $nonmeth);
    }

    # defining the first chromosome
    unless (defined $last_chr){
      $last_chr = $chr;
      # warn "Storing all covered cytosine positions for chromosome: $chr\n";
    }

    if ($chr eq $last_chr){
      $chr{$chr}->{$start}->{meth} = $meth;
      $chr{$chr}->{$start}->{nonmeth} = $nonmeth;
    }
    else{
      warn "Writing out detailed cytosine reports for chromosome $last_chr\n";
      warn "Stored ",scalar keys $chr{$last_chr}," positions for chromosome $last_chr\n";
      
      while ( $chromosomes{$last_chr} =~ /([CG])/g){
	
	my $tri_nt = '';
	my $di_nt= '';
	my $pos = pos$chromosomes{$last_chr};

	my $strand;
	my $meth = 0;
	my $nonmeth = 0;
	
	if ($1 eq 'C'){    # C on forward strand
	  $tri_nt = substr ($chromosomes{$last_chr},($pos-1),3);   # positions are 0-based!
	  $strand = '+';
	}
	elsif ($1 eq 'G'){ # C on reverse strand
	  $tri_nt = substr ($chromosomes{$last_chr},($pos-3),3);   # positions are 0-based!
	  $tri_nt = reverse $tri_nt;
	  $tri_nt =~ tr/ACTG/TGAC/;
	  $strand = '-';
	}

	$di_nt = substr ($tri_nt,0,2);

	if (exists $chr{$last_chr}->{($pos-1)}){ # stored positions are 0-based!
	  $meth =  $chr{$last_chr}->{$pos-1}->{meth};
	  $nonmeth = $chr{$last_chr}->{$pos-1}->{nonmeth};
	}
	
	if ($CpG_only){ # default

	  if ($di_nt eq 'CG'){
	    if ($dinucleotide_context){
	      if ($zero){ # zero based coordinates
		$pos -= 1;
		print "$last_chr\t$pos\t$strand\t$meth\t$nonmeth\t$tri_nt\t$di_nt\n";
	      }
	      else{ # default
		print "$last_chr\t$pos\t$strand\t$meth\t$nonmeth\t$tri_nt\t$di_nt\n";
	      }
	    }
	    else{ #default
	      if ($zero){ # zero based coordinates
		$pos -= 1;
		print "$last_chr\t$pos\t$strand\t$meth\t$nonmeth\t$tri_nt\n";
	      }
	      else{ # default
		print "$last_chr\t$pos\t$strand\t$meth\t$nonmeth\t$tri_nt\n";
	      }
	    }
 	  }

	}
	else{ ## all cytosines
	  if ($dinucleotide_context){
	    if ($zero){ # zero based coordinates
	      $pos -= 1;
	      print "$last_chr\t$pos\t$strand\t$meth\t$nonmeth\t$tri_nt\t$di_nt\n";
	    }
	    else{ # default
	      print "$last_chr\t$pos\t$strand\t$meth\t$nonmeth\t$tri_nt\t$di_nt\n";
	    }
	  }
	  else{
	    if ($zero){ # zero based coordinates
	      $pos -= 1;
	      print "$last_chr\t$pos\t$strand\t$meth\t$nonmeth\t$tri_nt\n";
	    }
	    else{ # default
	      print "$last_chr\t$pos\t$strand\t$meth\t$nonmeth\t$tri_nt\n";
	    }
	  }
	}
      }

      %chr = (); # resetting the hash

      # new first entry
      $last_chr = $chr;
      $chr{$chr}->{$start}->{meth} = $meth;
      $chr{$chr}->{$start}->{nonmeth} = $nonmeth;
    }
  }

  # Last found chromosome
  warn "Writing out detailed cytosine reports for chromosome $last_chr\n";
  warn "Stored ",scalar keys $chr{$last_chr}," positions for chromosome $last_chr\n";

  while ( $chromosomes{$last_chr} =~ /([CG])/g){
	
    my $tri_nt = '';
    my $di_nt= '';
    my $pos = pos$chromosomes{$last_chr};

    my $strand;
    my $meth = 0;
    my $nonmeth = 0;

    if ($1 eq 'C'){    # C on forward strand
      $tri_nt = substr ($chromosomes{$last_chr},($pos-1),3);   # positions are 0-based!
      $strand = '+';
    }
    elsif ($1 eq 'G'){ # C on reverse strand
      $tri_nt = substr ($chromosomes{$last_chr},($pos-3),3);   # positions are 0-based!
      $tri_nt = reverse $tri_nt;
      $tri_nt =~ tr/ACTG/TGAC/;
      $strand = '-';
    }

    $di_nt = substr ($tri_nt,0,2);

    if (exists $chr{$last_chr}->{($pos-1)}){ # stored positions are 0-based!
      $meth =  $chr{$last_chr}->{$pos-1}->{meth};
      $nonmeth = $chr{$last_chr}->{$pos-1}->{nonmeth};
    }
	
    if ($CpG_only){ # default

      if ($di_nt eq 'CG'){
	if ($dinucleotide_context){
	  if ($zero){ # zero based coordinates
	    $pos -= 1;
	    print "$last_chr\t$pos\t$strand\t$meth\t$nonmeth\t$tri_nt\t$di_nt\n";
	  }
	  else{ # default
	    print "$last_chr\t$pos\t$strand\t$meth\t$nonmeth\t$tri_nt\t$di_nt\n";
	  }
	}
	else{ #default
	  if ($zero){ # zero based coordinates
	    $pos -= 1;
	    print "$last_chr\t$pos\t$strand\t$meth\t$nonmeth\t$tri_nt\n";
	  }
	  else{ # default
	    print "$last_chr\t$pos\t$strand\t$meth\t$nonmeth\t$tri_nt\n";
	  }
	}
      }

    }
    else{ ## all cytosines
      if ($dinucleotide_context){
	if ($zero){ # zero based coordinates
	  $pos -= 1;
	  print "$last_chr\t$pos\t$strand\t$meth\t$nonmeth\t$tri_nt\t$di_nt\n";
	}
	else{ # default
	  print "$last_chr\t$pos\t$strand\t$meth\t$nonmeth\t$tri_nt\t$di_nt\n";
	}
      }
      else{
	if ($zero){ # zero based coordinates
	  $pos -= 1;
	  print "$last_chr\t$pos\t$strand\t$meth\t$nonmeth\t$tri_nt\n";
	}
	else{ # default
	  print "$last_chr\t$pos\t$strand\t$meth\t$nonmeth\t$tri_nt\n";
	}
      }
    }
  }
}

sub read_genome_into_memory{

  ## reading in and storing the specified genome in the %chromosomes hash
  chdir ($genome_folder) or die "Can't move to $genome_folder: $!";
  warn "Now reading in and storing sequence information of the genome specified in: $genome_folder\n\n";

  my @chromosome_filenames =  <*.fa>;

  ### if there aren't any genomic files with the extension .fa we will look for files with the extension .fasta
  unless (@chromosome_filenames){
    @chromosome_filenames =  <*.fasta>;
  }
  unless (@chromosome_filenames){
    die "The specified genome folder $genome_folder does not contain any sequence files in FastA format (with .fa or .fasta file extensions)\n";
  }

  foreach my $chromosome_filename (@chromosome_filenames){

    # skipping the tophat entire mouse genome fasta file
    next if ($chromosome_filename eq 'Mus_musculus.NCBIM37.fa');

    open (CHR_IN,$chromosome_filename) or die "Failed to read from sequence file $chromosome_filename $!\n";
    ### first line needs to be a fastA header
    my $first_line = <CHR_IN>;
    chomp $first_line;
    $first_line =~ s/\r//; # removing /r carriage returns

    ### Extracting chromosome name from the FastA header
    my $chromosome_name = extract_chromosome_name($first_line);
	
    my $sequence;
    while (<CHR_IN>){
      chomp;
      $_ =~ s/\r//; # removing /r carriage returns

      if ($_ =~ /^>/){
	### storing the previous chromosome in the %chromosomes hash, only relevant for Multi-Fasta-Files (MFA)
	if (exists $chromosomes{$chromosome_name}){
	  warn "chr $chromosome_name (",length $sequence ," bp)\n";
	  die "Exiting because chromosome name already exists. Please make sure all chromosomes have a unique name!\n";
	}
	else {
	  if (length($sequence) == 0){
	    warn "Chromosome $chromosome_name in the multi-fasta file $chromosome_filename did not contain any sequence information!\n";
	  }
	  warn "chr $chromosome_name (",length $sequence ," bp)\n";
	  $chromosomes{$chromosome_name} = $sequence;
	}
	### resetting the sequence variable
	$sequence = '';
	### setting new chromosome name
	$chromosome_name = extract_chromosome_name($_);
      }
      else{
	$sequence .= uc$_;
      }
    }

    if (exists $chromosomes{$chromosome_name}){
      warn "chr $chromosome_name (",length $sequence ," bp)\t";
      die "Exiting because chromosome name already exists. Please make sure all chromosomes have a unique name.\n";
    }
    else{
      if (length($sequence) == 0){
	warn "Chromosome $chromosome_name in the file $chromosome_filename did not contain any sequence information!\n";
      }
      warn "chr $chromosome_name (",length $sequence ," bp)\n";
      $chromosomes{$chromosome_name} = $sequence;
    }
  }
  warn "\n";
  chdir $parent_dir or die "Failed to move to directory $parent_dir\n";
}

sub extract_chromosome_name {
  ## Bowtie extracts the first string after the inition > in the FASTA file, so we are doing this as well
  my $fasta_header = shift;
  if ($fasta_header =~ s/^>//){
    my ($chromosome_name) = split (/\s+/,$fasta_header);
    return $chromosome_name;
  }
  else{
    die "The specified chromosome ($fasta_header) file doesn't seem to be in FASTA format as required!\n";
  }
}

sub usage{

  warn <<EOF

  SYNOPSIS:

  This script generates a cytosine methylation report for a genome of interest and a sorted methylation input file produced
  by the script "genome_methylation2bedGraph". By default, the output uses 1-based chromosome coordinates and reports CpG
  positions only (for both strands).

  The input file needs to have been generated with the genome_methylation2bedGraph script while using the '--counts' option,
  or otherwise be exactly in the following format:
 
  <chromosome>  <start position>  <end position>  <methylation percentage>  <count methylated>  <count unmethylated>



  USAGE: genome_wide_cytosine_report.pl [options] <input> > <output>


  --genome_folder <path>  -  Enter the genome folder you wish to use to extract sequences from (full path only). Accepted
                             formats are FastA files ending with '.fa' or '.fasta'. Default path: /data/public/Genomes/Mouse/NCBIM37/.

  --CX                    -  The output file contains information on every single cytosine in the genome irrespecitve of
                             its context. This applies to both forward and reverse strands. Please be aware that this will
                             generate output files with > 1.1 billion lines for a mammalian genome such as human or mouse.
                             Default: OFF (i.e. Default = CpG context only).

  --zero_based            -  Uses zero-based coordinates like used in e.g. bed files instead of 1-based coordinates.
                             Default: OFF.

  --dinucleotide_context  -  Reporting an additional column with dinucleotide cytosine context (CG, CT, CC or CA). As this
                             is already part of the tri-nucleotide context column and uses a fair amount of extra disk space,
                             the default is: OFF.


  --help                  -  Displays this help message and exits.


  OUTPUT FORMAT:

  The default output file is tab-delimited in the following format:

  <Chromosome>  <Position>  <Strand>  <Count methylated>  <Count non-methylated>  <Trinucleotide context>

  The option '--dinucleotide_context' writes out one additional column (while recognisably increasing the file size):

  <Chromosome>  <Position>  <Strand>  <Count methylated>  <Count non-methylated>  <Tri-nucleotide context>  <Di-nucleotide context>


                          Script last modified: 20 Sept 2012.

EOF
    ;
  exit 1;
}