#!/usr/bin/env perl

use strict;
use warnings;
use Carp;
use Vcf;

my $opts = parse_params();
do_consensus($opts);

exit;

#--------------------------------

sub error
{
    my (@msg) = @_;
    if ( scalar @msg )
    {
        croak @msg;
    }
    die
        "Usage: cat ref.fa | vcf-consensus [OPTIONS] in.vcf.gz > out.fa\n",
        "Options:\n",
        "   -h, -?, --help                   This help message.\n",
        "   -H, --haplotype <int>            Apply only variants for the given haplotype (1,2)\n",
        "   -s, --sample <name>              If not given, all variants are applied\n",
        "Examples:\n",
        "   samtools faidx ref.fa 8:11870-11890 | vcf-consensus in.vcf.gz > out.fa\n",
        "\n";
}


sub parse_params
{
    my $opts = { };
    while (my $arg=shift(@ARGV))
    {
        if ( $arg eq '-?' || $arg eq '-h' || $arg eq '--help' ) { error(); }
        if ( $arg eq '-s' || $arg eq '--sample' ) { $$opts{sample}=shift(@ARGV); next; }
        if ( $arg eq '-H' || $arg eq '--haplotype' ) { $$opts{haplotype}=shift(@ARGV); next; }
        if ( -e $arg && !exists($$opts{vcf_file}) ) { $$opts{vcf_file}=$arg; next; }
        error("Unknown parameter \"$arg\". Run -h for help.\n");
    }
    if ( exists($$opts{haplotype}) && !exists($$opts{sample}) ) { error("Expected -s option with -H.\n"); }
    return $opts;
}

sub do_consensus
{
    my ($opts) = @_;
    
    my $vcf = Vcf->new(file=>$$opts{vcf_file});
    $vcf->parse_header;
    if ( exists($$opts{sample}) )
    {
        if ( !exists($$vcf{has_column}{$$opts{sample}}) ) { error("No such sample: $$opts{sample}"); }
        $$opts{vcf} = $vcf; 
        $$opts{sample_col} = $$vcf{has_column}{$$opts{sample}};
    }
    my $chrs = $vcf->get_chromosomes();
    my %chrs = map { $_=>0 } @$chrs;

    my ($chr,$vcf_pos,$warned,$vcf_line);
    while (my $line=<STDIN>)
    {
        if ( $line=~/^>([^:\s]+)/ ) 
        {
            flush_fa_buffer($opts,0);
            $chr = $1;
            my $rest = $';
            $$opts{fa_pos} = ($rest=~/^:(\d+)-\d+$/) ? $1 : 1;
            $$opts{fa_idx} = 0;
            $$opts{fa_frz} = 0;
            if ( exists($chrs{$chr}) ) { $chrs{$chr}=1; }
            my $region = $$opts{fa_pos} > 1 ? "$chr:$$opts{fa_pos}" : $chr;
            $vcf->open(region=>$region);
            print $line;
            next;
        }

        chomp($line);
        $$opts{fa_buf} .= $line;
        $$opts{fa_len} += length($line);

        while ( defined($vcf_line = $vcf->next_data_array()) )
        {
            # can the beginning of the buffer be printed?
            if ( $$opts{fa_pos}+$$opts{fa_len}-$$opts{fa_idx}<=$$vcf_line[1] )
            {
                $vcf->_unread_line($vcf_line);
                flush_fa_buffer($opts,60);
                last;
            }
            # is the buffer long enough?
            if ( $$opts{fa_pos}+$$opts{fa_len}-$$opts{fa_idx}<=$$vcf_line[1]+length($$vcf_line[3]) )
            {
                $vcf->_unread_line($vcf_line);
                last;
            }
            apply_variant($opts,$vcf_line);
        }

        if ( !defined $vcf_line ) { flush_fa_buffer($opts,60); }
    }
    flush_fa_buffer($opts,0);

    for my $chr (keys %chrs)
    {
        if ( !$chrs{$chr} ) { warn("The sequence \"$chr\" not found in the fasta file.\n"); }
    }
}

sub flush_fa_buffer
{
    my ($opts,$len) = @_;
    while ( $$opts{fa_len} && $$opts{fa_len}>=60 )
    {
        print substr($$opts{fa_buf},0,60,''), "\n";
        $$opts{fa_len} -= 60;
        $$opts{fa_pos} += 60 - $$opts{fa_idx};
        $$opts{fa_idx}  = 0;
    }
    if ( $len or !$$opts{fa_len} ) { return; }
    print $$opts{fa_buf},"\n";
    $$opts{fa_pos} += $$opts{fa_len}-$$opts{fa_idx};
    $$opts{fa_len} = 0;
    $$opts{fa_buf} = '';
    $$opts{fa_idx} = 0;
}

sub apply_variant
{
    my ($opts,$vline) = @_;

    if ( $$vline[4] eq '.' ) { return; }

    my $hap = exists($$opts{haplotype}) ? $$opts{haplotype} : 0;
    my $alt;
    if ( !exists($$opts{sample_col}) )
    {
        my $idx;
        $alt = ($idx=index($$vline[4],','))==-1 ? $$vline[4] : substr($$vline[4],0,$idx);
    }
    else
    {
        my @als = $$opts{vcf}->split_gt($$vline[$$opts{sample_col}-1]);
        if ( $hap )
        {
            # Note: we are not checking the phase or phase blocks, assuming the VCF is perfect
            if ( $hap <= @als && $als[$hap-1] ne '0' ) 
            { 
                $alt = $$opts{vcf}->get_field($$vline[4],$als[$hap-1]-1,','); 
            }
        }
        else
        {
            for my $al (@als)
            {
                if ( $al eq '0' ) { next; }
                $alt = $$opts{vcf}->get_field($$vline[4],$al-1,',');
                last;
            }
        }
        if ( !defined $alt ) { return; }
    } 

    if ( $$vline[1] <= $$opts{fa_frz} )
    {
        print STDERR "Note: Conflicting variants at (or near) $$vline[0]:$$vline[1], cannot apply both.\n";
        return;
    }

    my $pos = $$vline[1] - $$opts{fa_pos} + $$opts{fa_idx};
    if ( $pos<0 or $pos>=$$opts{fa_len} ) { error("FIXME: $$vline[0]:$$vline[1] .. $$opts{fa_pos},$pos,$$opts{fa_len},$$opts{fa_frz}\n"); }

    # Sanity check
    my $ref_len = length($$vline[3]);
    if ( $$vline[3] ne substr($$opts{fa_buf},$pos,$ref_len) ) 
    { 
        error(sprintf "The fasta sequence does not match the REF at $$vline[0]:$$vline[1]. %s(%s) in .fa, %s in .vcf, frz=%d\n", 
            substr($$opts{fa_buf},$pos,$ref_len),
            substr($$opts{fa_buf},$pos+1,$ref_len+5),
            $$vline[3], $$opts{fa_frz}?$$opts{fa_frz}:0);
    }

    my $alt_len = length($alt);
    substr($$opts{fa_buf},$pos,$ref_len,$alt);
    $$opts{fa_len} += $alt_len - $ref_len;
    $$opts{fa_pos} += $ref_len;     # position with respect to the original reference sequence
    $$opts{fa_idx} += $alt_len;     # position in the modified sequence
    $$opts{fa_frz}  = $$vline[1] + $ref_len - 1;      # freeze changes until this position
}


