#!/usr/bin/perl -w

#
# tasuke_variant_multivcf.pl
# http://tasuke.dna.affrc.go.jp/
#
# Author: Norio Tabei
#
# Copyright (c) 2021 National Institute of Agrobiological Sciences
#

use strict;
use warnings;
use DBI;
use File::Basename 'basename', 'dirname';
use Getopt::Long;
use FindBin qw($Bin);
use lib "$Bin";
use POSIX 'WNOHANG';
use Data::Dumper;

my $RESTORE= 0;
my $DIR    = '';
my $DB     = '';
my $HOST   = 'localhost';
my $USER   = '';
my $PASS   = '';
my $THREADS = 4;
my $ACCMAPFILE = '';
my $CHECKONLY = 0;
my $MEASUREDB = 0;

################################################################
# Checking parameters and database connection
################################################################
GetOptions(     'r|restore'	=> \$RESTORE,
		'd|dir=s'	=> \$DIR,
		'db|database=s'	=> \$DB,
		'h|host=s'	=> \$HOST,
		'u|user=s'	=> \$USER,
		'p|password=s'	=> \$PASS,
		't|threads=i'	=> \$THREADS,
		'm|accmap=s'	=> \$ACCMAPFILE,
		'k|checkonly'	=> \$CHECKONLY,
		's|measure'	=> \$MEASUREDB
);

my $connect;
if ( !$DB || !$USER ) {
	_error_method();
	_exit(1);
}

if(!$DIR && !$ACCMAPFILE) {
	print STDERR "\"-d\" or \"-m\" (or both) is required.\n";
	_error_method();
	_exit(1);
} elsif($DIR && !(-d $DIR)) {
	print STDERR "$DIR not found.\n";
	_error_method();
	_exit(1);
} elsif($ACCMAPFILE && !(-f $ACCMAPFILE)) {
	print STDERR "$ACCMAPFILE not found.\n";
	_error_method();
	_exit(1);
}

$connect= DBI->connect('DBI:mysql:database='. $DB. ';host='. $HOST, $USER, $PASS, {AutoCommit => 0})
	or die "Can't connect to database: $DBI::errstr\n";

my @pids;
sub sigint_handler {
	print STDERR "\nGet SIGINT. All threads have been killed.\nDepth registered halfway can be deleted with the \"-r\" option.\n";
	_exit(1);
}
sub sigkill_handler {
	print STDERR "\nGet SIGKILL. All threads have been killed.\nDepth registered halfway can be deleted with the \"-r\" option.\n";
	killAllChildProcess(\@pids);
	_exit(1);
}
$SIG{'INT'} = 'sigint_handler';
$SIG{'KILL'} = 'sigkill_handler';
$SIG{'TERM'} = 'sigkill_handler';

###########################################################
# Checking base table
###########################################################
if ( ! defined _check_exists_table("sys_accession") ) {
    print "Some basic table does not exists.\n";
    print "Did you execute 'tasuke_init.pl' at the first?\n";
    _exit(1);
}

my ($dbAccInfoList, $dbAcc2InfoMap) = getDbAccInfo();
#print Dumper($dbAccInfoList);

#test_getAccMap($dbAccInfoList, $vcfSamples);
#test_getAccCsv($dbAccInfoList);

my $oriSample2accMap = {};
if($ACCMAPFILE) {
	setMapByFile($oriSample2accMap, $ACCMAPFILE);
} else {
	# FileName = AccessionID
	my @oriInfiles = glob("$DIR/*");
	foreach my $infile (@oriInfiles) {
		next if( !(-f $infile) );
		$infile =~ s/^$DIR//;
		$infile =~ s/^\///;

		my $acc = basename($infile);
		$acc =~ s/\.(gz|bz2)$//;
		$acc =~ s/\.tar$//;
		$acc =~ s/\.tsv$//i;
		$oriSample2accMap->{$infile} = $acc;
	}
}
my @infilesArray = sort {$a cmp $b} keys(%{$oriSample2accMap});
foreach my $infile (@infilesArray) {
	my $fullPath = sprintf('%s/%s', $DIR, $infile);
	if( !(-f $fullPath) ) {
		print STDERR "$fullPath not found\n";
		_exit(1);
	}
}

my @existSamples;
my @skipSamples;
my ($sample2accMap, $nonExistSamples) = deleteNonExistSamples($oriSample2accMap, $dbAcc2InfoMap);
{
	foreach my $sample (@infilesArray) {
		my $acc = "";
		if(exists($sample2accMap->{$sample})) {
			push @existSamples, $sample;
		} elsif(!exists($oriSample2accMap->{$sample})) {
			push @skipSamples, $sample;
		}
	}
}

my $existSamplesNum = @existSamples;
{
	if($RESTORE) {
		printf STDERR "> The following AccessionIDs is the target for deleting variant record: %d\n", $existSamplesNum;
	} else {
		printf STDERR "> The following AccessionIDs is the target for DB registration: %d\n", $existSamplesNum;
	}
	print STDERR "   +-----------------------------------------+\n";
	print STDERR "   |      AccessionID     <- TSV File  |\n";
	print STDERR "   +-----------------------------------------+\n";
	my $i = 1;
	foreach my $sample (@existSamples) {
		my $acc = $sample2accMap->{$sample};
		printf STDERR "   \| %-4d %-15s <- %-15s \| \n", $i, $acc, $sample;
		$i++;
	}
	print STDERR "   +-----------------------------------------+\n\n";
}

my $infiles = \@infilesArray;
my $notExistNum = @{$nonExistSamples};
if($notExistNum > 0) {
	printf STDERR "> !!! The AccessionIDs in the following samples are invalid: %d\n", $notExistNum;
	printf STDERR "> These will be skipped.\n";
	my $sorted = sortSampleVcfOrder($infiles, $nonExistSamples);
	my $count = 1;
	foreach my $nonExistSample (@{$sorted}) {
		my $acc = $oriSample2accMap->{$nonExistSample};
		printf STDERR "   \| %-4d %-15s <- %-15s \| \n", $count, $acc, $nonExistSample;
		$count++;
	}
	print STDERR "\n";
} else {
	printf STDERR "> A valid AccessionID was found in all the specified samples.\n\n";
}

my $skipNum = @skipSamples;
if($skipNum > 0) {
	printf STDERR "> The following VCF samples will be skipped (because it is not specified by user): %d\n", $skipNum;
	foreach my $sample (@skipSamples) {
		printf STDERR "$sample\n";
	}
	print STDERR "\n";
} 

my $dbAccsLen = @{$dbAccInfoList};
my $infilesLen =  @{$infiles};
my $SpecifiedIdsLen = keys(%{$oriSample2accMap});
my $statStr = sprintf "> Statistics\n> AllAccsInDB:%d   FilesInDir:%d   In it,SpecifiedSamples:%d   AccIdFound:%d   notFound(skip):%d",
        $dbAccsLen, $infilesLen, $SpecifiedIdsLen, $existSamplesNum, $notExistNum;

if($RESTORE) {
	print STDERR "$statStr\n\n";
	if($CHECKONLY) {
		print STDERR "> [This is CHECK-ONLY mode] it exits without deleting.\n";
	} else {
		print STDERR "> Delete only the depth record. AccessionID itself is \"not\" deleted.\n";
		print STDERR "> Are you sure deleteing?[y|n] ";
		while (defined(my $line = <STDIN>)){
			chomp($line);
			if ( $line eq "y" || $line eq "Y" ) { last; }
			elsif ( $line eq "n" || $line eq "N" ) { print STDERR "Canceled.\n"; _exit(0); }
			else { print STDERR "> [y|n] "; }
		}
		truncateTables($sample2accMap, $dbAcc2InfoMap);
	}
	_exit(0);
}

# Create command.
my $exOpt = "";
$exOpt .= " -p $PASS" if($PASS);
my $cmdBase = sprintf "perl %s/%s -h %s -db %s -u %s %s", $Bin, 'tasuke_tsv_db.pl', $HOST, $DB, $USER, $exOpt;

my $finalCode = 0;
if($CHECKONLY) {
	if($MEASUREDB) {
		measureDb($dbAccInfoList, $sample2accMap);
	}
	print STDERR "\n$statStr\n";
	print STDERR "\n> [This is CHECK-ONLY mode] Finished without DB registration.\n";
} else {
	my $endPid;
	my $remainSamplesLen = @existSamples;
	my $logCount = 0;
	my @logQueue;
	my $logPat = $DB. "_tsv_thread%03d.log";
	my %pid2log;
	do {
		while(1) {
			my $ts = @pids;
			if($ts < $THREADS) {
				my $infile = shift @existSamples;
				if($infile) {
					my $accId = $sample2accMap->{$infile};
					my $fullPath = sprintf "%s/%s", $DIR, $infile;
					$fullPath =~ s/^\/\//\//;
#					my $logFile = sprintf "%s_tsv_%04d_%s.log", $DB, $logCount, $accId;
					my $logFile = shift @logQueue;
					if(!$logFile) {
						# Generate new log.
						$logCount++;
						$logFile = sprintf $logPat, $logCount;
						unlink($logFile) if(-f $logFile);
					}
					my $cmd = sprintf "%s -n %s -f %s >>%s 2>&1", $cmdBase, $accId, $fullPath, $logFile;
					my $pid = fork;
					if ($pid == -1) {
						printf STDERR "Failed to generate Thread. [%s <- %s]\n", $accId, $infile;
						killAllChildProcess(\@pids);
						_exit(1);
					} elsif ($pid == 0) {
						# Child process.
						exec("exec $cmd");
					} else {
						# Parent process.
						printf STDERR "PID:%d\t%s\n", $pid, $cmd;
						$pid2log{$pid} = $logFile;
						push @pids, $pid;
					}
				} else {
					last;
				}
			} else {
				last;
			}
		}
		$remainSamplesLen = @existSamples;

		$endPid = waitpid(-1, WNOHANG);
		if($endPid != 0) {
			my $exit_value = $? >> 8;
			if(exists($pid2log{$endPid})) {
				print STDERR "PID:$endPid is ended and get status $exit_value";
				if($exit_value != 0) {
					printf STDERR "\t[CAUTION] status is invalid. Please check log \"%s\"", $pid2log{$endPid};
					$finalCode = $exit_value;
				}
				print "\n";
			}
			my @newPids;
			foreach my $pid (@pids) {
				push(@newPids, $pid) if($pid != $endPid);
			}
			@pids = @newPids;

			push(@logQueue, $pid2log{$endPid});
		}
		sleep 1;
	} while ($remainSamplesLen != 0 || $endPid != -1) ;
}

_exit($finalCode);

sub getFileHandle {
	my($VCF) = @_;
	my $fh;
	if ($VCF =~ m/\.tar\.gz$/) {
		open $fh, "tar zxf $VCF -O 2> /dev/null|" or die "file open error: $!";
	} elsif($VCF =~ m/\.tar\.bz2$/) {
		open $fh, "tar jxf $VCF -O 2> /dev/null|" or die "file open error: $!";
	} elsif($VCF =~ m/\.gz$/) {
		open $fh, "gzip -dc $VCF 2> /dev/null|" or die "file open error: $!";
	} elsif($VCF =~ m/\.bz2$/) {
		open $fh, "bzip2 -dc $VCF 2> /dev/null|" or die "file open error: $!";
	} else {
		open $fh, '<', $VCF or die "file open error: $!";
	}
	return $fh;
}
sub sortSampleVcfOrder {
	my ($vcfOrder, $samples) = @_;

	my $vcfOrderLen = @{$vcfOrder};
	my %sample2index;
	for(my $i=0; $i<$vcfOrderLen; $i++) {
		$sample2index{$vcfOrder->[$i]} = $i;
	}
	my @sorted =  sort {$sample2index{$a} <=> $sample2index{$b}} @{$samples};
	return \@sorted;
}
sub getVcfSamplelist {
	my($vcf) = @_;

	my @samples;
	local *FH;
#	open(FH, "<$vcf")  or die "$vcf: $!";
#	while(<FH>) {
	my $fh = getFileHandle($vcf);
	while(<$fh>) {
		chomp;
		my $datarow = trim($_);
		next if($datarow !~ m/\S/);

		if($datarow =~ m/#CHROM/i) {
			my @header = split /\t/, $datarow;
			my $headerLen = @header;
			for(my $i=9; $i<$headerLen; $i++) {
				push @samples, $header[$i];
			}
			last;
		}
	}
	close($fh);
#	close(FH);
	return \@samples;
}
sub getDbAccInfo {
	my $sth = $connect->prepare('select * from sys_accession');
	$sth->execute();
	my @array;
	my %hash;
	while(my $hashref = $sth->fetchrow_hashref) {
		push @array, $hashref;
		$hash{$hashref->{'accession_id'}} = $hashref;
	}
	$sth->finish();
	return (\@array, \%hash);
}
sub setMapByFile {
	my ($map, $accMapFile) = @_;

	local *FH;
	open(FH, "<$accMapFile") or die "$accMapFile: $!";
	while(<FH>) {
		chomp;
		my $datarow = trim($_);
		next if($datarow !~ m/\S/);

		my @vals = split /,/, $datarow;
		my $valsLen = @vals;
		for(my $i=0; $i<$valsLen; $i++) {
			$vals[$i] = trim($vals[$i]);
		}
#		my @vals = split /\s+/, $datarow;
		if($vals[0]) {
			if($vals[1]) {
				$map->{$vals[0]} = $vals[1];
			} else {
				$map->{$vals[0]} = $vals[0] 
			}
		}
	}
	close(FH);
}
sub deleteNonExistSamples{
	my ($map, $acc2Info) = @_;

	my %existsMap;
	my @nonExists;
	my @samples = keys(%{$map});
	foreach my $sample (@samples) {
		my $acc = $map->{$sample};
		if(!$acc || !exists($acc2Info->{$acc})) {
#			delete $map->{$sample};
			push @nonExists, $sample;
		} else {
			$acc2Info->{$acc}->{'sample'} = $sample;
			$existsMap{$sample} = $map->{$sample};
		}		
	}
	return (\%existsMap, \@nonExists);
}
sub truncateTables {
	my ($map, $acc2Info) = @_;

	my @samples = keys(%{$map});
	foreach my $sample (@samples) {
		my $acc = $map->{$sample};
		my $exOpt = "";
		$exOpt .= " -p $PASS" if($PASS);
		my $cmd = sprintf "perl %s/%s -h %s -db %s -u %s -n %s -r %s", $Bin, 'tasuke_tsv_db.pl', $HOST, $DB, $USER, $acc, $exOpt;
		my $status = 'success';
		$status = "failed" if(system($cmd) != 0);
=head
		my $db = $acc2Info->{$acc}->{'accession_db'};
		my $sql = sprintf "truncate table %s_accession", $db;
		my $status = 'success';
		my $sth = $connect->prepare($sql);
		$sth->execute() or $status = "failed";
=cut
		printf STDERR "   | %-15s <- %-15s|;	=> %s\n", $acc, $sample, $status;
	}
}
sub test_getAccMap {
	my ($dbAccInfos, $vcfSamples) = @_;

	foreach my $vcfSample (@{$vcfSamples}) {
		my $vcfAcc = "";
		foreach my $dbAccInfo (@{$dbAccInfos}) {
			foreach my $col ('accession_id', 'accession_name', 'origin', 'origin_2', 'other_1', 'other_2') {
				if($vcfSample eq $dbAccInfo->{$col}) {
					$vcfAcc = $dbAccInfo->{'accession_id'};
					last;
				}
			}
		}
		if(!$vcfAcc) {
			foreach my $dbAccInfo (@{$dbAccInfos}) {
				my $accName = $dbAccInfo->{'accession_name'};
				$accName =~ s/\s//g;
				if($vcfSample =~ m/$accName/i) {
					$vcfAcc = $dbAccInfo->{'accession_id'};
					last;
				}

			}
		}
		printf "%s\t%s\n", $vcfSample, $vcfAcc;
	}
}
sub test_getAccCsv {
	my ($dbAccInfos) = @_;

	my @cols = ('accession_id', 'accession_name', 'variety', 'sub_variety', 'origin', 'origin_2', 'type', 'other_1', 'other_2');
	foreach my $info (@{$dbAccInfos}) {
		my @onefow;
		foreach my $col (@cols) {
			push @onefow, $info->{$col};
		}
		printf "%s\n", join(',', @onefow);
	}
}
sub _check_exists_table {
	my ($t) = @_;
	my $val;
	$val = $connect->selectrow_array("show tables like '$t'");
	return $val;
}
sub trim {
	my($str) = @_;
	$str =~ s/^\s+//;
#	$str =~ s/[\s\W]+$//;
	$str =~ s/[\s]+$//;
	return $str;
}
sub killAllChildProcess {
	my($pids) = @_;
	foreach my $pid (@{$pids}) {
		kill(9, $pid);
#		kill(2, $pid);
	}
}
sub measureDb {
	my($accList, $sample2accMap) = @_;
	print STDERR "> Record count by accessionID\n\n";
	foreach my $accRef (@{$accList}) {
		my $sql = sprintf "select count(1) as cnt from %s_accession", $accRef->{'accession_db'};
		my $sth = $connect->prepare($sql);
		$sth->execute();
		my $arrayref = $sth->fetchrow_arrayref();
		my $sampleName = $accRef->{'sample'};
		my $targetOrNot = "";
		$targetOrNot = '*' if(exists($sample2accMap->{$sampleName}));
		printf STDERR "%s %s  <-  %s\t%s_accession\t%s\n", $targetOrNot, $accRef->{'accession_id'}, $sampleName, $accRef->{'accession_db'}, $arrayref->[0];
		$sth->finish();
	}
}

#### Disconnection and exit #######################################
sub _exit {
    my ($arg) = @_;
    my $code = 0;
    if(defined $arg){$code = $arg;}
    $connect->disconnect or warn "Disconnection error: $DBI::errstr\n" if($connect);
    exit($code);
}

sub _error_method {
	print STDERR "Depth registrations(by tasuke_tsv_db.pl) can be parallelized.\n";
	print STDERR "Automatically assign AccessionIDs to each threads.\n\n";
	print STDERR "Usage: $0  -db [database] -u [user] -p [password]\n\n";
	print STDERR "-d <DirPath> : Directory path where the TSV files are located. By default, filename(without extensions) is considered as AccessionID.\n";
	print STDERR "-m <Path>: Path to \"TsvFilePath > AccID\" correspondence table. If \"-d\" is specified, TsvFilePath is its relative path. TsvFilePath[,]AccID[\\n]...\n";
	print STDERR "Other options:\n";
	print STDERR "  -h  <remote host> : (default: localhost)\n";
	print STDERR "  -t <num> : Number of threads(default: 4)\n";
	print STDERR "  -k : TEST mode. Check AccessionIDs, but do not perform DB registration.\n";
	print STDERR "Sub-action:\n";
	print STDERR "  -r : Delete depth information for the specified AccessionIDs. AccessionID itself is not deleted.\n";
	print STDERR "\n";
}
