#!/usr/bin/perl -w

# legacy_tasuke_variant_vcf_multi.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 Getopt::Long;
use FindBin qw($Bin);
use lib "$Bin";
use POSIX 'WNOHANG';
use CommonLibTASUKE;
use Data::Dumper;

my $NAME   = '';
my $RESTORE= 0;
my $REGGT00 = 0;
my $VCF    = '';
my $DB     = '';
my $HOST   = 'localhost';
my $USER   = '';
my $PASS   = '';
my $THREADS = 4;
my $ACCMAPFILE = '';
my $CHECKONLY = 0;
my $MEASUREDB = 0;
my $FORCEFLG = 0;
my $APPENDFLG = 0;

my $commonMysql = new CommonMysql();

################################################################
# Checking parameters and database connection
################################################################
GetOptions(	'n|name=s'	=> \$NAME,
		'r|restore'	=> \$RESTORE,
		'z|reggt00'	=> \$REGGT00,
		'f|file=s'	=> \$VCF,
		'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,
		'c|force'	=> \$FORCEFLG,
		'a|append',	=> \$APPENDFLG
);

my $connect;
if ( !$DB || !$USER ) {
	_error_method();
	_exit(1);
	
}
if($VCF) {
	if(!(-f $VCF)) {
		print STDERR "VCF file not found: $VCF\n";
		_exit(1);
	}
} elsif(!$RESTORE) {
	print STDERR "VCF file(-f) is required.\n";
	_exit(1);
	#TODO VCF is not required for RESTORE mode.
}

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

my $MODE_IS_MULTIVCF = $commonMysql->modeIsPlusMultiVcf($connect);

my @pids;
sub sigint_handler {
	print STDERR "\nGet SIGINT. All threads have been killed.\nVariants registered halfway can be deleted with the \"-r\" option.\n";
	_exit(1);
}
sub sigkill_handler {
	print STDERR "\nGet SIGKILL. All threads have been killed.\nVariants 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 $vcfSamples = getVcfSamplelist($VCF);
#print Dumper($vcfSamples);

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

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

my $oriSample2accMap = {};
if($NAME) {
	setMapByNames($oriSample2accMap, $vcfSamples, $NAME)
} elsif($ACCMAPFILE) {
	setMapByFile($oriSample2accMap, $ACCMAPFILE);
} else {
	# SampleName = AccessionID
	foreach my $vcfSample (@{$vcfSamples}) {
		$oriSample2accMap->{$vcfSample} = $vcfSample;
	}
}

my @existSamples;
my @skipSamples;
my ($sample2accMap, $nonExistSamples) = deleteNonExistSamples($oriSample2accMap, $dbAcc2InfoMap);
{
	foreach my $sample (@{$vcfSamples}) {
		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     <- VCF SampleName  |\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 $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($vcfSamples, $nonExistSamples);
	my $count = 1;
	foreach my $nonExistSample (@{$sorted}) {
		my $acc = $oriSample2accMap->{$nonExistSample};
		printf STDERR "   \| %-4d %-15s <- %-15s \| \n", $count, $acc, $nonExistSample;
#		printf STDERR "$count\t$nonExistSample\n";
		$count++;
	}
	print STDERR "\n";
#	_exit(1);
} 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 $vcfSamplesLen =  @{$vcfSamples};
my $SpecifiedIdsLen = keys(%{$oriSample2accMap});
my $statStr = sprintf "> Statistics\n> AllAccsInDB:%d   VcfSamples:%d   In it,SpecifiedSamples:%d   AccIdFound:%d   notFound(skip):%d",
        $dbAccsLen, $vcfSamplesLen, $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 {
		if(!$FORCEFLG) { 
			print STDERR "> Delete only the Variant 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);
}

# Divide AccNum into Thread count.
my @divUnits;
my $unit = int($existSamplesNum / $THREADS);
my $divReminder = $unit > 0 ? ($existSamplesNum % $THREADS) : $existSamplesNum;
for(my $i=0; $i<$THREADS && $i<$existSamplesNum; $i++) {
	$divUnits[$i] = $unit;
	$divUnits[$i]++ if($divReminder > 0);
	$divReminder--;
}
#print STDERR Dumper(@divUnits);
my @divUnitsCopy = @{\@divUnits};	# for Log

# Generate "-n" option string for each tasuke_variant_vcf.pl threads.
#my $vcfSamplesLen =  @{$vcfSamples};
my $divUnitsLen = @divUnits;
my @divAccs;
# Initialize.
for(my $i=0; $i<$divUnitsLen; $i++) {
	$divAccs[$i] = [];
	for(my $j=0; $j<$vcfSamplesLen; $j++) {
		$divAccs[$i][$j] = "";
	}
}

my @fullAccs;
my $index = 0;
for(my $i=0; $i<$vcfSamplesLen && $index<$divUnitsLen; $i++) {
	my $vcfSample = $vcfSamples->[$i];
	if(exists($sample2accMap->{$vcfSample})) {
		$divAccs[$index][$i] = $sample2accMap->{$vcfSample};
		$divUnits[$index]--;
	}
	
	$index++ if($divUnits[$index] <= 0);

	$fullAccs[$i] = "";
	$fullAccs[$i] = $sample2accMap->{$vcfSample} if(exists($sample2accMap->{$vcfSample}));
}

# Create command.
my $exOpt = "";
$exOpt .= " -p $PASS" if($PASS);
$exOpt .= " -z" if($REGGT00);
$exOpt .= " -a" if($APPENDFLG);
my $cmdBase = sprintf "perl %s/%s -h %s -db %s -u %s -t %s -f %s -c %s", $Bin, 'tasuke_variant_vcf.pl', $HOST, $DB, $USER, 'gatkm', $VCF, $exOpt;
my $sumAccs = 1;
my @cmds;
for(my $i=0; $i<$divUnitsLen; $i++) {
	my $idsStr = join(',', @{$divAccs[$i]});
	$idsStr =~ s/,+$//;
	my $logFile = sprintf "%s_vcf_%03d-%03d.log", $DB, $sumAccs, ($sumAccs + $divUnitsCopy[$i] - 1);
	push @cmds, [
		sprintf("%s -n \"%s\"  >%s 2>&1", $cmdBase, $idsStr, $logFile),
		$logFile
	];
	$sumAccs += $divUnitsCopy[$i];
}

# Regist sys_variant info
if($MODE_IS_MULTIVCF) {
	my $logFile = sprintf "%s_vcf_sys_variant.log", $DB;
	push @cmds, [
#		sprintf("%s -m -n \"%s\"  >%s 2>&1", $cmdBase, $divAccs[0][0], $logFile),
		sprintf("%s -m >%s 2>&1", $cmdBase, $logFile),
		$logFile
	];
}

my $finalCode = 0;
if($CHECKONLY) {
#	print STDERR "> Correspondence of AccessionID from VCF SampleName order:\n"; 
#	printf STDERR "%s\n\n", join(",", @fullAccs);
	if($MEASUREDB) {
		measureDb($dbAccInfoList, $sample2accMap);
	}
	print STDERR "> [This is CHECK-ONLY mode] DB registration is not executed, only the command of each thread is output.\n\n";
	my $cmdsLen = @cmds;
	for(my $i=0; $i<$cmdsLen; $i++) {
		printf STDERR "# Thread%d:\n%s\n", ($i + 1), $cmds[$i][0];
	}
	print STDERR "\n$statStr\n";
	print STDERR "\n> [This is CHECK-ONLY mode] Finished without DB registration.\n";
} else {
	$| = 1;	# AutoFlush:On

	if(!$FORCEFLG) {
		print STDERR "> Checking if records exist in the target tables...\n";
		my @varTables;
		push @varTables, ['sys_variant','sys_variant'] if($MODE_IS_MULTIVCF);
		foreach my $sample (@{$vcfSamples}) {
			if(exists($sample2accMap->{$sample})) {
				my $accId = $sample2accMap->{$sample};
				my $varTable = $dbAcc2InfoMap->{$accId}->{'accession_db'}. '_variant';
				push @varTables, [$accId, $varTable];
			}
		}

		my @accsHasRecord;
		my $loop = 1;
		foreach my $varTableArr (@varTables) {
			printf STDERR "%d\t%s (%s)\t", $loop, $varTableArr->[0], $varTableArr->[1];
			if(hasRecord($varTableArr->[1])) {
				printf STDERR "RecordExists";
				push @accsHasRecord, $varTableArr->[0];
			} else {
				printf STDERR "Empty";
			}
			printf STDERR "\n";
			$loop++;
		}
		my $lenAccsHasRecord = @accsHasRecord;
		if($lenAccsHasRecord) {
			printf STDERR "> [!!!Please check!!!]\n%d tables below contain existing data.\n", $lenAccsHasRecord;
			print STDERR join(",", @accsHasRecord). "\n\n";
			if($APPENDFLG) {
				print STDERR "> This is \"Append\" mode, so it will add data to the existing table.\n";
				print STDERR "> [Note] Data with the same chr-pos will be registered multiple times.\n";
			} else {
				print STDERR "> Registration for the tables above will be skipped.\n";
			}
			print STDERR "> Are you sure proccessing?[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] "; }
			}
		} 
	}

	# If this isn't done, the threads below will stall.
	# Is the database locked? hasRecord() seems to be the problem, but the cause is unknown.
	$connect->disconnect or warn "Disconnection error: $DBI::errstr\n" if($connect);

	my %pid2log;
	foreach my $cmdRef (@cmds) { 
		my $cmd = $cmdRef->[0];
		my $logfile = $cmdRef->[1];

		my $pid = fork;
		if ($pid == -1) {
			killAllChildProcess(\@pids);
			print STDERR "Failed to generate Thread. The process ends.\n";
			_exit(1);
		} elsif ($pid == 0) {
			# Child process.
			exec("exec $cmd");	# This script ends here.
		} else {
			printf STDERR "PID:%d\t%s\n", $pid, $cmd;
			# Parent process.
			push @pids, $pid;
			$pid2log{$pid} = $logfile;
		}
	}

	# If the parent process ends first, the child process will become a zombie.
	print STDERR "\nMonitoring the end of threads. (If you want to stop, press Ctrl+C. All threads will stop. The same is true if you kill this process.)\n"; 
	my $killid;
	do {
		$killid = waitpid(-1, WNOHANG);
		my $exit_value = $? >> 8;
		if($killid != 0 && exists($pid2log{$killid})) {
			print STDERR "PID:$killid is ended and get status $exit_value";
			if($exit_value != 0) {
				printf STDERR "\t[CAUTION] status is invalid. Please check log \"%s\"", $pid2log{$killid};
				$finalCode = $exit_value;
			}
			print STDERR "\n";
		}
		sleep 1;
	} while ($killid != -1) ;

	$| = 0;		# AutoFlush:Off
}
_exit($finalCode);

sub hasRecord {
	my($table) = @_;
#	return $connect->selectrow_array("select * from $table limit 1");
	my $sth = $connect->prepare("select 1 from $table limit 1");
	$sth->execute() or die $!;
	my $res = $sth->fetchrow_arrayref();
	$sth->finish();
	return defined($res);
}
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 setMapByNames {
	my ($map, $vcfSamples, $namesStr) = @_;
	$namesStr =~ s/\s//g;
	my @names = split /,/, $namesStr;
	my $namesLen = @names;
	my $vcfSamplesLen = @{$vcfSamples};
	for(my $i=0; $i<$namesLen && $i<$vcfSamplesLen; $i++) {
		if($names[$i]) {
			$map->{$vcfSamples->[$i]} = $names[$i];
		}
	}
}
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 $db = $acc2Info->{$acc}->{'accession_db'};
		my $sql = sprintf "truncate table %s_variant", $db;
		my $status = 'success';
		my $sth = $connect->prepare($sql);
		$sth->execute() or $status = "failed";
		$sth->finish();
		printf STDERR "   | %-15s <- %-15s| %s;	=> %s\n", $acc, $sample, $sql, $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_variant", $accRef->{'accession_db'};
		my $sth = $connect->prepare($sql);
		$sth->execute();
		my $arrayref = $sth->fetchrow_arrayref();
		$sth->finish();
		my $sampleName = $accRef->{'sample'};
		my $targetOrNot = "";
		$targetOrNot = '*' if(exists($sample2accMap->{$sampleName}));
		printf STDERR "%s %s  <-  %s\t%s_variant\t%s\n", $targetOrNot, $accRef->{'accession_id'}, $sampleName, $accRef->{'accession_db'}, $arrayref->[0];
	}
}

#### 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 "Variant registrations(by tasuke_variant_vcf.pl) can be parallelized when the number of VCF samples is large.\n";
	print STDERR "Automatically assign AccessionIDs to each threads.\n\n";
#	print STDERR "Parallelize variant registrations(by tasuke_variant_vcf.pl), and automatically assign AccessionIDs to each thread.\n\n";
	print STDERR "Usage: $0  -db [database] -u [user] -p [password] -f [VCFfile]\n\n";
	print STDERR "Either required (*1):\n";
	print STDERR "  -n <IDs> : Comma-separated AccessionID list (Corresponds to the order of VCF sample names) (*2)\n";
	print STDERR "  -m <Path>: Path to \"SampleName > AccID\" correspondence table. VCFsampleName[,]AccID[\\n]...\n";
	print STDERR "  (none)   : Consider VCF sample name as AccessionID.\n";
	print STDERR "Other options:\n";
	print STDERR "  -h <remote host> : (default: localhost)\n";
	print STDERR "  -z : Register GT:0/0 variant(not by default)\n";
	print STDERR "  -t <num> : Number of threads(default: 4)\n";
	print STDERR "  -k : TEST mode. Check AccessionIDs, and output commands for each thread, but do not perform DB registration.\n";
	print STDERR "Sub-action:\n";
	print STDERR "  -a : Append variant information to a table that already contains records.(*3) Records with the same position will be registered multiple times.\n";
	print STDERR "  -r : Delete variant information for the specified AccessionIDs. AccessionID itself is not deleted.\n";
	print STDERR "\n";
	print STDERR "(*1) priority is n > m, Only the AccIDs specified here will be registered.\n";
	print STDERR "(*2) See description in tasuke_variant_vcf.pl on the Wiki.\n";
	print STDERR "(*3) By default, registration is skipped for tables where records already exist.\n";
	print STDERR "\n";
}
