#!/usr/bin/perl -w 
use strict;

my $usage = "Usage\n $0 <bed file> <genome fasta> <sequence buffer>\n";
my $bedFile = $ARGV[0] or die $usage;
my $genomeFasta = $ARGV[1] or die $usage;
my $buffer = defined $ARGV[2] ? $ARGV[2] : die $usage;

my($fileBase) = $bedFile =~ /\/?(.*).bed/;
my $loci=readBedFile($bedFile);
my $genome = loadGenome($genomeFasta);
printFastaFile($loci,$genome,$fileBase,$buffer);

sub printFastaFile {
    my($loci,$genome,$fileBase,$buffer) = @_;
    my $outputFile = $fileBase . ".fasta";
    open(FASTA,">$outputFile") or die "Could not open $outputFile for writing.\n";
 
    foreach my $chrom (keys %{$loci}) {
	foreach my $locus (@{$loci->{$chrom}}) {
	    my($start,$stop)=@{$locus};
	    $start -= $buffer;
	    $stop += $buffer;
	    my $strand = "+";
	    my $location = "$chrom:$start..$stop:$strand";
	    my $sequence = getSequence($location,$genome);
	    print FASTA ">$location\n$sequence\n";
	}
    }
}

#####################
# INPUT SUBROUTINES #
#####################

sub readBedFile {
    my $bedFile = shift;
    my %bedList;
    open(BEDFILE,$bedFile);
    while(<BEDFILE>) {
	unless(/\#/) {
	    chomp;
	    my($chrom,$start,$stop)=split(/\t/,$_);	
	    $chrom =~ s/chrMT/chrM/;
	    push(@{$bedList{$chrom}},[$start,$stop]);
	}
    }
    return \%bedList;
}

sub loadGenome {
    my $genomeFile = shift;

    my %sequences;
    my $defline;
    open(GFILE,$genomeFile) or die "could not open $genomeFile\n";
    while(<GFILE>) {
        chomp;
        if(/>/) {
            s/>//g;
            my @terms = split;
	    $defline = shift(@terms);
        } else {
            $sequences{$defline} .= $_;
        }
    }
    return \%sequences;
}

sub getSequence {
    my($location,$sequences)=@_;
    my($chrom,$start,$stop,$strand)=parseLocation($location);
    if($sequences->{$chrom}) {
	#print "trying to get ", $stop-$start+1, "\n";
        my $string =  substr($sequences->{$chrom},$start-1,$stop-$start+1);
        if($strand eq "+") {
            return $string;
        } else {
	    return reverseComplement($string);
        }
    } else {
        die "could not find $chrom in function getSequence\nLoaded from $location\n";
    }
}

sub parseLocation {
    my($location)=@_;
    my($chrom,$start,$end,$strand);
    
    if($location =~ /(.*)\:(-?\d+)\-(-?\d+)\:(.*)/) {
        $chrom=$1;
        $start=$2;
        $end=$3;
        $strand=$4;
    } elsif($location =~ /(.*)\:(-?\d+)\-(-?\d+)/) {
        $chrom=$1;
        $start=$2;
        $end=$3;
        $strand="+";
    } elsif($location =~ /(.*)\:(-?\d+)\.\.(-?\d+)\:(.*)/) {
	$chrom=$1;
        $start=$2;
        $end=$3;
        $strand=$4;   	
    } elsif($location =~ /(.*)\:(-?\d+)\.\.(-?\d+)/) {
	$chrom=$1;
        $start=$2;
        $end=$3;
        $strand="+";   	
    }
    return ($chrom,$start,$end,$strand);
}
