#!/usr/bin/perl

use strict;


my $lookup = { "120904_sample_ID.csv"=>"National Institute of Genetics-Koide Mouse 02oct2012_FinalReport.txt",
	    "130108_sample_ID.csv"=>"National Institute of Genetics-Koide Mouse 07feb2013_FinalReport.txt",
	    "130305_sample_ID.csv"=>"Koide Mouse 01apr2013_FinalReport.txt",
	    "131118_sample_ID.csv"=>"Koide Mouse 09dec2013_FinalReport.txt" };

my $mega=undef;

foreach my $idfile ( keys %$lookup ) {
  print "$lookup->{$idfile}, $idfile, $mega \n";
  $mega = readRawMEGAMUGA( $lookup->{$idfile}, $idfile, $mega ) ;
}

my $founders = ["MSM", "HMI", "BLG2", "PGN2", "KJR", "CHD", "NJL", "BFM2"];


#$mega = mergeGenotypes( $mega, $raw );
#writeMergedData( $mega, $dir, $prefix);

sub readRawMEGAMUGA {
  my ($file, $idfile, $mega) = @_;
  my $reading = 0;
  my $thresh=0.7;
  my $snpthresh=0.4;
  my $snpc;
  my $samplec;
  my $snpgc;
  my $samplegc;
  my $geno;
  my $ngeno;
  my $newgeno;
  my $lines=0;
  my $samp_id;

  if ( $mega ) {
    $newgeno = $mega->{geno};
    $is_sample = $mega->{is_sample};
    $samplegc = $mega->{samplegc};
    $snpgc = $mega->{snpgc};
    $is_snp = $mega->{is_snp};
    $id = $mega->{id};
  }

  open (IDS, $idfile ) || die "Could not open $idfile\n";
  while(<IDS>) {
    chomp;
    my ( $num, $id ) = split(/\,/);
    $sampid->{$num} = $id;
  }
  close(IDS);

  open(FILE, $file ) || die "Could not open $file\n";
  print "reading $file\n";
  while( <FILE>) {
    chomp;
    if ( /\S+/ ) {
      print "$lines $reading $_\n" if ( $lines++ < 20 );
      if ( $reading )  {
	my ($snp, $samp, $a1, $a2, $x, $y, $g ) = split;
	my $sample = $samp_id->{$samp};
	die " unknown id $samp\n" if ( !$sample );
	$snpc->{$snp}++;
	$samplec->{$sample}++;
	if ( $g > $thresh ) {
	  $geno->{$sample}{$snp} = [" $a1 $a2", $g];
	  my $ggg = $a1 eq $a2 ? $a1 : "H";
	  $ggg = "N" if ( $a1 eq "-" );
	  $ggg = "N" if ( ! $gg->{$a1} );
	  $ngeno->{$sample}{$snp} = $ggg;
	  $snpgc->{$snp}++;
	  $samplegc->{$sample}++;
	}
      }
      else {
	if ( /^\[Data\]$/ ) {
	  $_ = <FILE>;
	  $reading =1;
	}
      }
    }
  }
  my $bad=0;

  foreach my $snp ( keys %$snpc ) {
    $snpgc->{$snp} /= $snpc->{$snp};
    if ( $snpgc->{$snp} < $snpthresh ) {
      $bad ++;
    } else {
      $is_snp->{$snp} = 1;
    }
  }
  print "$bad bad snps\n";
  my $nd=0;
  foreach my $sample ( sort keys %$samplec ) { 

    $samplegc->{$sample} /= ($samplec->{$sample}+1.0e-10);
    print "$sample $samplegc->{$sample}\n";

    $samplegc->{$sample} = $samplec->{$sample} = 0;
    foreach my  $snp ( keys %{$geno->{$sample}} ) {
      if ( $snpgc->{$snp} < $snpthresh) {
	delete( $geno->{$sample}{$snp} );
	$nd++;
      }
      else {
	$samplec->{$sample}++;
	$samplegc->{$sample} += ($geno->{$sample}{$snp}[1] > $thresh) ;
      }
    }
    $samplegc->{$sample} /= ($samplec->{$sample}+1.0e-10);
    print "$sample $samplegc->{$sample}\n";

  }
  print "deleted $nd snps\n";
  my $is_sample;

  foreach my $sample ( keys %$samplegc ) {
    my ( $n, $t ) = $sample =~ /IL\-(\d+)([AB]*)\-\d+/;
    $n += 500 if ( $t eq "B" );
    $id->{$sample} = "IL.$n";
    $newgeno->{$id->{$sample}} = $ngeno->{$sample};
    $is_sample->{$id->{$sample}} = $id->{$sample} if ( $n>0);
    print "sample $id->{$sample} $sample $ngeno->{$id->{$sample}}\n";    
  }
  
  return( { geno=>$newgeno, file=>$file, is_sample=>$is_sample, samplegc=>$samplegc, snpgc=>$snpgc, is_snp=>$is_snp, id=>$id });
}

sub readMEGAMUGAchrs {
  my ( $prefix ) = @_;

  my $geno;
  my $fgeno;
  my $map;
  my $is_snp;
  my $is_sample;
  my $is_founder;
  my $snps=0;
  my $totalsnps=0;
  my $nfounders = 8;
  
  my @files = glob ("/data/www/preCC/MEGA_MUGA/MegaMUGAsnps.Flines.TAUlines.chr*.csv");
  print "@files\n";
  foreach my $file ( @files ) {
    open(FILE, $file ) || die "Could not open $file\n";
    print "reading $file\n";
    $_ = <FILE>;
    chomp;
    my ( $snp, $chr, $bp, @header) = split(/,/);
    my @founder;
    my @cc;
    my $k;
    my $sex;
    
    for( $k=0;$k<$nfounders;$k++) {
      my ($prefix,$suffix) = $header[$k]=~ /^(.+)(\(.+\))$/;
#      $is_founder->{$header[$k]} = $prefix;
      print "prefix $prefix suffix $suffix\n";
      $is_founder->{$prefix} = $header[$k];
      $founder[$k] =$prefix;
    }
    for( my $i=$k;$i<=$#header;$i++) {
      $cc[$i-$k] = $header[$i];
      my ($prefix,$suffix) = $header[$i]=~ /^(.+)(\(.+\))$/;
      if ( $prefix =~/IL(\d+)([mf])/ ) {
	$prefix = "IL.$1";
	$sex = $2;
	$is_sample->{$header[$i]} =$prefix;
      }
    }
    
    my $skipped=0;
    
    while(<FILE>) {
      chomp;
      my ( $snp, $chr, $bp, @a ) = split(/,/);
      
      my $het = 0;
      my $miss = 0;
      my $na=0;
      my $H = "";
      my %al;
      for( my $i=0;$i<$nfounders;$i++) {
	if ( $a[$i] eq "N" ) {
	  $miss ++;
	} elsif ( $a[$i] eq "H" ) {
	  $het ++;
	}
	else {
	  $al{$a[$i]}++;
	}
      }
      foreach my $a ( sort keys %al ) {
	$na ++;
	$H .= " $a";
      }
      
      if ( $skipped < 10 ) {
#	print "$nfounders miss $miss het $het | ". join("-",@a) . " | na $na | $H |" . (keys %al) . "\n";
      }
      if ( $miss + $het == 0 && $na==2 ) {
	push @{$map->{$chr}}, [$bp, $snp, $H];
	$snps++;
#	print "$snps\n" if ( ! ($snps % 100 ));
	$is_snp->{$snp}++;
	my $i;
	for( $i=0;$i<$nfounders;$i++) {
	  $fgeno->{$founder[$i]}{$snp} = $a[$i];
	}
	for( my $j=$i;$j<=$#a;$j++) {
	  $geno->{$header[$j]}{$snp} = $a[$j];
	}
      } else {
	$skipped++;
      }
    }
    print "$file $snps $skipped\n";
    $totalsnps += $snps;
  }
  print "total snps $snps\n";
 
  my @s = keys %$is_sample;
  print "@s\n";
  my $samples = 1+$#s;
  @s = values %$is_sample;  print "@s\n";
  print "total samples $samples " . join( ",", sort keys %$is_sample ) . "\n";
  return( {file=>\@files, map=>$map, fgeno=>$fgeno, geno=>$geno, is_snp=>$is_snp, is_sample=>$is_sample, is_founder=>$is_founder} );
}
  
sub readMEGAMUGA_more {
  my ( $file, $data ) = @_;
  my $geno = $data->{geno};
  my $fgeno = $data->{fgeno};
  my $map = $data->{map};
  my $is_snp = $data->{is_snp};
  my $is_sample = $data->{is_sample};
  my $is_founder = $data->{is_founder};
  my $snps=0;
  my $nfounders = 0;

  open(FILE, $file ) || die "Could not open $file\n";
  print "reading $file\n";
  $_ = <FILE>;
  chomp;
  my ( $marker, $chr, $bp, @header) = split(/,/);
  my @cc;
  my $k;
  my $sex;
 
  for( my $i=0;$i<=$#header;$i++) {
    $cc[$i] = $header[$i];
    my ($prefix,$suffix) = $header[$i]=~ /^(.+)(\(.+\))$/;
    if ( $prefix =~/IL(\d+)([mf])/ ) {
      $prefix = "IL.$1";
      $sex = $2;
    }
    $is_sample->{$header[$i]} =$prefix;
  }

  my $skipped=0;
  while(<FILE>) {
    chomp;
    my ( $snp, $chr, $bp, @a ) = split(/,/);
    if ( $is_snp->{$snp} ) {
      for( my $j=0;$j<=$#a;$j++) {
	$geno->{$header[$j]}{$snp} = $a[$j];
      }
    }
  }
  my @s = keys %$is_sample;
  my @v = values %$is_sample;
  my $samples = $#s+1;
  print "more samples $samples @s @v\n";
  return( {file=>$file, map=>$map, fgeno=>$fgeno, geno=>$geno, is_snp=>$is_snp, is_sample=>$is_sample, is_founder=>$is_founder} );
}
  
sub mergeGenotypes { 

  my ( $set, @sets ) = @_;
  my $geno = $set->{geno};
  my $fgeno = $set->{fgeno};
  my $map = $set->{map};
  my $is_snp = $set->{is_snp};
  my $sample_count;
  my $merged;
  my $oksnp;
  my $d=0;
  my $new;
  my $nsnps=0;

  foreach my $sample ( keys %{$set->{is_sample}} ) {
    my $cc = $set->{is_sample}{$sample};
    if ( exists( $merged->{$cc} ) ) {
      warn "$sample $cc duplicated\n";
    }
    $merged->{$cc} = $geno->{$sample};
  }
  foreach my $s ( @sets ) {
    foreach my $sample ( keys %{$s->{is_sample}} ) {
      my $cc = $s->{is_sample}{$sample};
      if ( exists( $merged->{$cc} ) ) {
	warn "$sample duplicated\n";
      } else {
	$merged->{$cc} = $s->{geno}{$sample};
      }
    }
  }
    
  foreach  my $s ( @sets ) {
    foreach my $snp ( keys %{$set->{is_snp}} ) {
      $set->{is_snp}{$snp} += $s->{is_snp}{$snp};
    }
  }
  my $n = $#sets+2;
  my $k=0;
  foreach my $snp ( keys %{$set->{is_snp}} ) {
    print "$snp $set->{is_snp}{$snp} $n\n" if ( $k++<20);
    if ( $set->{is_snp}{$snp} == $n ) {
      $oksnp ++;
    }
    else {
      delete( $set->{is_snp}{$snp});
      $d++;
    }
  }

  print "$oksnp snps in common $d deleted\n";
  $nsnps += $oksnp;

  my @chr = sort keys %$map;
  print "merge chr @chr\n";
 foreach my $chr ( @chr ) {
   my @map = @{$map->{$chr}};
   my @m;
   foreach my $ref ( @map ) {
     push @m, $ref if ( $set->{is_snp}{$ref->[1]});
   }
   $new->{$chr} = \@m;
 }

  print "$nsnps total snps\n";
 return( { geno=>$merged, is_snp=>$is_snp, fgeno=>$fgeno, map=>$new });
}

sub writeMergedData {
  my ( $data, $dir, $prefix ) = @_;
  my $h;
  my @chr = sort keys  %{$data->{map}} ;
  my @samples = sort keys %{$data->{geno}} ;
  system( "mkdir $dir") if ( ! -d $dir );
  print "chr @chr\n";
  foreach my $chr ( @chr ) {
    my $pedfile = "$dir/chr$chr.$prefix.data";
    my $allelesfile = "$dir/chr$chr.$prefix.alleles";
    my $mapfile = "$dir/chr$chr.$prefix.map";
    print "$pedfile $allelesfile $mapfile\n";
    my @snps;
    foreach my  $ref ( @{$data->{map}{$chr}} ) {
      push @snps, $ref->[1];
      $h->{$ref->[1]} = $ref->[2];
    }
    open(PED, ">$pedfile") || die "Could not write $pedfile\n";
    foreach my $s ( @samples ) {
      print PED join( " ", "CC", $s, "0", "0", 1, "NA");
      my @g;
      foreach my $snp ( @snps ) {
	$gg->{H} = $h->{$snp};
	push @g, $gg->{$data->{geno}{$s}{$snp}};
      }
      print PED join ( " ", @g) . "\n";
    }
    close(PED);

    my $fgeno = $data->{fgeno};
    my @map = @{$data->{map}{$chr}};
    my @strains = sort keys %$fgeno;
    my $nmarkers = $#snps +1;
    open(ALLELES, ">$allelesfile") || die "Could not open $allelesfile\n";
    print ALLELES "markers $nmarkers strains 8\n";
    print ALLELES "strain_names @strains\n";
    for( my $i=0;$i<=$#snps; $i++) {
      my $snp = $snps[$i];
      my $cm = $map[$i]->[0]/2.0e6;
      my %al;
      my $g;
      foreach my $s ( @strains ) {
	$al{$fgeno->{$s}{$snp}}++;
	$g->{$fgeno->{$s}{$snp}}{$s} = 1;
      }
      my @a = sort keys %$g;
      my $na = $#a +2;
      print ALLELES "marker $snp $na $chr $cm\n";
      foreach my $a ( @a ) {
	print ALLELES "allele $a ";
	foreach my $s ( @strains ) {
	  my $p = $g->{$a}{$s}/$al{$a};
	  printf ALLELES " %.4f", $p;
	}
	print ALLELES  "\n"
      }
      print ALLELES "allele NA ";
      foreach my $s ( @strains ) {
	printf ALLELES " %.4f", 1/8;
      }
      print ALLELES  "\n";
    }
    close(ALLELES);
  
    open(MAP, ">$mapfile") || die "Could not write $mapfile\n";
    print MAP "marker	chromosome	bp\n";
    
    foreach my  $ref ( @{$data->{map}{$chr}} ) {
      print MAP "$ref->[1]\t$chr\t$ref->[0]\n";
    }
    close(MAP);
  }
}

  
