#!/usr/bin/perl

use strict;


my $subdir = $ARGV[0];
my $dir = $ARGV[1];
my $combset = $ARGV[2];


my @chrs = (1..19, 'X', 'Y', 'MT');
print("@chrs\n");

my @prel_indices = (0, 1, 3, 4);
my @data_indices = (9..28);


my ($fname, $line, $start, $ch, $nskip);
my ($chr, $bppos, $ref, $alt, $Malt, $mouse, $qual, $geno, $allele_string, $depth_string);
my ($gt_id, $rd_id, $new_ind);

my (@row, @mice, @mouse_info, @geno_vec);


my $sumfile = $dir.'/alleles.'.$combset.'.sumout.txt';

open(OUS, '>', $sumfile) or die "Cannot open file $sumfile for writing, reason $!\n";
print(OUS "chr Nmult.known Nmult.new Nmult.all\n");


foreach $ch ( @chrs ) {
    print("\nchromosome $ch\n");

    my $kname = $subdir.'/genotypes.'.$combset.'.known.chr'.$ch.'.vcf';
    my $nname = $subdir.'/genotypes.'.$combset.'.new.chr'.$ch.'.vcf';
    my $aname = $subdir.'/SNPs.'.$combset.'.chr'.$ch.'.vcf';

    my $kfile = $dir.'/alleles.'.$combset.'.known.chr'.$ch.'.txt';
    my $nfile = $dir.'/alleles.'.$combset.'.new.chr'.$ch.'.txt';
    my $afile = $dir.'/alleles.'.$combset.'.chr'.$ch.'.txt';

    my $mkfile = $dir.'/alleles.'.$combset.'.multALTs.known.chr'.$ch.'.txt';
    my $mnfile = $dir.'/alleles.'.$combset.'.multALTs.new.chr'.$ch.'.txt';
    my $mafile = $dir.'/alleles.'.$combset.'.multALTs.chr'.$ch.'.txt';


    my $Nmult_knw = 0;
    my $Nmult_new = 0;
    my $Nmult_all = 0;


    open(IN, '<', $kname) or die "Cannot open file $kname, reason $!\n";

    $line = <IN>;
    $start = substr( $line, 0, 2);

    while ( $start eq '##' ) {
	$line = <IN>;
	$start = substr( $line, 0, 2);
    } # end of headers while loop

    my $header_line = substr( $line, 1);
    chomp($header_line);
#    print("$header_line\n");
    my @header = split("\t", $header_line);

    my @GT;
    my @RD;
    foreach my $id ( @header[@data_indices] ) {
	$gt_id = join('.', 'GT', $id);
	push(@GT, $gt_id);
	$rd_id = join('.', 'RD', $id);
	push(@RD, $rd_id);
    }
    my $new_header_line = join("\t", @header[@prel_indices], 'MultALTs', @GT, @RD);


    open(OUT, '>', $kfile) or die "Cannot open file $kfile for writing, reason $!\n";
    print(OUT "$new_header_line\n");

    open(OUM, '>', $mkfile) or die "Cannot open file $mkfile for writing, reason $!\n";
    print(OUM "$new_header_line\n");

    while ( defined($line = <IN>) ) {

	chomp($line);
	@row = split("\t", $line);

	$chr = $row[0];
	$bppos = $row[1];
	$ref = $row[3];
	$alt = $row[4];
	$Malt = 'N';

	@mice = @row[ 9..$#row ];

	my @alleles;
	my @depths;


	if ( length($alt) > 1 ) {
	  $Malt = 'Y';
	  $Nmult_knw++;

	  foreach $mouse ( @mice ) {
	    @mouse_info = split(/:/, $mouse);

	    $qual = $mouse_info[2];
	    $geno = $mouse_info[0];

	    push(@alleles, $geno);
	    push(@depths, $qual);

	  } # end of loop over mice

	  $allele_string = join("\t", @alleles);
	  $depth_string = join("\t", @depths);

	  print(OUM "$chr\t$bppos\t$ref\t$alt\t$Malt\t$allele_string\t$depth_string\n");

	} else {

	  foreach $mouse ( @mice ) {
	    @mouse_info = split(/:/, $mouse);

	    $qual = $mouse_info[2];

	    @geno_vec = split(/\//, $mouse_info[0]);
	    if ( $geno_vec[0] eq '.' ) {
		if ( $geno_vec[1] eq '.' ) {
		    $geno = '.';
		    $qual = '.';
		} else {
		    print("missing allele @geno_vec\n");
		    exit;
		}
	    } elsif ( $geno_vec[0] == 0 || $geno_vec[0] == 1 ) {
		if ( $geno_vec[1] eq '.' ) {
		    print("missing allele @geno_vec\n");
		    exit;
		} elsif ( $geno_vec[1] == 0 || $geno_vec[1] == 1 ) {
		    $geno = ( $geno_vec[0] + $geno_vec[1] );
		} else {
		    print("unrecognised allele $geno_vec[1]\n");
		    exit;
		}
	    } else {
		print("unrecognised allele $geno_vec[0]\n");
		exit;
	    }

	    push(@alleles, $geno);
	    push(@depths, $qual);
	} # end of loop over mice

	$allele_string = join("\t", @alleles);
	$depth_string = join("\t", @depths);

	print(OUT "$chr\t$bppos\t$ref\t$alt\t$Malt\t$allele_string\t$depth_string\n");

	} # end of mult ALT loop

    } # end of while loop

    close(IN);
    close(OUT);


    my %new;

    open(IN, '<', $nname) or die "Cannot open file $nname, reason $!\n";

    $line = <IN>;
    $start = substr( $line, 0, 2);

    while ( $start eq '##' ) {
	$line = <IN>;
	$start = substr( $line, 0, 2);
    } # end of headers while loop

    my $header_line = substr( $line, 1);
    chomp($header_line);
    my @header = split("\t", $header_line);

    my @GT;
    my @RD;
    foreach my $id ( @header[@data_indices] ) {
	$gt_id = join('.', 'GT', $id);
	push(@GT, $gt_id);
	$rd_id = join('.', 'RD', $id);
	push(@RD, $rd_id);
    }
    my $new_header_line = join("\t", @header[@prel_indices], 'MultALTs', @GT, @RD);


    open(OUT, '>', $nfile) or die "Cannot open file $nfile for writing, reason $!\n";
    print(OUT "$new_header_line\n");

    open(OUM, '>', $mnfile) or die "Cannot open file $mnfile for writing, reason $!\n";
    print(OUM "$new_header_line\n");


    while ( defined($line = <IN>) ) {

	chomp($line);
	@row = split("\t", $line);

	$chr = $row[0];
	$bppos = $row[1];
	$ref = $row[3];
	$alt = $row[4];
	$Malt = 'N';

	$new{ $bppos }{'REF'} = $ref;

	@mice = @row[ 9..$#row ];

	my @alleles;
	my @depths;


	if ( length($alt) > 1 ) {
	  $Malt = 'Y';
	  $Nmult_new++;

	  foreach $mouse ( @mice ) {
	    @mouse_info = split(/:/, $mouse);

	    $qual = $mouse_info[2];
	    $geno = $mouse_info[0];

	    push(@alleles, $geno);
	    push(@depths, $qual);

	  } # end of loop over mice

	  $allele_string = join("\t", @alleles);
	  $depth_string = join("\t", @depths);

	  print(OUM "$chr\t$bppos\t$ref\t$alt\t$Malt\t$allele_string\t$depth_string\n");

	} else {

	foreach $mouse ( @mice ) {
	    @mouse_info = split(/:/, $mouse);

	    $qual = $mouse_info[2];

	    @geno_vec = split(/\//, $mouse_info[0]);
	    if ( $geno_vec[0] eq '.' ) {
		if ( $geno_vec[1] eq '.' ) {
		    $geno = '.';
		    $qual = '.';
		} else {
		    print("missing allele @geno_vec\n");
		    exit;
		}
	    } elsif ( $geno_vec[0] == 0 || $geno_vec[0] == 1 ) {
		if ( $geno_vec[1] eq '.' ) {
		    print("missing allele @geno_vec\n");
		    exit;
		} elsif ( $geno_vec[1] == 0 || $geno_vec[1] == 1 ) {
		    $geno = ( $geno_vec[0] + $geno_vec[1] );
		} else {
		    print("unrecognised allele $geno_vec[1]\n");
		    exit;
		}
	    } else {
		print("unrecognised allele $geno_vec[0]\n");
		exit;
	    }

	    push(@alleles, $geno);
	    push(@depths, $qual);
	} # end of loop over mice

	$allele_string = join("\t", @alleles);
	$depth_string = join("\t", @depths);

	print(OUT "$chr\t$bppos\t$ref\t$alt\t$allele_string\t$depth_string\n");

	} # end of mult ALT loop

    } # end of while loop

    close(IN);
    close(OUT);



    open(IN, '<', $aname) or die "Cannot open file $aname, reason $!\n";

    $line = <IN>;
    $start = substr( $line, 0, 2);

    while ( $start eq '##' ) {
	$line = <IN>;
	$start = substr( $line, 0, 2);
    } # end of headers while loop

    my $header_line = substr( $line, 1);
    chomp($header_line);
    my @header = split("\t", $header_line);

    my @GT;
    my @RD;
    foreach my $id ( @header[@data_indices] ) {
	$gt_id = join('.', 'GT', $id);
	push(@GT, $gt_id);
	$rd_id = join('.', 'RD', $id);
	push(@RD, $rd_id);
    }
    my $new_header_line = join("\t", @header[@prel_indices], 'MultALTs', 'NEW', @GT, @RD);


    open(OUT, '>', $afile) or die "Cannot open file $afile for writing, reason $!\n";
    print(OUT "$new_header_line\n");

    open(OUM, '>', $mafile) or die "Cannot open file $mafile for writing, reason $!\n";
    print(OUM "$new_header_line\n");


    while ( defined($line = <IN>) ) {

	chomp($line);
	@row = split("\t", $line);

	$chr = $row[0];
	$bppos = $row[1];
	$ref = $row[3];
	$alt = $row[4];
	$Malt = 'N';

	if ( exists( $new{ $bppos } )) {
	    $new_ind = 'Y';
	} else {
	    $new_ind = 'N';
	}

	@mice = @row[ 9..$#row ];

	my @alleles;
	my @depths;


	if ( length($alt) > 1 ) {
	  $Malt = 'Y';
	  $Nmult_all++;

	  foreach $mouse ( @mice ) {
	    @mouse_info = split(/:/, $mouse);

	    $qual = $mouse_info[2];
	    $geno = $mouse_info[0];

	    push(@alleles, $geno);
	    push(@depths, $qual);

	  } # end of loop over mice

	  $allele_string = join("\t", @alleles);
	  $depth_string = join("\t", @depths);

	  print(OUM "$chr\t$bppos\t$ref\t$alt\t$Malt\t$allele_string\t$depth_string\n");

	} else {

	foreach $mouse ( @mice ) {
	    @mouse_info = split(/:/, $mouse);
	    
	    $qual = $mouse_info[2];

	    @geno_vec = split(/\//, $mouse_info[0]);
	    if ( $geno_vec[0] eq '.' ) {
		if ( $geno_vec[1] eq '.' ) {
		    $geno = '.';
		    $qual = '.';
		} else {
		    print("missing allele @geno_vec\n");
		    exit;
		}
	    } elsif ( $geno_vec[0] == 0 || $geno_vec[0] == 1 ) {
		if ( $geno_vec[1] eq '.' ) {
		    print("missing allele @geno_vec\n");
		    exit;
		} elsif ( $geno_vec[1] == 0 || $geno_vec[1] == 1 ) {
		    $geno = ( $geno_vec[0] + $geno_vec[1] );
		} else {
		    print("unrecognised allele $geno_vec[1]\n");
		    exit;
		}
	    } else {
		print("unrecognised allele $geno_vec[0]\n");
		exit;
	    }

	    push(@alleles, $geno);
	    push(@depths, $qual);
	} # end of loop over mice

	$allele_string = join("\t", @alleles);
	$depth_string = join("\t", @depths);

	print(OUT "$chr\t$bppos\t$ref\t$alt\t$new_ind\t$allele_string\t$depth_string\n");

	} # end of mult ALT loop

    } # end of while loop

    close(IN);
    close(OUT);

    print(OUS "$ch $Nmult_knw $Nmult_new $Nmult_all\n");

} # end of chr loop

close(OUS);
