#!/usr/bin/perl -w

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

use strict;
use warnings;
use DBI;
use Getopt::Long;
use File::Basename;
use Cwd;
use Scalar::Util qw(looks_like_number);
use FindBin qw($Bin);
use lib "$Bin";
use CommonLibTASUKE;
use TasukeVcfEncode;
use Data::Dumper;

my $NAME   = '';
my $CREATE = 0;
my $RESTORE= 0;
my $LIST   = 0;
my $REGGT00 = 0;
my $VCF    = '';
my $DB     = '';
my $HOST   = 'localhost';
my $USER   = '';
my $PASS   = '';
#my @NAMELIST;
my $NUMSAMPLE = 0;
my $SQLLDCOL = '';
my $FORCEFLG = 0;
my $APPENDFLG = 0;
#my $CHECKONLY = 0;
my $COMMONONLY = 0;
my $OUTDIR = '';

my $UNIT_COMMON_SQL = 100;	# Common: Bulk insert SQL Records. 
				# The EFF information is large. There is a limit to the size of the query, so it cannot be made large.
my $UNIT_COMMON_COMMIT = 100;	# Common: Number of SQL executions to commit
my $UNIT_M_SQL = 4000;		# SampleInfo: Bulk insert SQL Records. 15col * 4000row < 65535limit
my $UNIT_M_COMMIT = 2;		# SampleInfo: Number of SQL executions to commit

my $commonMysql = new CommonMysql();
my $vcfEncode = new TasukeVcfEncode();

################################################################
# Checking parameters and database connection
################################################################
GetOptions( 'n|name=s'      => \$NAME,
            'l|list'        => \$LIST,
            'r|restore'     => \$RESTORE,
            'z|reggt00'     => \$REGGT00,
            'f|file=s'      => \$VCF,
            'o|outdir=s'    => \$OUTDIR,
            'db|database=s' => \$DB,
            'h|host=s'      => \$HOST,
            'u|user=s'      => \$USER,
            'p|password=s'  => \$PASS,
            'c|force'       => \$FORCEFLG,
            'a|append'      => \$APPENDFLG,
            'm|commononly'  => \$COMMONONLY
);

if ($DB eq '' || ($RESTORE && !$NAME && !$COMMONONLY)) {
    _error_method();  
    exit(1);
}
if(!$RESTORE && !$LIST && (!$VCF || !(-e $VCF))) {
	print "Query file/dir is required: $VCF\n";
	exit(1);
}
if($OUTDIR && !(-d $OUTDIR)) {
	print "Outdir(-o) not found: $OUTDIR\n";
	exit(1);
}

$vcfEncode->{'useZero'} = $REGGT00 if($REGGT00);

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

my $modeIsMultivcf = $commonMysql->modeIsPlusMultiVcf($connect);
if(!$modeIsMultivcf) {
	print "This program is for \"Multi-sample VCF\" mode only.";
	exit(1);
}

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

my $TABLE;
if($COMMONONLY) {
	if(!(defined _check_exists_table("sys_variant"))) {
		print "\"sys_variant\" table does not exists.\n";
		print "Did you execute 'tasuke_init.pl' at the first?\n";
		_exit(1);
	} else {
		$TABLE = "sys_variant";
		if ( !$FORCEFLG && !$RESTORE && defined _check_var_table() ) {
			print "\"sys_variant\" data is already installed.\n";
			if($APPENDFLG) {
				print "Since it's in \"Append\" mode, data will be added to the existing table. This message will not be displayed in forced mode(\"-c\").\n";
				print "Do you want to continue? [y|n] ";
				confirmYN();
			} else {
				print "If you want to install again, \nplease run this tool again with \"-r\" option for deleting the variants.\n";
				_exit(1);
			}
		}
	}
	if($FORCEFLG) {
		print "> Since force flag(-c) is attached, registration is forced.\n";
		} else {
		if($RESTORE) {
			print "> Are you sure deleting sys_variant(site-common variant information) table?[y|n] ";
		} else {
			print "> Are you sure registering sys_variant(site-common variant information) table?[y|n] ";
		}
		confirmYN();
	}
}

###########################################################
# Listing track
###########################################################
if ($LIST) {
    print "[List of accessions] :\n";
    foreach my $l (_lists()) {
        print "$l\n";
    }
    _exit();
}

# for tasuke_variant_multivcf.pl. Wait for all parallel Threads to be generated successfully.
sleep 1;

###########################################################
# Restoring sys_variant table
###########################################################
if($RESTORE && $COMMONONLY) {
	_dbexe("TRUNCATE TABLE sys_variant");
	print "\"sys_variant\" table is restored.\n";
	_exit();
}

my ($chrId2info, $chrSeq2Info) = getChrInfoMap();
my $finalCode = 0;
if($COMMONONLY) {
	if(checkTableRegistration('sys_variant')) {
		$vcfEncode->prepareCommon($VCF, $OUTDIR) if($OUTDIR);

		_dbexe("LOCK TABLES sys_variant WRITE");
		$finalCode =_insertVarGATKm_common($chrId2info, $VCF);

		# Optimize
		if(1) {
			my $sqlOpt = "optimize table sys_variant";
			my $sthOpt = $connect->prepare($sqlOpt);
			$sthOpt->execute();
			my $resOpt = $sthOpt->fetchrow_arrayref();
			print join("\t", @{$resOpt}). "\n";
		}		
		_dbexe("UNLOCK TABLES");

		$vcfEncode->writePropFile() if($OUTDIR);	# Create prop file.
	}
} else {
	###########################################################
	# Checking accession id
	###########################################################
	$TABLE =_check_db();
	my @nameArr = split(/,/,$NAME);
	my @NAMELIST;
	for(my $i=0; $i<scalar(@nameArr); $i++){
		$NAME = $nameArr[$i];
		next if(!$NAME);

		push @NAMELIST, [$NAME, $i];
		my $accdb =_check_db();
		if(!defined $accdb){
			print "'$nameArr[$i]' is not exists on list.\n\n";
			print "- Here is list you loaded\n";
			foreach my $l (_lists()) {
				print "$l\n";
			}
			_exit(1);
		}
	}
	if(scalar(@NAMELIST) == 0){
		print "'$NAME' is not exists on list.\n\n";
		print "- Here is list you loaded\n";
		foreach my $l (_lists()) {
			print "$l\n";
		}
		_exit(1);
	}

	###########################################################
	# Restoring variant table
	###########################################################
=head
	if ($RESTORE) {
		_dbexe("TRUNCATE TABLE $TABLE");
		print "VARIANT $NAME table is restored.\n";
		_exit();
	}
=cut

	###########################################################
	# Checking ref-seq already installed or not
	###########################################################
	  # Changed to skip registered tables.
	  if(!$FORCEFLG) {
	    foreach my $nm (@NAMELIST) {
	      next if(!$nm->[0]);
	      $NAME = $nm->[0];
	      $TABLE =_check_db();
	      $TABLE .= "_variant";
	      if(defined _check_var_table()){
	        print "$nm: The variants are already installed.\n";
		if($APPENDFLG) {
			print "Since it's in \"Append\" mode, data will be added to the existing table.\n> [y|n] ";
			confirmYN();
		} else { 
			print "If you want to install again, \nplease run this tool again with \"-r\" option for deleting the variants.\n";
			_exit(1);
		}
	      }
	    }
	  }

	###########################################################
	# Creating variants table
	###########################################################
	$vcfEncode->prepare($VCF);
	foreach my $ref (@NAMELIST){
		if($ref) {
			$NAME = $ref->[0];
			_print_time("*[".$NAME." variant] begin");
			$TABLE = _check_db();
			$TABLE .= "_variant";

			if(checkTableRegistration($TABLE)) {
				_dbexe("LOCK TABLES $TABLE WRITE,sys_variant READ");
				$finalCode =_insertVarGATKm_m($ref->[1]);
				_dbexe("UNLOCK TABLES");
			}
			_print_time(" [".$NAME." variant] complete");
		}
		$NUMSAMPLE++;
	}
}
$vcfEncode->finish();

_exit($finalCode);

sub checkTableRegistration {
	my($table) = @_;
	my $execRegist = 0;
	if(checkAlreadyRegisted($table)) {
		print "The variant is already installed.\n";
		if($APPENDFLG) {
			print "Since it's in \"Append\" mode, data will be added.\n";
			$execRegist = 1;
		} elsif($FORCEFLG) {
			print "Registration will be skipped. To update, restore(-r) the table first.\n";
		}
	} else {
		$execRegist = 1;
	}
	return $execRegist;
}

## Update variants ##############################
sub _insertVarGATKm_common {
	my ($chrId2Info, $file) = @_;
	my $code = 0;

	my $currentUnitSql = 0;
	my $currentUnitCommit = 0;

	$| = 1;		# AutoFlush:On

	my $sqlCol = "insert into sys_variant (chr,pos_vcf,REF,ALTS,VDB,MQ,FQ,SP,EFF,AF1,AC1,QUAL) values ";
	my $sqlVal = "(?,?,?,?,?,?,?,?,?,?,?,?),";
	my @sqlParams;
	my @valsTemplate = ('',undef,'','',undef,undef,undef,undef,'','','',undef);

	my $fh = _getFileHandle($file);
	my $total = 0;
	while(<$fh>) {
		chomp;
		my $datarow = $_;
		next if($datarow !~ m/\S/ || $datarow =~ m/^#/);

		my @line = split(/\t/, $datarow);
		my $CHR = $line[0];
		next if(!exists($chrId2Info->{$CHR}));
		my $POS = int($line[1]);
		my $REF = $line[3];
		my $ALTS = $line[4];
		my $QUAL = int($line[5]);
		my %valsHash;
		my $eff = '';
		$eff = $2 if($line[7] =~ m/(ANN|EFF)\=([^\;]+)/);

		# Reduce momory.
		$line[0] = '';
		$line[1] = '';
		$line[2] = '';
		$line[3] = '';
		$line[4] = '';
		$line[5] = '';
		$line[6] = '';
		$line[7] = '';
=head
		my @infoArr = split(/;/, $line[7]);
		foreach my $info (@infoArr) {
			my @keyVal = split(/\=/, $info);
			if($keyVal[0] eq 'ANN' || $keyVal[0] eq 'EFF') {
				$eff = $keyVal[1];
				$eff =~ s/['"]/ /g;
				last;
			}
		}
		undef @infoArr;	# Release momory.
=cut

		my $refLen = length($REF);
		my @vals = @valsTemplate;	# Deep copy
		$vals[0] = $CHR;
		$vals[1] = $POS;
		$vals[2] = $REF;
		$vals[3] = $ALTS;
		$vals[4] = $valsHash{'VDB'} if(exists($valsHash{'VDB'}));
		$vals[5] = $valsHash{'MQ'} if(exists($valsHash{'MQ'}));
		$vals[6] = $valsHash{'FQ'} if(exists($valsHash{'FQ'}));
		$vals[7] = $valsHash{'SP'} if(exists($valsHash{'SP'}));
		$vals[8] = $eff;
		$vals[9] = $valsHash{'AF'} if(exists($valsHash{'AF'}));
		$vals[10] = $valsHash{'AC'} if(exists($valsHash{'AC'}));
		$vals[11] = $QUAL;
		$total++;
		$currentUnitSql++;

		push @sqlParams, [\@line, \@vals];
		if($currentUnitSql >= $UNIT_COMMON_SQL) {
			_execVarGATKm_Common($sqlCol, $sqlVal, \@sqlParams);

			@sqlParams = ();
			$currentUnitSql = 0;

			$currentUnitCommit++;
			if($currentUnitCommit >= $UNIT_COMMON_COMMIT) {
				printf("Progress: %s\t%d\t RegisteredVariants:%d\t", $CHR, $POS, $total);

				$connect->commit();
				print " ->Commit";
				$currentUnitCommit = 0;

				$vcfEncode->writeSampleInfo();
				print " ->WriteTmpFile\n";
			}
		}

	}
	_execVarGATKm_Common($sqlCol, $sqlVal, \@sqlParams);
	@sqlParams = ();
	$connect->commit();
	$vcfEncode->writeSampleInfo();

	close($fh);

	$| = 0;		# AutoFlush:Off

	return $code;
}
sub _execSql {
	my($sqlCol, $sqlVal, $sqlParams) = @_;

	my $sql = $sqlCol;
	my @arr;
	my $num = scalar(@{$sqlParams});
	return if($num == 0);
	for(my $i=0; $i<$num; $i++) {
		my $ref = $sqlParams->[$i];
		$sql .= $sqlVal;
		push @arr, @{$ref->[1]};
	}
	$sql =~ s/,$//;
	my $sth = $connect->prepare($sql);
	my $stat = $sth->execute(@arr);
	my $seqIdTop = -1;
	if($stat) {
		$seqIdTop = $sth->{mysql_insertid};
	} else {
		print "Error: $!\n". Dumper($sqlParams). "\n";
	}
	$sth->finish();
	return $seqIdTop;
}
sub _execVarGATKm_Common {
	my($sqlCol, $sqlVal, $sqlParams) = @_;

	my $seqId = _execSql($sqlCol, $sqlVal, $sqlParams);
	if($seqId > 0 && $OUTDIR) {
		foreach my $ref (@{$sqlParams}) {
			$vcfEncode->encodeOnerow($seqId, $ref->[0]);
			$seqId++;
		}
	}
}
sub _insertVarGATKm_m {
	my ($idx) = @_;
	my $code = 0;

	my $INDEL = "Dels=";

	my $currentUnitSql = 0;
	my $currentUnitCommit = 0;

	$| = 1;	# AutoFlush:On


	my $sqlCol = "insert into $TABLE (chr,pos_vcf,pos,variant,ALT,DP,DP1,DP2,DP3,DP4,GT,GTROW,PL,GQ,EFF_GRP) values ";
	my $sqlVal = "(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?),";
	my @sqlParams;

	my $sqlAppend = '';
#	if($APPENDFLG) {
	if(1) {
		my $fhPre = $vcfEncode->openBin($idx);
		my $valHash = $vcfEncode->decodeSampleInfo($fhPre);
		$sqlAppend = sprintf "where seq>=%d", $valHash->{'seqId'};
		close($fhPre);
	}
	my $sqlSysvariant = "select seq,chr,pos_vcf,REF,ALTS,EFF from sys_variant $sqlAppend order by seq asc";
	my $sthSysvariant = $connect->prepare($sqlSysvariant);
	$sthSysvariant->execute();

	my $fhBin = $vcfEncode->openBin($idx);
	my $total = 0;
	while(my $arrRef = $sthSysvariant->fetchrow_arrayref()) {
		$total++;

		my $seqId = $arrRef->[0];
		my $CHR = $arrRef->[1];
		my $POS = $arrRef->[2];
		my $REF = $arrRef->[3];
		my $ALTS = $arrRef->[4];
		my $EFF = $arrRef->[5];

		my $valHash = $vcfEncode->decodeSampleInfo($fhBin);
		last if(!$valHash);
		if($valHash->{'seqId'} != $seqId) {
			printf "There is an inconsistency between DB and  temporary file: %d / %d\n", $valHash->{'seqId'}, $seqId;
#			_exit(1);
			return 1;
		} elsif(!exists($valHash->{'GT'})) {
			# Invalid site -> skip
			next;
		}

		$vcfEncode->setInfo(
			{'POS' => $POS, 'REF' => $REF, 'ALTS' => $ALTS, 'EFF' => $EFF},
			$valHash
		);

		my @valsArr = ($CHR, $POS, $valHash->{'pos'}, $valHash->{'variant'},
			undef,undef,undef,undef,undef,undef,$valHash->{'GT'},undef,undef,undef,undef);
		$valsArr[4] = $valHash->{'ALT'} if(exists($valHash->{'ALT'}));
		$valsArr[5] = $valHash->{'DP'} if(exists($valHash->{'DP'}));
		if(exists($valHash->{'DP1'})) {
			$valsArr[6] = $valHash->{'DP1'};
			$valsArr[7] = $valHash->{'DP2'};
			$valsArr[8] = $valHash->{'DP3'};
			$valsArr[9] = $valHash->{'DP4'};
		}
		$valsArr[11] = $valHash->{'GTROW'} if(exists($valHash->{'GTROW'}));
		$valsArr[12] = $valHash->{'PL'} if(exists($valHash->{'PL'}));
		$valsArr[13] = $valHash->{'GQ'} if(exists($valHash->{'GQ'}));
		$valsArr[14] = $valHash->{'EFF_GRP'} if(exists($valHash->{'EFF_GRP'}));

		$currentUnitSql++;
		push @sqlParams, [undef, \@valsArr];
		if($currentUnitSql >= $UNIT_M_SQL) {
#			printf("Progress: %s\t%d\t RegisteredVariants:%d\n", $CHR, $POS, $total);
			_execSql($sqlCol, $sqlVal, \@sqlParams);
			@sqlParams = ();
			$currentUnitSql = 0;

			$currentUnitCommit++;
			if($currentUnitCommit >= $UNIT_M_COMMIT) {
				printf("Progress: %s\t%d\t RegisteredVariants:%d\t", $CHR, $POS, $total);
				print " ->Commit\n";
				$connect->commit();
				$currentUnitCommit = 0;
			}
		}
	}
	_execSql($sqlCol, $sqlVal, \@sqlParams);
	@sqlParams = ();
	$connect->commit();
	$sthSysvariant->finish();
	close($fhBin);

	$| = 0;	# AutoFlush:Off

	return $code;
}
### Checking id on the list ####################################
sub _check_db {
    my $val = $connect->selectrow_array("select accession_db from sys_accession where accession_id='$NAME'");
    return $val;
}

# Accession list
sub _lists {
    my $val;
    my @lists;
    my $query = "SELECT accession_id FROM sys_accession";
    my $query_handle = $connect->prepare($query);
    $query_handle->execute();
    $query_handle->bind_columns(undef, \$val);
    while($query_handle->fetch()) {
        push @lists, $val;
    }
    return @lists;
}

sub getChrInfoMap {
	my %id2info;
	my %seq2Info;
	my $sth = $connect->prepare("SELECT * FROM sys_chromosome");
	$sth->execute();
	while(my $hashref = $sth->fetchrow_hashref()) {
		$id2info{$hashref->{'id'}} = $hashref;
		$seq2Info{$hashref->{'seq'}} = $hashref;
	}
	return (\%id2info, \%seq2Info);
}

sub _print_time {
    my ($str) = @_;
    my ($sec,$min,$hour,$mday,$mon,$year,$wday,$yday,$isdst) = localtime;
    $year += 1900;
    $mon += 1;
    my $datetime = sprintf "%04d-%02d-%02d %02d:%02d:%02d", $year, $mon, $mday, $hour, $min, $sec;
    print "$str : $datetime\n";
}
### read input file ##########################################
sub _getFileHandle {
    my($file) = @_;

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


### Execute DB ################################################
sub _dbexe {
    my ($query) = @_;
    my $query_handle = $connect->prepare($query);
    $query_handle->execute();
    $query_handle->finish();
}

sub _error_method {
    print "\nRegistering variant table\n";
    print "tasuke_variant_vcf.pl -n [id] -f [variant_file(.vcf)] -t [samtools or gatk or gatkm] -db [database] -u [user] -p [password] \n( To register GT:0/0 variant: -z) \n\n";
    print "Append variant table\n";
    print "tasuke_variant_vcf.pl -a -n [id] -f [variant_file(.vcf)] -t [samtools or gatk or gatkm] -db [database] -u [user] -p [password]\n\n";
    print "Restoring table\n";
    print "tasuke_variant_vcf.pl -r -n [name] -db [database] -u [user] -p [password]\n\n";
    print "Restoring site-common table(sys_variant) for \"mvcf\" DB mode\n";
    print "tasuke_variant_vcf.pl -r -m -db [database] -u [user] -p [password]\n\n";
    print "Listing accession\n";
    print "tasuke_variant_vcf.pl -l -db [database] -u [user] -p [password]\n\n";
}

### Checking table ####################################
sub _check_exists_table {
    my ($t) = @_;
    my $val;
    $val = $connect->selectrow_array("show tables like '$t'");
    return $val;
}

### Checking tablespace ####################################
sub _check_var_table {
    my $res;
    $res = $connect->selectrow_array("select * from $TABLE limit 1");
    return $res;
}
sub checkAlreadyRegisted {
        my($table) = @_;
        my $sql = sprintf 'select * from %s limit 1', $table;
        my $sth = $connect->prepare($sql);
        $sth->execute();
        my $res = $sth->fetchrow_arrayref();
        return $res;
}

#### Disconnection and exit #######################################
sub _exit {
    my ($arg) = @_;
    my $code = 0;
    if(defined $arg){$code = $arg;}
    $connect->disconnect or warn "Disconnection error: $DBI::errstr\n";
    exit($code);
}
=head
sub _isNumber {
	my($str) = @_;
	if($str =~ m/[^\+\-\d\.e]/i) {
		return 0;
	} else {
		return 1;
	}
}
=cut
sub confirmYN {
                while (defined(my $line = <STDIN>)){
                    chomp($line);
                    if ( $line eq "y" || $line eq "Y" ) {
                        last;
                    } elsif ( $line eq "n" || $line eq "N" ) {
                        print "finished.\n";
                        _exit();
                    } else {
                        print "> [y|n] ";
                    }
                }
}
