#!/usr/bin/perl -w

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

use strict;
use DBI;
use Getopt::Long;
use File::Basename qw/basename dirname/;
use File::Temp;
use FindBin qw($Bin);
use lib "$Bin";
use TasukeVariantExistMap;

my $BIN_DMATRIX = 'converter/makeBinaryDistanceMatrix';

my $DB		= '';
my $HOST	= 'localhost';
my $USER	= '';
my $PASS	= '';
my $OUTFILE	= '';
my $ORDERFILE	= '';
my $CHRS	= '';
my $CHECK_DEPTH	= 0;
my $DEPTH_IN_MEM = 0;
my $SKIP_NA_COL	= 0;
my $AS_HAPLOID = 0;
my $NOT_DELETE_BINTABLE = 0;
my $METHOD = '';
my $VTYPES = '';
my $TMPDIR = '';
my $PHYLIPDIR = '';

################################################################
# Checking parameters and database connection
################################################################
GetOptions(
	'db|database=s'	=> \$DB,
	'h|host=s'	=> \$HOST,
	'u|user=s'	=> \$USER,
	'p|password=s'	=> \$PASS,
	'o|outfile=s'	=> \$OUTFILE,
	'r|orderfile=s'	=> \$ORDERFILE,
	'c|chrs=s'	=> \$CHRS,
	'a|checkdepth'	=> \$CHECK_DEPTH,
	'b|depthinmem'	=> \$DEPTH_IN_MEM,
	'n|skipnacol'	=> \$SKIP_NA_COL,
	'd|haploid'	=> \$AS_HAPLOID,
	'l|nodel'	=> \$NOT_DELETE_BINTABLE,
	'm|method=s'	=> \$METHOD,
	'v|vtypes=s'	=> \$VTYPES,
	't|tmpdir=s'	=> \$TMPDIR,
	'y|phylip=s'	=> \$PHYLIPDIR
);
$DEPTH_IN_MEM = 1;

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

if( !(-f "$Bin/$BIN_DMATRIX") ) {
	print STDERR "Dependency cannot be found: $Bin/$BIN_DMATRIX\n";
	exit(1);
}
if(!$OUTFILE) {
	print STDERR "[Warning!!!] Outfile(-o) is not specified. Output result to STDOUT.\n";
}

my %methods = ('jaccard' => 1, 'dice' => 1, 'simple' => 1, 'soergel' => 1);
if($METHOD && !exists($methods{$METHOD})) {
	print STDERR "Unknown method type: $METHOD\n";
	exit(1);
} elsif(!$METHOD) {
	$METHOD = 'simple';
}

my $obj = TasukeVariantExistMap->new('dbh' => $connect, 'db' => $DB);
$obj->verbose(1);

if($VTYPES) {
	if($obj->replaceTargetVariantTypesFromStr($VTYPES)) {
		print STDERR "Abnormal Varaint type: -t $VTYPES\n";
		exit(1);
	}
}

if(!$TMPDIR) {
	$TMPDIR = '/tmp';
} else {
	if( !(-d $TMPDIR) ) {
		print STDERR "Tempdir not found: $TMPDIR\n";
		exit(1);
	}
}

if($PHYLIPDIR) {
	if(-d $PHYLIPDIR && -d $PHYLIPDIR. "/exe" && -f $PHYLIPDIR. "/exe/neighbor") {
		# Phylip path OK
		if( !(-f $PHYLIPDIR. "/exe/retree") ) {
			printf STDERR "Phylip retree(%s/exe/retree) not found. No Mid-point rooting will be performed.\n", $PHYLIPDIR;
		}
	} else {
		print STDERR "Incorrect Phylip directory. Specify a directory that contains \"exec/neighbor\": $PHYLIPDIR\n";
		exit(1);
	}
}

my $binaryTableFile = $TMPDIR. '/tasuke_bittable'. $$;
my $binaryTableFileCreated = 0;
$SIG{'INT'} = '_SIGINThandler'; # Delete temporary files by terminating with Ctrl+C
$SIG{'KILL'} = '_SIGINThandler';
$SIG{'TERM'} = '_SIGINThandler';

eval {

###############################################################
# Scan DB Start and make CSV
###############################################################
if($CHRS) {
	my @chrArray = split(/\s*,\s*/, $CHRS);
	checkChromosome(\@chrArray);
	$obj->setChrs(\@chrArray);
}
if($ORDERFILE) {
	if( !(-f $ORDERFILE) ) {
		print STDERR "Accession list file(common name is \"order.conf\") cannot be found: $ORDERFILE\n";
		die;
	} else {
		$obj->setAccsByOrderFile($ORDERFILE);
	}
}

$obj->setCheckDepth($CHECK_DEPTH);
$obj->setDepthInMem($DEPTH_IN_MEM);
my $workDir = File::Temp->newdir(DIR => $TMPDIR, CLEANUP => 1);
$obj->setTempDir($workDir);

my $binaryTableFileFh;
open($binaryTableFileFh, ">$binaryTableFile") or die "Can't open $binaryTableFile: $!";
$obj->output($binaryTableFileFh);
close($binaryTableFileFh);
$binaryTableFileCreated = 1;

$connect->disconnect();

=head
if($NOT_DELETE_BINTABLE) {
	eval {
		my $csvFile;
		if($OUTFILE) {
			$csvFile = $OUTFILE. '.btbl';
		} else {
			$csvFile = './'. basename($binaryTableFile). '.btbl';
		}
		system("mv $binaryTableFile $csvFile") == 0 or die "$!\nCant't move $binaryTableFile. Proceed without file moving.\n";
		$binaryTableFile = $csvFile;
	};
}
=cut

###############################################################
# Convert CSV -> PhylipDistanceMatrix
###############################################################
print STDERR "Createing distance matrix\n";
my $cmd = "$Bin/$BIN_DMATRIX -i $binaryTableFile";
$cmd .= " -n" if($SKIP_NA_COL);
$cmd .= sprintf(" -m %s", $METHOD) if($METHOD);
$cmd .= " -d" if($AS_HAPLOID);
$cmd .= " > $OUTFILE" if($OUTFILE);
print STDERR "Command: $cmd\n";
system($cmd) == 0 or die "Can't execute $cmd: $!";

if($PHYLIPDIR) {
	# In phylip, sample name can be up to 10 characters long.
	# If the name exceeds this, it must be rewritten.
	my $maxLenSmplName = 0;
	local *FH;
	open(FH, "<$OUTFILE");
	<FH>;	# Discard header(number of samples)
	while(<FH>) {
		chomp;
		my $datarow = $_;
		next if($datarow !~ m/\S/);

		m/^(\S+)/;
		my $smplNamelen = length($1);
		$maxLenSmplName = $smplNamelen if($maxLenSmplName < $smplNamelen);
	}

	my $newickPath = sprintf "%s.nwk", $OUTFILE;
#	if($maxLenSmplName > 10) {
		my $stat = $obj->makeNewickByPhylipiViaRename($PHYLIPDIR, $OUTFILE, $newickPath);
#	} else {
#		my $stat = $obj->makeNewickByPhylip($PHYLIPDIR, $OUTFILE, $newickPath);
#	}
}

###############################################################
# Finished
###############################################################
#if($NOT_DELETE_BINTABLE) {
#	print STDERR "\n\nSaved binaryTableFile: $binaryTableFile\n";
#	print STDERR "If you do not need the above file, be sure to delete it.\n";
#}
print STDERR "\n  -----------\n";
print STDERR "   Next step\n";
print STDERR "  -----------\n\n";
print STDERR "   1. Place the created distance matrix file in any directory(the path permitted to read by www user).\n";
print STDERR "   2. Set that path in conf/config.php as shown below.\n";
print STDERR "      \$distanceMatrixPath = \"PathTo/DistanceMatrixFile\";\n";
print STDERR "\n";
print STDERR "   *Tips: If you keep(\"-l\") binary table file, For example, you can recreate distance matrix as follows.\n";
printf STDERR "   %s -m simple -i %s.btbl > %s2\n", $BIN_DMATRIX, $OUTFILE, $OUTFILE;
print STDERR "\n#######################################################################\n";

endProcessing();
#if($NOT_DELETE_BINTABLE) {
#	print STDERR "\nSaved binaryTableFile: $binaryTableFile\n";
#	print STDERR "If you do not need the above file, be sure to delete it.\n";
#}
};
if ($@) {
	print STDERR "Exception occur: $@";
	endProcessing();
}

exit(0);

sub endProcessing {
	my($forceDel) = @_;
	if(-f $binaryTableFile) {
		if($NOT_DELETE_BINTABLE && $binaryTableFileCreated) {
			eval {
				my $csvFile;
				if($OUTFILE) {
					$csvFile = $OUTFILE. '.btbl';
				} else {
					$csvFile = './'. basename($binaryTableFile). '.btbl';
				}
				my $cmd = "mv $binaryTableFile $csvFile";
#				printf STDERR "$cmd\n";
				system($cmd) == 0 or die "$!\nCant't move $binaryTableFile. Proceed without file moving. $binaryTableFile remains\n";
				$binaryTableFile = $csvFile;
			};
			if (my $error = $@) {
				print STDERR "$error\n";
			}
		} else {
			printf STDERR "Delete binary table: %s\n", $binaryTableFile;
			unlink($binaryTableFile);
		}
	}
#	unlink($binaryTableFile) if((!$NOT_DELETE_BINTABLE || $forceDel) && -f $binaryTableFile);
}
sub checkChromosome {
	my($chrsRef) = @_;

	my $sql = "select 1 from sys_chromosome where id=?";
	my $sth = $connect->prepare($sql);
	foreach my $chr (@{$chrsRef}) {
		$sth->execute($chr);
		if(!$sth->fetchrow_hashref()) {
			print STDERR "$chr is not registered in DB\n";
			die;
		}
	}
	$sth->finish();
}
sub _error_method {
	print STDERR "Create binary(1/0) table based on presence or absence of variants and create a binary distance matrix.\n\n";
	print STDERR "Usage: tasuke_tree_dmatrix.pl -db [database] -u [user] -p [password] -o [outfile]\n";
	print STDERR "Other options (Not required):\n";
	print STDERR "  -h <str>: Host\n";
	print STDERR "  -r <file>: Target accession list file(Generally [tasukeHome]/conf/order.conf) (Default: all Accessions)\n";
	print STDERR "  -y <dir>: Path to Phylip directory(having \"exe/neighbor\" under it). Create an additional Newick with name \"[outfile].nwk\".\n";
	print STDERR "  -c <str>: Target chromosomes separated by commas(Default: all Chromosomes)\n";
	print STDERR "  -v <str>: Specify target variant type. s=SNP, i=INS, d=DEL, order does not matter. ex:\"-v s\" \"-v di\" (Default: sid)\n";
	print STDERR "  -m <str>: Distance matrix calculation method. simple(default)/jaccard/dice/soergel\n";
	print STDERR "  -a : Check DEPTH=NULL and if so, set \"NA\" to that position. It takes a lot of time. If DEPTH info is not registered, This option will have no effect.\n";
	print STDERR "  -b : (This option is invalid, but left for compatibility. DEPTH=NULL positions are always kept in memory.)\n";
#	print STDERR "  -b : [Use with \"-a\"] Store DEPTH=NULL positions on memory. Can be faster but uses more memory\n";
	print STDERR "  -n : [Use with \"-a\"] Cross all accessions and skip Column where NA exists\n"; 
#	print STDERR "  -d : Simply calculate as haploid. 0:Ref(0/0)  1:Alt(1/1, 1/2, 0/1)  NA:DEPTH=NULL\n";
	print STDERR "  -l : Leave binary table with name \"[outfile].btbl\". This file can be reused for distance calculation by $BIN_DMATRIX\n";
	print STDERR "  -t <dir>: Temporary directory for creating binary table. Large datasets may require tens to hundreds GB.(Default: /tmp)\n";
	print STDERR "\n";
}

### SIGINT ###
sub _SIGINThandler{
	endProcessing(1);
	print STDERR "\n\nFinished.\n";
	exit(1);
}
