#!/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 $zero;

GetOptions("h|help" => \$help,
	   "genome_folder=s" => \$genome_folder,
	   "CX" => \$CX_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);
}

### 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 cytosine reports for chromosome $last_chr (stored ",scalar keys %{$chr{$last_chr}}," different covered positions)\n";

      while ( $chromosomes{$last_chr} =~ /([CG])/g){
	
	my $tri_nt = '';
	my $context = '';
	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 = '-';
	}
	next if (length$tri_nt < 3); # trinucleotide sequence could not be extracted

	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};
	}

	### determining cytosine context	
	if ($tri_nt =~ /^CG/){
	  $context = 'CG';
	}
	elsif ($tri_nt =~ /^C.{1}G$/){
	  $context = 'CHG';
	}
	elsif ($tri_nt =~ /^C.{2}$/){
	  $context = 'CHH';
	}
	else{ # if the context can't be determined the positions will not be printed (it will equally not have been reported by Bismark)
	  warn "The sequence context could not be determined (found: '$tri_nt'). Skipping.\n";
	  next;
	}

	if ($CpG_only){
	  if ($tri_nt =~ /^CG/){ # CpG context is the default
	    if ($zero){ # zero based coordinates
	      $pos -= 1;
	      print join ("\t",$last_chr,$pos,$strand,$meth,$nonmeth,$context,$tri_nt),"\n";
	    }
	    else{ # default
	      print join ("\t",$last_chr,$pos,$strand,$meth,$nonmeth,$context,$tri_nt),"\n";
	    }
	  }
	}
	else{ ## all cytosines, specified with --CX
	  if ($zero){ # zero based coordinates
	    $pos -= 1;
	    print join ("\t",$last_chr,$pos,$strand,$meth,$nonmeth,$context,$tri_nt),"\n";
	  }
	  else{ # default
	    print join ("\t",$last_chr,$pos,$strand,$meth,$nonmeth,$context,$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 cytosine reports for chromosome $last_chr (stored ",scalar keys %{$chr{$last_chr}}," different covered positions)\n";

  while ( $chromosomes{$last_chr} =~ /([CG])/g){
	
    my $tri_nt;
    my $context;
    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 = '-';
    }

    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};
    }

    next if (length$tri_nt < 3); # trinucleotide sequence could not be extracted

    ### determining cytosine context	
    if ($tri_nt =~ /^CG/){
      $context = 'CG';
    }
    elsif ($tri_nt =~ /^C.{1}G$/){
      $context = 'CHG';
    }
    elsif ($tri_nt =~ /^C.{2}$/){
      $context = 'CHH';
    }
    else{ # if the context can't be determined the positions will not be printed (it will equally not have been reported by Bismark)
      warn "The cytosine context could not be determined (found: '$tri_nt'). Skipping.\n";
      next;
    }
	
    if ($CpG_only){
      if ($tri_nt =~ /^CG/){ # CpG context is the default
	if ($zero){ # zero-based coordinates
	  $pos -= 1;
	  print join ("\t",$last_chr,$pos,$strand,$meth,$nonmeth,$context,$tri_nt),"\n";
	}
	else{ # default
	  print join ("\t",$last_chr,$pos,$strand,$meth,$nonmeth,$context,$tri_nt),"\n";
	}
      }
    }
    else{ ## all cytosines, specified with --CX
      if ($zero){ # zero based coordinates
	$pos -= 1;
	print join ("\t",$last_chr,$pos,$strand,$meth,$nonmeth,$context,$tri_nt),"\n";
      }
      else{ # default
	print join ("\t",$last_chr,$pos,$strand,$meth,$nonmeth,$context,$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 irrespective 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.

  --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>  <C context>  <trinucleotide context>



                          Script last modified: 21 Sept 2012.

EOF
    ;
  exit 1;
}