#!/usr/bin/perl -w

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

#TODO Support for 32bit INT(L) when block size is large(>500,000).

use strict;
use warnings;
use DBI;
use Getopt::Long;
use FindBin qw($Bin);
use lib "$Bin";
use POSIX;
#use CommonLibTASUKE;
#use Data::Dumper;

my $DB     = '';
my $HOST   = 'localhost';
my $USER   = '';
my $PASS   = '';
my $OUTFILE = '';
my $ORDERFILE   = '';
my $BLOCKW = 100000;
my $GQ = 1;
my $DP = 1;
my $DP34 = 1;
my $NOLIMIT = 0;

#my $commonMysql = new CommonMysql();

################################################################
# 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,
	'w|blockw=i'	=> \$BLOCKW,
	'x|nolimit'	=> \$NOLIMIT
);

my $connect;
if ( !$DB || !$USER ) {
	_error_method();
	_exit(1);
}
my $validIds = undef;
if($ORDERFILE) {
	if( !(-f $ORDERFILE) ) {
		print STDERR "Accession list file(common name is \"order.conf\") cannot be found: $ORDERFILE\n";
		die;
	} else {
		$validIds = {};
		my $cnt = 0;
		local *FH;
		open(FH, "<$ORDERFILE") or die "Cant open file: $ORDERFILE\n";
		while(<FH>) {
			chomp;
			my $onerow = $_;
			next if($onerow !~ m/\S/ || $onerow =~ m/^\#/);

			$onerow =~ s/^\s+//;
			my @vals = split /\s+/, $onerow;
			my $acc = $vals[0];
			$acc =~ s/\W+$//;
			$validIds->{$acc} = 1;
			$cnt++;
		}
		print STDERR "Use order file. Accs: $cnt\n";
	}
}

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

my $accArr = getAccs($connect, $validIds);
my $chrArr = getChrs($connect);
my $outHandle = *STDOUT;
if($OUTFILE) {
	open($outHandle, ">$OUTFILE");
}

my $packFormatBlockW = 'L';	# Block-width. Unsigned 32bit int(0~4,294,967,295).
my $packFormatDens = 'S';	# Block-value(density). Unsigned 16bit int(0~65,535).
my $packFormatBlocks = 'L';	# Block-num in CHR. Unsigned 32bit int(0~4,294,967,295).

print $outHandle pack($packFormatBlockW, $BLOCKW);
print STDERR "BlockW: $BLOCKW\n";

printStrsHeader($outHandle, $accArr, 'accession_db');
printStrsHeader($outHandle, $chrArr, 'id');

my $i = 0;
foreach my $accRef (@{$accArr}) {
	my $sql = "select FLOOR(pos/". $BLOCKW. ") as bucket,count(1),".
		"sum(CASE WHEN EFF_GRP=3 THEN 1 ELSE 0 END) ".
#		"sum(CASE WHEN EFF_GRP=2 THEN 10 WHEN EFF_GRP=3 THEN 60 ELSE EFF_GRP+1 END) ".
		"from ". $accRef->{'accession_db'}. "_variant where chr=? and GT!='./.' ";
	$sql .= "and GQ>=? and DP>=? and (DP3+DP4)>=? " if(!$NOLIMIT);
	$sql .= "group by bucket";
	my $sth = $connect->prepare($sql);
	my $j = 0;
	foreach my $chrRef (@{$chrArr}) {
		my @chrBlockArr;
		my $chrBlocks = ceil($chrRef->{'size'} / $BLOCKW);
		for(my $i=0; $i<$chrBlocks; $i++) {
			$chrBlockArr[$i] = [0, 0];
		}

#		printf STDERR "%s\t%s\n", $accRef->{'accession_id'}, $chrRef->{'id'};
		my @params = ($chrRef->{'id'});
		push @params, ($GQ, $DP, $DP34) if(!$NOLIMIT);
		$sth->execute(@params);
		printf STDERR "%d-%02d\t%s\t%s\n", ($i+1), ($j+1), $accRef->{'accession_id'}, $chrRef->{'id'};
		while(my $arrRef = $sth->fetchrow_arrayref()) {
#printf STDERR "%s\t%s\t%d\t%d\n", $accRef->{'accession_id'}, $chrRef->{'id'}, $arrRef->[0], $arrRef->[1];
			my $cnt = $arrRef->[1] + 0;	# "+0" is integer conversion.
				$cnt = 65535 if($cnt > 65535);
			my $eff = $arrRef->[2] + 0;
				$eff = 65535 if($eff > 65535);
			$chrBlockArr[$arrRef->[0] + 0] = [$cnt, $eff];
		}
		print $outHandle pack($packFormatBlocks, $chrBlocks);
		foreach my $dens (@chrBlockArr) {
			print $outHandle pack($packFormatDens, $dens->[0]);
			print $outHandle pack($packFormatDens, $dens->[1]);
		}
		$j++;
	}
	$sth->finish();

	$i++;
}

close($outHandle);
_exit(0);

sub getChrs {
	my($connect) = @_;

	my @array;
	my $sql = "select * from sys_chromosome order by seq";
	my $sth = $connect->prepare($sql);
	$sth->execute();
	while(my $hashref = $sth->fetchrow_hashref()) {
		push @array, $hashref;
	}
	$sth->finish();
	return \@array;
}
sub getAccs {
	my($connect, $validIds) = @_;
	
	my @array;
	my $sql = "select * from sys_accession";
	my $sth = $connect->prepare($sql);
	$sth->execute();
	while(my $hashref = $sth->fetchrow_hashref()) {
		next if(defined $validIds && !exists($validIds->{$hashref->{"accession_id"}}));
		push @array, $hashref;
	}
	my @arraySorted  = sort {getAccDbNum($a->{'accession_db'}) <=> getAccDbNum($b->{'accession_db'})} @array;
	return \@arraySorted;
}
sub printStrsHeader {
	my($outHandle, $arr, $key) = @_;
	my $num  = @{$arr};
	print $outHandle pack('S', $num);
	foreach my $ref (@{$arr}) {
		my $str = $ref->{$key};
		printf $outHandle "%s%s", pack('S', length($str)), $str;	#TODO For multi-byte str.
	}
} 
sub getAccDbNum {
	my($accdb) = @_;
	if($accdb =~ m/^d(\d+)$/) {
		return $1;
	} else {
		return 0;
	}
}
=head
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);
}
=cut
sub _error_method {
	print STDERR "Create a variant count file for TASUKE+ Macro Viewer.\n\n";
	print STDERR "Usage: tasuke_mv_variantcount.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 "  -w <int>:  The block size (bp) for counting variants.(*1) (Default: $BLOCKW)\n";
	print STDERR "  -x: Also counts variants with missing GQ or DP.(*2)\n";
	print STDERR "\n";
	print STDERR "(*1) A smaller value will result in higher data resolution, but at the cost of larger file sizes and slower loading times.\n"; 
	print STDERR "(*2) By default, each must have a value >= 1.\n";
	print STDERR "\n";
}

