#!/usr/bin/perl -w
#  A perl script to design MIP probes to capture CpG islands.
# USAGE: ./getCpGProbes.pl cpgIslandExt.txt chrN > chrN_H1H2_list.txt
#       cpgIslandExt.txt was downloaded from UCSC
#
# Dependencies:
#   (1) The repeat masked human genome sequence downloaded from UCSC;
#   (2) A nearest-neighbor parameter file for Tm calculation;
#   (3) A pre-compiled file containing the frequency of 12-mers in the bisulfite converted human genome;
#   (4) Three Perl modules from CPAN (Bioperl, File::Temp, Sort::Array)
#
# This perl script requires at least 4Gb of memory. It runs quite slow, so it is 
#   recommended to do probe design for one chromosome at a time. 
#
# Written by Kun Zhang (kzhang@bioeng.ucsd.edu), March 2008
#

use strict;

#NEED TO SET THE PATH FOR PERL LIBRARY IF NOT INSTALLED IN THE DEFAULT DIRECTORIES
#use lib qw(/home/kun/PerlLib);
use Bio::SeqIO;
use File::Temp qw/ tempfile tempdir /;
use Sort::Array qw(
      Sort_Table
      Discard_Duplicates
  );

#THE FOLLOWING VARIABLES NEED TO BE CHANGED
my $HsDir="/home/kun/HsGenome/";
my $NNParamFile = "/home/kun/HsGenome/NN_param.txt";
my $KmerDistFile = "/home/kun/HsGenome/unmasked/HsGenome18.bis.S12.dist";


my $KmerSize = 12;

my $primerMaxLen = 30;
my $primerMinLen = 20;
my $primerOptLen = 25;
my $primerMaxTm = 67;
my $primerMinTm = 62;
my $targetOptLen = 200;
my $targetMaxLen = 225;
my $targetMinLen = 175;

my %cmpTable = ("A", "T", "T", "A", "G", "C", "C", "G", "N", "N");
my $flankingLen = 50;
#Some constants
my $R = 1.987;
my %deltaH;
my %deltaS;
my @kMerDist;
my %oligoFreqBuffer;
&readNNParam($NNParamFile);
&readKmerDist($KmerDistFile);

#open the CpG islands file and load all information into a table;
my %cpgTable;
open(INFILE, $ARGV[0]) || die("Error in opening gene file $ARGV[0]!");
while(my $line = <INFILE>){
	next if($line !~ /^chr/);
	chop($line);
	my @fields = split(/[\t ]+/, $line);	
	next if($fields[0] =~ /random/);
	my $start = sprintf("%09d", $fields[1]);
	$cpgTable{$fields[0]}->{$start}->{'end'}=$fields[2];
}
close(INFILE);

#extract genomic sequences from the human genome
my ($seq_in,$seqObj);
my $last_chr='chr';
my ($totalTargetSize, $totalUniqueTargetSize, $totalCapturedTargets, $totalUniqueCapturedTargets);
foreach my $chr (keys(%cpgTable)){
	next if($ARGV[1] && $chr ne $ARGV[1]);
	if ($chr ne $last_chr)	{
		undef($seqObj);
		$seq_in = Bio::SeqIO->new(-file => $HsDir . $chr.".fa.masked" , -format =>'Fasta');	
		$seqObj = $seq_in->next_seq();
		$last_chr = $chr;
	}
	
	foreach my $cpgStart (sort keys(%{$cpgTable{$chr}})){
		my $cpgEnd = $cpgTable{$chr}->{$cpgStart}->{'end'};
		my $gSeq = $seqObj->subseq($cpgStart-$flankingLen, $cpgEnd+$flankingLen);
		
		my $N_counts=0;
		my $gSeqLen = length($gSeq);
		for (my $i=0; $i<$gSeqLen; $i++){
			$N_counts ++ if(substr($gSeq,$i,1) eq 'N');
		}
		$totalTargetSize += $gSeqLen;
		$totalUniqueTargetSize += $gSeqLen - $N_counts;
		next if($N_counts > $gSeqLen *0.8); #ignore CGI that contain >80% repetitive sequences
		
		my @probeInfo = &getProbes($gSeq);
		
		my ($setCoverage, $setUniqueCoverage) = &calcSetCoverage(\@probeInfo, $gSeq);
		$totalCapturedTargets += $setCoverage;
		$totalUniqueCapturedTargets += $setUniqueCoverage;
		
		foreach my $probe (@probeInfo){		
			print "$chr:$cpgStart-$cpgEnd\t$probe\n";			
		}
		
		#trim the oligo buffer when it grows to over 5000 entries
		my @entry = keys(%oligoFreqBuffer);
		if(@entry > 5000){
			for(my $i=0; $i<2000; $i++){
				my $index = int(rand(5000));
				delete($oligoFreqBuffer{$entry[$index]}) if($oligoFreqBuffer{$entry[$index]});
			}
		}
	}
}

print "total_CGI=$totalTargetSize\ttotal_unique_CGI=$totalUniqueTargetSize\ttotal_targets=$totalCapturedTargets\ttotal_unique_targets=$totalUniqueCapturedTargets\n";

# Find a set of probes that maximizes the coverage of the input sequence. 
sub getProbes(){
	my $seq = shift;
	my ($btFwdSeq, $btRevSeq) = &bisulfiteCvt($seq);
	my @probeListFwd=&getProbeList($btFwdSeq, 'w');	
	my @probeListRev=&getProbeList($btRevSeq, 'c');	

	return if(scalar(@probeListFwd)<1 || scalar(@probeListRev) < 1);
	my @candidateProbeSets;
	
	if(scalar(@probeListFwd)){
		for(my $i=0; $i<scalar(@probeListFwd); $i++){
			my @currentSet;
			push(@currentSet, $probeListFwd[$i]);
			for(my $j=$i+1; $j<scalar(@probeListFwd); $j++){
				if(&isCompatibleProbes(\@currentSet, $probeListFwd[$j])==1){
					push(@currentSet, $probeListFwd[$j]);
				}	
			}
			for(my $j=0; $j<scalar(@probeListRev); $j++){
				if(&isCompatibleProbes(\@currentSet, $probeListRev[$j])==1){
					push(@currentSet, $probeListRev[$j]);
				}	
			}
			my $setScore = &calcSetScore(\@currentSet);
			push(@candidateProbeSets, join(";", @currentSet) . "\t$setScore");
			undef(@currentSet);		
		}
	}
	
	if(scalar(@probeListRev)){
		for(my $i=0; $i<scalar(@probeListRev); $i++){
			my @currentSet;
			push(@currentSet, $probeListRev[$i]);
			for(my $j=$i+1; $j<scalar(@probeListRev); $j++){
				if(isCompatibleProbes(\@currentSet, $probeListRev[$j])==1){
					push(@currentSet, $probeListRev[$j]);
				}	
			}
			for(my $j=0; $j<scalar(@probeListFwd); $j++){
				if(isCompatibleProbes(\@currentSet, $probeListFwd[$j])==1){
					push(@currentSet, $probeListFwd[$j]);
				}	
			}
			my $setScore = calcSetScore(\@currentSet);
			push(@candidateProbeSets, join(";", @currentSet) . "\t$setScore");
			undef(@currentSet);		
		}
	}
	
	return if(scalar(@candidateProbeSets)<1);
	
	my @sorted_probeSets = Sort_Table(
                            	cols      => '2',
                                field     => '2',
                                sorting   => 'descending',
                                structure => 'csv',
                                separator => '\t',
                                data      => \@candidateProbeSets);	
   #extract primers
   	my ($bestSet, $bestScore) = split(/\t/, $sorted_probeSets[0]);
   	my @bestProbes = split(/;/, $bestSet);
   	my $seqLen = length($btRevSeq);
   	my @probeInfo;
	foreach my $probe (@bestProbes){
	   	my ($a_start, $a_len, $a_Tm, $a_oligo, $b_start, $b_len, $b_Tm, $b_oligo, $score, $strand) = split(/,/, $probe);
	   	my ($LP, $TmLP, $RP, $TmRP, $target);
		if($strand =~ /w/i){
			$LP = $a_oligo;
			$RP = $b_oligo;

			$target = substr($btFwdSeq, $a_start, $b_start-$a_start+$b_len);
			$TmLP = $a_Tm;
			$TmRP = $b_Tm;
		}else{
			$LP = $b_oligo;
			$RP = $a_oligo;
			$target = substr($btRevSeq, $seqLen-$b_start-$b_len, $b_start-$a_start+$b_len);
			$TmLP = $b_Tm;
			$TmRP = $a_Tm;
		}
		my $targetStart = $a_start+$a_len;
		push(@probeInfo,  "$targetStart\t$b_start\t$LP\t$TmLP\t$RP\t$TmRP\t$score\t$strand\t$target");
   	}
   	my @sorted_probeInfo = Sort_Table(
                            	cols      => '9',
                                field     => '1',
                                sorting   => 'ascending',
                                structure => 'csv',
                                separator => '\t',
                                data      => \@probeInfo);	
   	return @sorted_probeInfo;
}

# Calculate the coverage for a set of probes.
#   The score is the sum of the scores for all probes, plus the total length of the targets coverred multiplied by a weight.

sub calcSetCoverage(){
	my $h_probeInfo = shift;
	my $gSeq = shift;
	my @baseCoverred;
	my $totalCoverage = 0;
	my $totalUniqueCoverage = 0;
	foreach my $probe (@{$h_probeInfo}){
	   	my ($a_start, $b_start, $a_seq, $a_Tm, $b_seq, $b_Tm, $score, $strand, $target) = split(/\t/, $probe);
		for(my $i=$a_start; $i<$b_start; $i++){
				$baseCoverred[$i]=1;
		}	
	}
	for(my $i=0; $i<scalar(@baseCoverred); $i++){
		if($baseCoverred[$i]){
			$totalCoverage += 1;
			$totalUniqueCoverage += 1 if(substr($gSeq,$i,1) ne 'N');
		}
	}
	return ($totalCoverage, $totalUniqueCoverage);	
}

# Calculate the combined score for a set of probes.
#   The score is the sum of the scores for all probes, plus the total length of the targets coverred multiplied by a weight.

sub calcSetScore(){
	my $h_probeSets = shift;
	my @fwdBaseCoverred;
	my @revBaseCoverred;

	my $totalScore = 0;
	foreach my $probe (@{$h_probeSets}){
		next if(!$probe);
	   	my ($a_start, $a_len, $a_Tm, $a_oligo, $b_start, $b_len, $b_Tm, $b_oligo, $score, $strand) = split(/,/, $probe);
	   	next if(!$a_len);
		for(my $i=$a_start+$a_len; $i<$b_start; $i++){
				$fwdBaseCoverred[$i]=1 if($strand =~ /w/i);
				$revBaseCoverred[$i]=1 if($strand =~ /c/i);
		}	
		$totalScore += $score;
	}
	for(my $i=0; $i<scalar(@fwdBaseCoverred); $i++){
		$totalScore += $fwdBaseCoverred[$i]*5 if($fwdBaseCoverred[$i]);
	}
	for(my $i=0; $i<scalar(@revBaseCoverred); $i++){
		$totalScore += $revBaseCoverred[$i]*5 if($revBaseCoverred[$i]);
	}
	return $totalScore;	
}

# Check whether a new probe is compatible with an existing probe set or not.		
# Rules:
#  1. No overlap between two probes on the same strand
#  2. A probe should cover a gap of at least $minBaseCoverred bases on the opposite strand
#
sub isCompatibleProbes(){
	my $h_probeSets = shift;
	my $new_probe = shift;
	my $minBaseCoverred = 50;
	return -1 if(!$new_probe);
	
	my (@fwdBaseCoverred, @revBaseCoverred);
	
	my  ($a_start, $a_len, $a_Tm, $a_oligo, $b_start, $b_len, $b_Tm, $b_oligo, $score, $strand);
	foreach my $probe (@{$h_probeSets}){
		next if(!$probe);
		($a_start, $a_len, $a_Tm, $a_oligo, $b_start, $b_len, $b_Tm, $b_oligo, $score, $strand) = split(/,/, $probe);
		next if($strand !~ /[wc]/i || $a_start<0 || $b_start<0);
		if($strand =~ /w/i){
			for(my $i=0; $i<$a_len; $i++){
				$fwdBaseCoverred[$i+$a_start]=-1;
			}	
			for(my $i=0; $i<$b_len; $i++){
				$fwdBaseCoverred[$i+$b_start]=-1;
			}		
			for(my $i=$a_start+$a_len; $i<$b_start; $i++){
				$fwdBaseCoverred[$i]=1;
			}	
		}else{
			for(my $i=0; $i<$a_len; $i++){
				$revBaseCoverred[$i+$a_start]=-1;
			}	
			for(my $i=0; $i<$b_len; $i++){
				$revBaseCoverred[$i+$b_start]=-1;
			}		
			for(my $i=$a_start+$a_len; $i<$b_start; $i++){
				$revBaseCoverred[$i]=1;
			}	
		}
	} 
	
	($a_start, $a_len, $a_Tm, $a_oligo, $b_start, $b_len, $b_Tm, $b_oligo, $score, $strand) = split(/,/, $new_probe);
	return -1 if($strand !~ /[wc]/i);

	if($strand =~ /w/i){
		for(my $i=$a_start; $i<$b_start+$b_len; $i++){
			return -1 if($fwdBaseCoverred[$i]);
		}
		my $base_coverred=0;
		for(my $i=$a_start+$a_len; $i<$b_start; $i++){
			$base_coverred++ if(!$revBaseCoverred[$i] || $revBaseCoverred[$i] != 1);
		}
		return -1 if($base_coverred < $minBaseCoverred);
	}else{
		for(my $i=$a_start; $i<$b_start+$b_len; $i++){
			return -1 if($revBaseCoverred[$i]);
		}
		my $base_coverred=0;
		for(my $i=$a_start+$a_len; $i<$b_start; $i++){
			$base_coverred++ if(!$fwdBaseCoverred[$i] || $fwdBaseCoverred[$i] != 1);
		}  
		return -1 if($base_coverred < $minBaseCoverred);
	}
	return 1;	
}

# Generate a list of candidate probes which contain two capturing arms of 
#   appropriate Tm within a appropriate distance
#
sub getProbeList(){
	my $bis_seq = shift;
	my $strand = shift;
	my $seqLen = length($bis_seq);
	my @oligoList;
	my %bestProbes;
	for(my $i=0; $i<$seqLen; $i++){
		for(my $j=$primerMinLen; $j<=$primerMaxLen; $j++){
			my $oligo = substr($bis_seq, $i, $j);
			my @CGs = split(/C/, $oligo);
			last if($oligo =~ /[Nn]/          #no Ns (repeat masked sequences)
					|| $oligo =~ /GGGG/       #no four Gs in a run
					|| scalar(@CGs>3));      #at most two CpGs.
			my $Tm = shortOligoTm($oligo, 200, 1.5, 50, 0.2, 0.0, 50);
			next if($Tm < $primerMinTm 
					|| $Tm > $primerMaxTm);
					
			my @polyTs = split(/TTTTT/, $oligo); #calculate the number of poly-T runs (more than four Ts).
			my @polyAs = split(/AAAAA/, $oligo); #calculate the number of poly-A runs (more than four As).

			my $score =  $j-$primerMinLen -100*(scalar(@CGs)-1) - 50*(scalar(@polyTs)-1) - 50*(scalar(@polyAs)-1);
			my $oligo_len = length($oligo);
			push(@oligoList, "$i,$oligo_len,$score,$Tm,$oligo");
		}
	}
	my @probeList;
	for(my $i=0; $i<scalar(@oligoList); $i++){
		my ($a_start, $a_len, $a_score, $a_Tm, $a_oligo) = split(/,/, $oligoList[$i]);
		for(my $j=$i+1; $j<scalar(@oligoList); $j++){
			my ($b_start, $b_len, $b_score, $b_Tm, $b_oligo) = split(/,/, $oligoList[$j]);
			my $gap = $b_start-$a_start-$a_len;
			next if($gap < $targetMinLen 
					|| $gap>$targetMaxLen
					|| $a_oligo =~ /AAAAAA$/
					|| $a_oligo =~ /TTTTTT$/
					|| $b_oligo =~ /^AAAAAA/
					|| $b_oligo =~ /^TTTTTT/);
			

			my $p_score = $a_score+$b_score-abs($targetOptLen-$gap) - 5*($b_Tm - $primerMinTm + $primerMaxTm - $a_Tm);
			my $a_freq = oligoFreq($a_oligo,$KmerSize,'R');
			my $b_freq = oligoFreq($b_oligo,$KmerSize,'L')*2; #more weight for the extension arm
			$p_score -= log($a_freq+$b_freq)*50;
			
			my $CG_penalty=0;
			my $weight = -500;
			for(my $k=$a_len-2; $k>=0; $k--){
				$CG_penalty += $weight if(substr($a_oligo,$k,2) eq "CG");
				$weight *= 0.75;
			}
			$weight = -500;
			for(my $k=0; $k<$b_len-1; $k++){
				$CG_penalty += $weight if(substr($b_oligo,$k,2) eq "CG");
				$weight *= 0.75;
			}
			$p_score += $CG_penalty;
			
			my $candidate = ($strand =~ /w/i) ? 
			      sprintf("%d,%d,%d,%s,%d,%d,%d,%s,%d,W", $a_start, $a_len, $a_Tm, $a_oligo, $b_start, $b_len, $b_Tm, $b_oligo, $p_score) 
			     :sprintf("%d,%d,%d,%s,%d,%d,%d,%s,%d,C", $seqLen-$b_start-$b_len, $b_len, $b_Tm, $b_oligo, $seqLen-$a_start-$a_len, $a_len, $a_Tm, $a_oligo, $p_score);
			
			my $bin = int(($a_start+$b_start+$b_len)/20);
			if(!$bestProbes{$bin} || $bestProbes{$bin}->{'s'} < $p_score){
				$bestProbes{$bin}->{'p'}=$candidate;
				$bestProbes{$bin}->{'s'}=$p_score;
			}
		}
	}
	foreach my $bin (keys(%bestProbes)){
		push(@probeList, $bestProbes{$bin}->{'p'});
	}
	
	my @sorted_probeList = Sort_Table(
                            	cols      => '10',
                                field     => '9',
                                sorting   => 'descending',
                                structure => 'csv',
                                separator => ',',
                                data      => \@probeList);
                                
	return @sorted_probeList;
}

# Generate a reverse complementary sequence.
sub RevComp(){
    my $seq = shift;
	my $seqLen = length($seq);
	my $revcom = $cmpTable{substr($seq,0,1)};
	for(my $i = 1; $i < $seqLen; $i++){
		die($seq . ": $i=". substr($seq, $i, 1)) if(!$cmpTable{substr($seq, $i, 1)});
		$revcom = $cmpTable{substr($seq, $i, 1)} . $revcom;
	}
    return $revcom;
}

#In silico bisulfite conversion
#CG is not converted since whether "C" will be converted depends on the methylation status.
sub bisulfiteCvt(){
	my $fwdSeq = shift;
	$fwdSeq = uc($fwdSeq);
	my $revSeq = &RevComp($fwdSeq);
	$fwdSeq =~ s/CG/NG/g;
	$fwdSeq =~ s/C/T/g;
	$fwdSeq =~ s/NG/CG/g;
	
	$revSeq =~ s/CG/NG/g;
	$revSeq =~ s/C/T/g;
	$revSeq =~ s/NG/CG/g;
	return($fwdSeq, $revSeq);
}

sub shortOligoTm(){
        my $seq = shift;
        my $C_primer = shift;           # nM
        my $C_Mg = shift;               # mM
        my $C_MonovalentIon = shift;    #mM
        my $C_dNTP = shift;             #mM
        my $percentage_DMSO = shift;
        my $percentage_annealed = shift; #percentage of templates that anneal to primers

        $seq =~ s/[ \t\n]+//g;
        $percentage_annealed = 50.0 if (!$percentage_annealed);
        $percentage_annealed /= 100.0;

        my $C_SodiumEquivalent = $C_MonovalentIon + 120 * sqrt($C_Mg-$C_dNTP);
        my $seqLength = length($seq);
        my $dH = $deltaH{'pm'}->{substr($seq, 0, 1)} + $deltaH{'pm'}->{substr($seq, $seqLength-1, 1)};
        my $dS = $deltaS{'pm'}->{substr($seq, 0, 1)} + $deltaS{'pm'}->{substr($seq, $seqLength-1, 1)};
        $seq = uc($seq);
        for(my $i = 0; $i < $seqLength - 1; $i ++){
                $dH += $deltaH{'pm'}->{substr($seq, $i, 2)};
                $dS += $deltaS{'pm'}->{substr($seq, $i, 2)};
        }
        $dS += 0.368 * $seqLength * log($C_SodiumEquivalent/1000.0);
        my $Tm = sprintf("%5.2f", ($dH * 1000) / ($dS + $R * (log($C_primer*(1-$percentage_annealed)/$percentage_annealed)-21.4164)) - 273.15 - 0.75*$percentage_DMSO);
        return $Tm;
}

# Load the dH, dG and dS parameters for Tm calculation.
sub readNNParam(){
	my $NNParamFile = shift;
	open(NNFILE, $NNParamFile) || die("Error in opening NN parameter file!");
	my $line = <NNFILE>;
	while($line = <NNFILE>){
		$line =~ s/\r//;
		chop($line);
		my ($seqF, $seqR, $dH, $dS, $mismatch) = split(/[:\t]/, $line);
		if(!$mismatch){
			$deltaH{'pm'}->{$seqF} = $dH;
			$deltaS{'pm'}->{$seqF} = $dS;
		}else{
			$deltaH{'mm'}->{$seqF}->{$seqR} = $dH;
			$deltaS{'mm'}->{$seqF}->{$seqR} = $dS;			
		}
	}
	close(NNFILE);
}

# Load the pre-computed K-mer frequencies
sub readKmerDist(){
	my $KmerDistFile = shift;
	open(KMERFILE, $KmerDistFile) || die("Error in opening Kmer Distribution file!");
	while(my $line = <KMERFILE>){
		chop($line);
		my ($decCode, $count, $lFreq, $rFreq) = split(/[\t ]/, $line);
		$kMerDist[$decCode]=$count;
	}
	close(KMERFILE);
}

# Convert DNA sequence to a decimal code
sub seq2dec(){
	my $seq = shift;	
	my ($i,$weight, $decCode, $seqLen);
	$weight = 1;
	$decCode = 0;
	$seqLen = length($seq);	
	my %base2value;
	$base2value{'A'}=0;
	$base2value{'T'}=1;
	$base2value{'G'}=2;
	$base2value{'C'}=3;
	while(length($seq) > 0){
		$decCode += $weight * $base2value{chop($seq)};
		$weight *= 4;
	}
	return $decCode;
}

# Estimate the abundance of a short oligo sequence based on the distribution of k-mers
sub oligoFreq(){
	my ($seq, $wordSize, $direction) = @_;
	my $seqLen = length($seq);
	my $range = $seqLen-$wordSize+1;
	my $weight = 1;
	my $freq;	
	if($direction =~ /L/){
		if($oligoFreqBuffer{$seq."L"}){
			$freq=$oligoFreqBuffer{$seq."L"};
		}else{
			for(my $i=0; $i<$range; $i++){
				my $index = &seq2dec(substr($seq,$i,$wordSize));			
				my $oligoFreq = $kMerDist[$index];
				$freq += $oligoFreq*$weight if($oligoFreq);
				$weight *= 0.75;
			}
			$freq = $freq ? $freq/1000 : 1;
			$oligoFreqBuffer{$seq."L"} = $freq;
		}
	}else{
		if($oligoFreqBuffer{$seq."R"}){
			$freq=$oligoFreqBuffer{$seq."R"};
		}else{
			for(my $i=$range-1; $i>=0; $i--){
				my $index = &seq2dec(substr($seq,$i,$wordSize));
				my $oligoFreq = $kMerDist[$index];
				$freq += $oligoFreq*$weight if($oligoFreq);
				$weight *= 0.75;
			}
			$freq = $freq ? $freq/1000 : 1;
			$oligoFreqBuffer{$seq."R"} = $freq;
		}
	}
	return $freq;
}
