#!/usr/bin/perl

use strict;


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


my @chrs = (1..19, 'X', 'Y', 'MT');
#print("@chrs\n");
my @indices = (0, 1, 3, 4, 9..28);


my ($fname, $line, $start, $ch);

my ($chr, $bppos, $mouse, $checksum, $den);
my ($count_ref, $count_alt, $count_mis, $ind_ref, $ind_alt, $ind_mis);
my ($count_hom_ref, $count_hom_alt, $count_het, $count_geno_mis, $geno_ref, $geno_het, $geno_alt, $geno_mis);
my ($freq_ref, $freq_alt, $freq_hom_ref, $freq_hom_alt, $freq_het, $prop_miss_all, $prop_miss_geno);
my (@row, @new_row, @mice, @mouse_info, @geno_vec);

my $ts;


my %tstv;
$tstv{'A'} = 'G';
$tstv{'C'} = 'T';
$tstv{'G'} = 'A';
$tstv{'T'} = 'C';



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

    my $Kname = $subdir.'/genotypes.known.chr'.$ch.'.vcf';
    my $Nname = $subdir.'/genotypes.new.chr'.$ch.'.vcf';

    my $Kfile = $dir.'/stats.known.chr'.$ch.'.txt';
    my $Nfile = $dir.'/stats.new.chr'.$ch.'.txt';
    my $Afile = $dir.'/stats.chr'.$ch.'.txt';

    my %data;


    open(OUTK, '>', $Kfile) or die "Cannot open file $Kfile for writing, reason $!\n";

    print(OUTK  "CHR\tBP\tREF\tALT\tTS\tMAF\tMAF_allele\tProp_missing_genotypes\tProp_missing_alleles\tAllfreq_ref\tAllfreq_alt\tGenofreq_ref\tGenofreq_het\tGenofreq_alt\tAllcount_ref\tAllcount_alt\tGenocount_ref\tGenocount_het\tGenocount_alt\n");

    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);
#    my @new_header = @header[@indices];

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

	chomp($line);
	@row = split("\t", $line);
	@new_row = @row[@indices];

	if ( length($new_row[3]) == 1 ) {

	$chr = $new_row[0];
	$bppos = $new_row[1];
	$data{ $bppos }{'new'} = 'N';
	$data{ $bppos }{'REF'} = $new_row[2];
	$data{ $bppos }{'ALT'} = $new_row[3];

	if ( $tstv{ ${$data{$bppos}}{'REF'} } eq ${$data{$bppos}}{'ALT'} ) {
	    $ts = 'Y';
#    print("ref ${$data{$bppos}}{'REF'}  Ts $tstv{ ${$data{$bppos}}{'REF'} }  alt ${$data{$bppos}}{'ALT'}\n");
	} else {
#    print("ref ${$data{$bppos}}{'REF'}   alt ${$data{$bppos}}{'ALT'}\n");
	    $ts = 'N';
	}
	$data{ $bppos }{'TS'} = $ts;


	@mice = @new_row[ 4..$#new_row ];

	$count_ref = 0;
	$count_alt = 0;
	$count_mis = 0;
	$count_hom_ref = 0;
	$count_hom_alt = 0;
	$count_het = 0;
	$count_geno_mis = 0;
	foreach $mouse ( @mice ) {
	    @mouse_info = split(/:/, $mouse);
	    @geno_vec = split(/\//, $mouse_info[0]);
#	    print("genotype = @geno_vec\n");

	    $ind_ref = 0;
	    $ind_alt = 0;
	    $ind_mis = 0;
	    if ( $geno_vec[0] eq '0' ) {
		$ind_ref = ( $ind_ref + 1 );
		$count_ref = ( $count_ref + 1 );
	    } elsif ( $geno_vec[0] eq '1' ) {
		$ind_alt = ( $ind_alt + 1 );
		$count_alt = ( $count_alt + 1 );
	    } elsif ( $geno_vec[0] eq '.' ) {
		$ind_mis = ( $ind_mis + 1 );
		$count_mis = ( $count_mis + 1 );
	    } else {
		print("ERROR - unrecognised allele $geno_vec[0] at pos $bppos\n");
		exit;
	    }
	    if ( $geno_vec[1] eq '0' ) {
                $ind_ref = ( $ind_ref + 1 );
                $count_ref = ( $count_ref + 1 );
            } elsif ( $geno_vec[1] eq '1' ) {
                $ind_alt = ( $ind_alt + 1 );
                $count_alt = ( $count_alt + 1 );
            } elsif ( $geno_vec[1] eq '.' ) {
                $ind_mis = ( $ind_mis + 1 );
                $count_mis = ( $count_mis + 1 );
            } else {
                print("ERROR - unrecognised allele $geno_vec[1] at pos $bppos\n");
                exit;
            }
#	    print("ind counts (ref, alt, mis) = $ind_ref $ind_alt $ind_mis\n");

	    $geno_mis = 0;
	    $geno_ref = 0;
	    $geno_alt = 0;
	    $geno_het = 0;
	    if ( $ind_mis > 0 ) {
		$geno_mis = ( $geno_mis + 1 );
		$count_geno_mis = ( $count_geno_mis + 1 );
	    } elsif ( $ind_ref == 2 ) {
		$geno_ref = ( $geno_ref + 1 );
		$count_hom_ref = ( $count_hom_ref + 1 );
	    } elsif ( $ind_ref == 1 ) {
		$geno_het = ( $geno_het + 1 );
		$count_het = ( $count_het + 1 );
	    } elsif ( $ind_ref == 0 ) {
		$geno_alt = ( $geno_alt + 1 );
		$count_hom_alt = ( $count_hom_alt + 1 );
	    } else {
		print("ERROR - unrecognised genotype $geno_vec[0] $geno_vec[1] at pos $bppos\n");
	    }
#	    print("ind genotype counts (ref, het, alt, mis) = $geno_ref $geno_het $geno_alt $geno_mis\n");
	} # end of loop over individual mice

	$checksum = ( $count_ref + $count_alt + $count_mis );
	if ( $checksum != 40 ) {
	    print("ERROR - number of alleles $checksum at pos $bppos is not equal to 40\n");
	    print("$bppos : allele counts (ref, alt, mis) = $count_ref, $count_alt, $count_mis\n");
	}

	$data{ $bppos }{ 'allele counts' } = [ $count_ref, $count_alt, $count_mis ];
	$den = ( $checksum - $count_mis );
	$freq_ref = $count_ref/$den;
	$freq_alt = $count_alt/$den;
	$data{ $bppos }{ 'allele freqs' } = [ $freq_ref, $freq_alt ];

	if ( $freq_ref < $freq_alt ) {
	    $data{ $bppos }{ 'MAF' } = $freq_ref;
	    $data{ $bppos }{ 'MAF allele' } = 'REF';
	} else {
	    $data{ $bppos }{ 'MAF' } = $freq_alt;
	    $data{ $bppos }{ 'MAF allele' } = 'ALT';
	}
#	print("MAF = ${$data{ $bppos }}{'MAF'}  MAF allele = ${$data{ $bppos }}{'MAF allele'}\n");

	$checksum = ( $count_hom_ref + $count_het + $count_hom_alt + $count_geno_mis );
	if ( $checksum != 20 ) {
	    print("ERROR - number of genotypes $checksum at pos $bppos is not equal to 20\n");
	    print("$bppos : genotype counts (ref, het, alt, mis) = $count_hom_ref $count_het $count_hom_alt $count_geno_mis\n");
	}

	$data{ $bppos }{ 'genotype counts' } = [ $count_hom_ref, $count_het, $count_hom_alt, $count_geno_mis ];
        $den = ( $checksum - $count_geno_mis );
        $freq_hom_ref = $count_hom_ref/$den;
	$freq_het = $count_het/$den;
        $freq_hom_alt = $count_hom_alt/$den;
	$data{ $bppos }{ 'genotype freqs' } = [ $freq_hom_ref, $freq_het, $freq_hom_alt ];

	$prop_miss_all = $count_mis/40;
	$prop_miss_geno = $count_geno_mis/20;
	$data{ $bppos }{ 'missing' } = [ $prop_miss_all, $prop_miss_geno ];

#	print("missing data: alleles ${$data{ $bppos }}{'missing'}[0]  genotypes ${$data{ $bppos }}{'missing'}[1]\n");

	print(OUTK "$chr\t$bppos\t${$data{$bppos}}{'REF'}\t${$data{$bppos}}{'ALT'}\t${$data{$bppos}}{'TS'}\t${$data{$bppos}}{'MAF'}\t${$data{$bppos}}{'MAF allele'}\t${$data{$bppos}}{'missing'}[1]\t${$data{$bppos}}{'missing'}[0]\t${$data{$bppos}}{'allele freqs'}[0]\t${$data{$bppos}}{'allele freqs'}[1]\t${$data{$bppos}}{'genotype freqs'}[0]\t${$data{$bppos}}{'genotype freqs'}[1]\t${$data{$bppos}}{'genotype freqs'}[2]\t${$data{$bppos}}{'allele counts'}[0]\t${$data{$bppos}}{'allele counts'}[1]\t${$data{$bppos}}{'genotype counts'}[0]\t${$data{$bppos}}{'genotype counts'}[1]\t${$data{$bppos}}{'genotype counts'}[2]\n");

      } # end of mult ALTs loop

    } # end of while loop
    close(IN);
    close(OUTK);

    print("known stats file written\n");


    open(OUTN, '>', $Nfile) or die "Cannot open file $Nfile for writing, reason $!\n";

    print(OUTN  "CHR\tBP\tREF\tALT\tTS\tMAF\tMAF_allele\tProp_missing_genotypes\tProp_missing_alleles\tAllfreq_ref\tAllfreq_alt\tGenofreq_ref\tGenofreq_het\tGenofreq_alt\tAllcount_ref\tAllcount_alt\tGenocount_ref\tGenocount_het\tGenocount_alt\n");

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


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

      chomp($line);
      @row = split(/\t/, $line);
      @new_row = @row[@indices];

      if ( length($new_row[3]) == 1 ) {

      $chr = $new_row[0];
      $bppos = $new_row[1];
      $data{ $bppos }{'new'} = 'Y';
      $data{ $bppos }{'REF'} = $new_row[2];
      $data{ $bppos }{'ALT'} = $new_row[3];

      if ( $tstv{ ${$data{$bppos}}{'REF'} } eq ${$data{$bppos}}{'ALT'} ) {
	  $ts = 'Y';
#  print("ref ${$data{$bppos}}{'REF'}  Ts $tstv{ ${$data{$bppos}}{'REF'} }  alt ${$data{$bppos}}{'ALT'}\n");
      } else {
#  print("ref ${$data{$bppos}}{'REF'}   alt ${$data{$bppos}}{'ALT'}\n");
	  $ts = 'N';
      }
      $data{ $bppos }{'TS'} =$ts;


      @mice = @new_row[ 4..$#new_row ];

      $count_ref = 0;
      $count_alt = 0;
      $count_mis = 0;
      $count_hom_ref = 0;
      $count_hom_alt = 0;
      $count_het = 0;
      $count_geno_mis = 0;
      foreach $mouse ( @mice ) {
	@mouse_info = split(/:/, $mouse);
	@geno_vec = split(/\//, $mouse_info[0]);
	#	    print("genotype = @geno_vec\n");

	$ind_ref = 0;
	$ind_alt = 0;
	$ind_mis = 0;
	if ( $geno_vec[0] eq '0' ) {
	  $ind_ref = ( $ind_ref + 1 );
	  $count_ref = ( $count_ref + 1 );
	} elsif ( $geno_vec[0] eq '1' ) {
	  $ind_alt = ( $ind_alt + 1 );
	  $count_alt = ( $count_alt + 1 );
	} elsif ( $geno_vec[0] eq '.' ) {
	  $ind_mis = ( $ind_mis + 1 );
	  $count_mis = ( $count_mis + 1 );
	} else {
	  print("ERROR - unrecognised allele $geno_vec[0] at pos $bppos\n");
	  exit;
	}
	if ( $geno_vec[1] eq '0' ) {
	  $ind_ref = ( $ind_ref + 1 );
	  $count_ref = ( $count_ref + 1 );
	} elsif ( $geno_vec[1] eq '1' ) {
	  $ind_alt = ( $ind_alt + 1 );
	  $count_alt = ( $count_alt + 1 );
	} elsif ( $geno_vec[1] eq '.' ) {
	  $ind_mis = ( $ind_mis + 1 );
	  $count_mis = ( $count_mis + 1 );
	} else {
	  print("ERROR - unrecognised allele $geno_vec[1] at pos $bppos\n");
	  exit;
	}
	#	    print("ind counts (ref, alt, mis) = $ind_ref $ind_alt $ind_mis\n");

	$geno_mis = 0;
	$geno_ref = 0;
	$geno_alt = 0;
	$geno_het = 0;
	if ( $ind_mis > 0 ) {
	  $geno_mis = ( $geno_mis + 1 );
	  $count_geno_mis = ( $count_geno_mis + 1 );
	} elsif ( $ind_ref == 2 ) {
	  $geno_ref = ( $geno_ref + 1 );
	  $count_hom_ref = ( $count_hom_ref + 1 );
	} elsif ( $ind_ref == 1 ) {
	  $geno_het = ( $geno_het + 1 );
	  $count_het = ( $count_het + 1 );
	} elsif ( $ind_ref == 0 ) {
	  $geno_alt = ( $geno_alt + 1 );
	  $count_hom_alt = ( $count_hom_alt + 1 );
	} else {
	  print("ERROR - unrecognised genotype $geno_vec[0] $geno_vec[1] at pos $bppos\n");
	}
	#	    print("ind genotype counts (ref, het, alt, mis) = $geno_ref $geno_het $geno_alt $geno_mis\n");
      } # end of loop over individual mice

      $checksum = ( $count_ref + $count_alt + $count_mis );
      if ( $checksum != 40 ) {
	print("ERROR - number of alleles $checksum at pos $bppos is not equal to 40\n");
	print("$bppos : allele counts (ref, alt, mis) = $count_ref, $count_alt, $count_mis\n");
      }

      $data{ $bppos }{ 'allele counts' } = [ $count_ref, $count_alt, $count_mis ];
      $den = ( $checksum - $count_mis );
      $freq_ref = $count_ref/$den;
      $freq_alt = $count_alt/$den;
      $data{ $bppos }{ 'allele freqs' } = [ $freq_ref, $freq_alt ];

      if ( $freq_ref < $freq_alt ) {
	$data{ $bppos }{ 'MAF' } = $freq_ref;
	$data{ $bppos }{ 'MAF allele' } = 'REF';
      } else {
	$data{ $bppos }{ 'MAF' } = $freq_alt;
	$data{ $bppos }{ 'MAF allele' } = 'ALT';
      }
      #	print("MAF = ${$data{ $bppos }}{'MAF'}  MAF allele = ${$data{ $bppos }}{'MAF allele'}\n");

      $checksum = ( $count_hom_ref + $count_het + $count_hom_alt + $count_geno_mis );
      if ( $checksum != 20 ) {
	print("ERROR - number of genotypes $checksum at pos $bppos is not equal to 20\n");
	print("$bppos : genotype counts (ref, het, alt, mis) = $count_hom_ref $count_het $count_hom_alt $count_geno_mis\n");
      }

      $data{ $bppos }{ 'genotype counts' } = [ $count_hom_ref, $count_het, $count_hom_alt, $count_geno_mis ];
      $den = ( $checksum - $count_geno_mis );
      $freq_hom_ref = $count_hom_ref/$den;
      $freq_het = $count_het/$den;
      $freq_hom_alt = $count_hom_alt/$den;
      $data{ $bppos }{ 'genotype freqs' } = [ $freq_hom_ref, $freq_het, $freq_hom_alt ];

      $prop_miss_all = $count_mis/40;
      $prop_miss_geno = $count_geno_mis/20;
      $data{ $bppos }{ 'missing' } = [ $prop_miss_all, $prop_miss_geno ];

      #	print("missing data: alleles ${$data{ $bppos }}{'missing'}[0]  genotypes ${$data{ $bppos }}{'missing'}[1]\n");

      print(OUTN "$chr\t$bppos\t${$data{$bppos}}{'REF'}\t${$data{$bppos}}{'ALT'}\t${$data{$bppos}}{'TS'}\t${$data{$bppos}}{'MAF'}\t${$data{$bppos}}{'MAF allele'}\t${$data{$bppos}}{'missing'}[1]\t${$data{$bppos}}{'missing'}[0]\t${$data{$bppos}}{'allele freqs'}[0]\t${$data{$bppos}}{'allele freqs'}[1]\t${$data{$bppos}}{'genotype freqs'}[0]\t${$data{$bppos}}{'genotype freqs'}[1]\t${$data{$bppos}}{'genotype freqs'}[2]\t${$data{$bppos}}{'allele counts'}[0]\t${$data{$bppos}}{'allele counts'}[1]\t${$data{$bppos}}{'genotype counts'}[0]\t${$data{$bppos}}{'genotype counts'}[1]\t${$data{$bppos}}{'genotype counts'}[2]\n");

    } # end of mult ALTs loop

    } # end of while loop
    close(IN);
    close(OUTN);

    print("new stats file written\n");


    my @srted_keys = sort { $a <=> $b } keys %data;

    open(OUTA, '>', $Afile) or die "Cannot open file $Afile for writing, reason $!\n";
    print(OUTA "CHR\tBP\tREF\tALT\tNEW\tTS\tMAF\tMAF_allele\tProp_missing_genotypes\tProp_missing_alleles\tAllfreq_ref\tAllfreq_alt\tGenofreq_ref\tGenofreq_het\tGenofreq_alt\tAllcount_ref\tAllcount_alt\tGenocount_ref\tGenocount_het\tGenocount_alt\n");

    foreach my $loc ( @srted_keys ) {
	print(OUTA "$chr\t$loc\t${$data{$loc}}{'REF'}\t${$data{$loc}}{'ALT'}\t${$data{$loc}}{'new'}\t${$data{$loc}}{'TS'}\t${$data{$loc}}{'MAF'}\t${$data{$loc}}{'MAF allele'}\t${$data{$loc}}{'missing'}[1]\t${$data{$loc}}{'missing'}[0]\t${$data{$loc}}{'allele freqs'}[0]\t${$data{$loc}}{'allele freqs'}[1]\t${$data{$loc}}{'genotype freqs'}[0]\t${$data{$loc}}{'genotype freqs'}[1]\t${$data{$loc}}{'genotype freqs'}[2]\t${$data{$loc}}{'allele counts'}[0]\t${$data{$loc}}{'allele counts'}[1]\t${$data{$loc}}{'genotype counts'}[0]\t${$data{$loc}}{'genotype counts'}[1]\t${$data{$loc}}{'genotype counts'}[2]\n");
    }
    close(OUTA);

    print("combined stats file written\n");

} # end of chr loop
