#!/usr/bin/perl -w

#
# tasuke_ref.pl
# http://tasuke.dna.affrc.go.jp/
#
# Author: Jungsok Kim, Ryutaro Itoh
#
# Copyright (c) 2013 National Institute of Agrobiological Sciences   
#

use strict;
use DBI;
use Getopt::Long;
use File::Copy;
use File::Basename;
use Cwd;
use FindBin qw($Bin);
use lib "$Bin";
use CommonLibTASUKE;

my %CHR2TABLE;
my %TABLE2CHR;

my $TABLE   = "sys_fasta";
my $CREATE  = 0;
my $DELETE  = 0;
my $FILE    = '';
my $DB      = '';
my $HOST    = 'localhost';
my $USER    = '';
my $PASS    = '';
my $RDBMS = '';
my $TABLESPACE = '';

my $commonMysql = new CommonMysql();

################################################################
# Checking parameters and database connection
################################################################
GetOptions( 
            'r|restore'     => \$DELETE,
            'f|file=s'      => \$FILE,
            'db|database=s' => \$DB,
            'h|host=s'      => \$HOST,
            'u|user=s'      => \$USER,
            'b|backend=s'   => \$RDBMS,
            's|space=s'   => \$TABLESPACE,
            'p|password=s'  => \$PASS
);

if ( !$DB || !$USER || (!$DELETE && !-e$FILE)) {
    _error_method();
    exit;
}
if(!$RDBMS){
    $RDBMS = 'mysql';
}

my $connect;
if($RDBMS eq "oracle"){
    if(!$TABLESPACE){
        _error_method();
        exit;
    }
    $connect = DBI->connect('DBI:Oracle:host='.$HOST.';SERVICE_NAME='.$DB.';',$USER,$PASS)
    or die "Can't connect to database: $DBI::errstr\n";
    if(!_check_exists_tablespace($TABLESPACE)){
        print "Tablespace: '$TABLESPACE' is not found.";
        _exit();
    }
}else{
    $connect= DBI->connect('DBI:mysql:database='.$DB.';host='.$HOST,$USER,$PASS)
    or die "Can't connect to database: $DBI::errstr\n";
}

my (@chrs) = _chrs();

# Set Table <=> Chr map
my $col2tableRef = $commonMysql->getCol2TableMap($connect, $DB, 'sys_fasta');
%CHR2TABLE = %{$col2tableRef};
my $numCols = 0;
foreach my $chr (keys(%CHR2TABLE)) {
	my $table = $CHR2TABLE{$chr};
	$TABLE2CHR{$table} = [] if(!exists($TABLE2CHR{$table}));
	push @{$TABLE2CHR{$table}}, $chr;
	$numCols++;
}
#setTable2Chr(\@chrs);

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

###########################################################
# Drop table
###########################################################
if ($DELETE) {
	foreach my $table (keys(%TABLE2CHR)) {
		_dbexe("TRUNCATE TABLE $table");
		print "Reference sequence $table was deleted.\n";
	}
#    _dbexe("TRUNCATE TABLE $TABLE");
#    print "Reference sequence was deleted.\n";
    _exit();
}

###########################################################
# Checking ref-seq already installed or not
###########################################################
if(defined _check_seq_table()){
    print "Reference sequence is already installed.\n";
    print "If you want to install again, \nplease run this tool again with \"-r\" option for restore reference sequence table.\n";
    _exit(1);
}

###########################################################
# Checking ref-seq length per line
###########################################################
if(!defined _check_fasta_length()){
    print "This file has too many nucleotide characters per line. \n";
    print "Set a line break character to the file every 200 characters.\n";
    _exit(1);
}

###########################################################
# Insert fasta
###########################################################
_print_time("*[fasta]  begin");
#my (@chrs) = _chrs();
if($RDBMS eq "oracle"){
    _sqlloder(@chrs);
}else{
#    _createSysFasta(\@chrs);
    _createQry(@chrs);

    my @tables = _getUniTables();
    my $tablesSql = join(" WRITE,", @tables);
    $tablesSql .= " WRITE";
    my $lockSql = "LOCK TABLES $tablesSql";
    print "$lockSql\n";
    _dbexe($lockSql);
    _updateRef(\@chrs);
    _dbexe("UNLOCK TABLES");

=head
    _dbexe("LOCK TABLES $TABLE WRITE;");
    _updateRef(\@chrs);
    _dbexe("UNLOCK TABLES;");
=cut
}
_print_time(" [fasta]  complete");
_exit();

###########################################################
# Functions
###########################################################
sub _getFileHandle {
    my @reffile = split(/\./,$FILE);
    my $ext = $reffile[$#reffile];
    my $fh;
    if ($ext eq "gz") {
        open $fh, "gzip -dc $FILE |" or die "file open error: $!";
    } elsif ($ext eq "bz2") {
        open $fh, "bzip2 -dc $FILE |" or die "file open error: $!";
    } else {
        open $fh, '<', $FILE or die "file open error: $!";
    }
    return($fh);
}
=head
sub setTable2Chr {
	my($chrs) = @_;
	my @c = @{$chrs};
	if(!$CHR_DIVIDE_UNIT || @c <= $CHR_DIVIDE_UNIT) {
		my $s = "";
		for ( my $i=0; $i<=$#c; $i++) {
			$CHR2TABLE{$c[$i]} = $TABLE;
			$TABLE2CHR{$TABLE} = [] if(!exists($TABLE2CHR{$TABLE}));
			push @{$TABLE2CHR{$TABLE}}, $c[$i];
		}
	} else {
		my $serial = 1;
		for ( my $i=0; $i<=$#c; $i++) {
			my $tablename = $TABLE. $serial;

			$CHR2TABLE{$c[$i]} = $tablename;
			$TABLE2CHR{$tablename} = [] if(!exists($TABLE2CHR{$tablename}));
			push @{$TABLE2CHR{$tablename}}, $c[$i];

			if( ($i != 0 && (!(($i+1)%$CHR_DIVIDE_UNIT) || $i == $#c)) || $CHR_DIVIDE_UNIT == 1 ) {
				$serial++;
			}
		}
	}
}
sub _createSysFasta {
	my $baseCreateSql = "CREATE TABLE %s (pos INT PRIMARY KEY AUTO_INCREMENT,%s) ENGINE MyISAM";
	my @tables = _getUniTables();
	foreach my $table (@tables) {
		my $s = "";
		foreach my $chr (@{$TABLE2CHR{$table}}) {
			my $colname = "`". $chr. "_f`";
			$s .= "$colname CHAR(1),";
		}
		$s =~ s/,$//;
		my $sql = sprintf($baseCreateSql, $TABLE, $s);
		_dbexe($sql);
	}
}
=cut
sub _updateRef {
    my ($refchrs)= @_;
    my (@c) = @$refchrs;
    my $POS=0;
    my $CHR="";
    my $GO=0;

    $| = 1;	# AutoFlush:On

    my $refFh = _getFileHandle();
    while(<$refFh>) {
        chomp($_);
        next if($_ !~ m/\S/);
        if ($_ =~ /^>/i){
            my @chrline;
            $CHR=$_;
            $CHR=~s/>//g;
            $CHR =~ s/^\s*(.*?)\s*$/$1/;
            @chrline = split(/\s+/, $CHR);
            $CHR=$chrline[0];
            $POS=0;
            $GO=0;
            for(my $i=0; $i<=$#c; $i++) {
                if ($c[$i] eq $CHR ) {
                    $GO=1;
                    print $CHR."\n";
                }
            }
        } elsif ($GO) {
            my $line=$_;
            $line =~ s/^\s*(.*?)\s*$/$1/;
            my @p;
            my $colname = $CHR. "_f";
            my $tablename = $CHR2TABLE{$colname};
            my $query = sprintf( "UPDATE %s SET `%s`=case pos ", $tablename, $colname );
#            my $query="UPDATE $TABLE SET `". $CHR. "_f`=case pos ";
            for (my $j=0; $j < length($line); $j++){
                $POS++;
                push @p, "$POS";
                my $FASTA = substr($line, $j, 1);
                $query .= "when $POS then '$FASTA' ";
            }
            $query .= "end WHERE pos in(";
            for (my $j=0; $j < length($line); $j++){
                $query .= $p[$j];
                $query .= "," if $j < length($line)-1;
            }
            $query.=")";
            _dbexe($query);
        }
    }
    close ($refFh);

    $| = 0;	# AutoFlush:Off
}
sub _getUniTables {
    my @tables = keys(%TABLE2CHR);
    @tables = sort {$a cmp $b} @tables;
    return @tables;
}
sub _createQry {
    my @tables = _getUniTables();
    my $m = _max();
    print "Max sequence length: $m\n";
    my $str="";
    if($RDBMS eq "oracle"){
        #for (my $i=1; $i<=$m; $i++) {
        #    if ( $i == $m || !($i%2000) ) {
        #        $str.=" into $TABLE(pos) values($i)";
        #        _dbexe("INSERT all $str select * from dual");
        #        $str="";
        #    } else {
        #        $str.=" into $TABLE(pos) values($i)";
        #    }
        #}
    }else{
      foreach my $tablename (@tables) {
        print "Initializing table: $tablename\n";
        for (my $i=1; $i<=$m; $i++) {
            if ( $i == $m || !($i%2000) ) {
                $str.="($i)";
                _dbexe("INSERT into $tablename(pos) values $str");
                $str="";
            } else {
                $str.="($i),";
            }
        }
      }
    }
} 

### Inserting for oracle ################################################
sub _sqlloder{
    my $m = _max();
    my $crtdir = Cwd::getcwd();
    my @chrs;
    my $tempdir = $FILE.".temp";
    my $temp_base_name = basename($FILE);
    
    _mkTmpDir($tempdir);
    
    open IN, '<', $FILE or die "file open error: $!";
    my $reading = 0;
    my $written = 0;
    while (<IN>) {
        chomp;
        if($_ =~ /^>/){
            my @line = split(/>/, $_);
            if($reading == 1){
                if($written < $m){
                    for (my $b=$written+1; $b<=$m; $b++) {
                        print OUT "\\N\n";
                    }
                }
                close (OUT);
                open OUT, '>', $tempdir."/".$temp_base_name.".".$line[1] or die "file open error: $!";
                push(@chrs,$tempdir."/".$temp_base_name.".".$line[1]);
                $reading = 0;
                $written = 0;
            }else{
                open OUT, '>', $tempdir."/".$temp_base_name.".".$line[1] or die "file open error: $!";
                push(@chrs,$tempdir."/".$temp_base_name.".".$line[1]);
                $reading = 1;
            }
        }else{
            my $len = length($_);
            my @letters = split("", $_);
            for (my $a=0; $a<$len; $a++) {
                print OUT $letters[$a]."\n";
            }
            $written += $len;
        }
        if(eof IN){
            if($written < $m){
                for (my $b=$written+1; $b<=$m; $b++) {
                    print OUT "\\N\n";
                }
            }
            close (OUT);
        }
    }
    close (IN);
    
    my $control_file = $tempdir."/".$temp_base_name.".ctl";
    my $tsv = $tempdir."/".$temp_base_name.".tsv";
    my $logfile = $tempdir."/log.log";
    my $badfile = $tempdir."/bad.bad";
    my $num_chrs = @chrs;
    my $ctlstr = "";
    $ctlstr .= "OPTIONS( DIRECT = TRUE, MULTITHREADING = TRUE) ";
    $ctlstr .= "LOAD DATA INFILE '$tsv' BADFILE '$tempdir\/bad' DISCARDFILE '$tempdir\/disc' ";
    $ctlstr .= "APPEND INTO TABLE sys_fasta FIELDS TERMINATED BY X'09' (pos, ";
    for(my $c=0; $c<$num_chrs; $c++){
        my @chr_name = split(/\./,$chrs[$c]);
        if($c == $num_chrs-1 ){
            $ctlstr .= $chr_name[$#chr_name]."_f NULLIF ".$chr_name[$#chr_name]."_f=\'\\\\N\'";
        }else{
            $ctlstr .= $chr_name[$#chr_name]."_f NULLIF ".$chr_name[$#chr_name]."_f=\'\\\\N\',";
        }
    }
    $ctlstr .= ")";
    
    open OUTCTL, '>', $control_file or die "file open error: $!";
    print OUTCTL $ctlstr;
    close (OUTCTL);
    
    open OUT, '>', $tempdir."/".$temp_base_name.".pos" or die "file open error: $!";
    for (my $t=1; $t<=$m; $t++) {
        print OUT $t."\n";
    }
    close (OUT);
    unshift(@chrs,$tempdir."/".$temp_base_name.".pos");
    
    my $param1= join("\t", @chrs);
    my $cmd = "paste $param1 > $tsv";
    if (system($cmd)){
        print  "Error: Merging sequence files.\n";
        _exit();
    }
    my $cmd2 = "sqlldr userid=\"$USER/$PASS\@$DB\" control=\'$control_file\' log='$tempdir\/log' errors=0";
    push(@chrs,$logfile);
    push(@chrs,$badfile);
    push(@chrs,$control_file);
    push(@chrs,$tsv);
    if (system($cmd2)){
        print  "Error: Loading the reference sequence is failed.\n";
        _exit();
    }
    clean_dir($crtdir, \@chrs, $crtdir);
    rmdir($tempdir);
}

### temp directory ################################################
sub _mkTmpDir{
    my $OUTDIR = shift;
    if(!-e "$OUTDIR"){
        mkdir("$OUTDIR") or die "$!: $OUTDIR";
    } else{
        print "$OUTDIR already exists.\n";
        exit(0);
    }
}

### Removing temp files ###############################
sub clean_dir{
    my $work_dir = shift;
    my $chr_sc = shift;
    my $CD = shift;
    my @files= @$chr_sc;
    chdir("$work_dir") or die "$! : $work_dir";
    foreach my $file(@files){    
        if(-e $file ){
            unlink $file;
        }
    }
    chdir($CD) or die "$! : $CD";
}

### Set chromosome ################################################
sub _chrs {
    my $val="";
    my @c;
    my $query = "SELECT id FROM sys_chromosome";
    my $query_handle = $connect->prepare($query);
    $query_handle->execute();
    $query_handle->bind_columns(undef, \$val);
    while($query_handle->fetch()) {
        push @c, $val;
    }
    return @c;
}

### Set max of chromosome ################################################
sub _max {
    my $val;
    my $query;
    if($RDBMS eq "oracle"){
        $query = "SELECT max(seq_size) FROM sys_chromosome";
    }else{
        $query = "SELECT max(size) FROM sys_chromosome";
    }
    my $query_handle = $connect->prepare($query);
    $query_handle->execute();
    $val = $query_handle->fetchrow_hashref;
    if($RDBMS eq "oracle"){
        return $val->{uc('max(seq_size)')};
    }else{
        return $val->{'max(size)'};
    }
}

### Execute DB ################################################
sub _dbexe {
    my ($query) = @_;
    my $query_handle = $connect->prepare($query);
#    $query_handle->execute() or warn "Excute error: $DBI::errstr\n";
    $query_handle->execute() or die "Excute error: $DBI::errstr\n";
}

#### Usage #######################################
sub _error_method {
    print "\nCreating reference sequence track\n";
    print "./tasuke_ref.pl -f [.fasta] -db [database] -u [user] -p [password]\n( For oracle database: -b oracle -s [tablespace])\n\n";
    print "Restoring track\n";
    print "./tasuke_ref.pl -r -db [database] -u [user] -p [password]\n\n";
}

#### Time #######################################
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";
}

### Checking table ####################################
sub _check_exists_table {
    my ($t) = @_;
    my $val;
    if($RDBMS eq "oracle"){
        $t = uc($t);
        $val = $connect->selectrow_array("select table_name from user_tables where tablespace_name = 'TSK_SPACE' and table_name = '$t'");
    }else{
        $val = $connect->selectrow_array("show tables like '$t'");
    }
    return $val;
}

### Checking tablespace ####################################
sub _check_exists_tablespace {
    my ($t) = @_;
    $t = uc($t);
    my $val = $connect->selectrow_array("select tablespace_name from user_tablespaces where tablespace_name = '$t'");
    return $val;
}

### Installed or not ####################################
sub _check_seq_table {
    my $val;
    my $tbl = 'sys_fasta';
    if($RDBMS eq "oracle"){
        $val = $connect->selectrow_array("select table_name from user_tables where tablespace_name = 'TSK_SPACE' and table_name = 'SYS_FASTA'");
    }else{
	my @tables = keys(%TABLE2CHR);
	if(@tables > 0) {
		$tbl = $tables[0];
		$val = $connect->selectrow_array("show tables like '". $tbl. "'");
	}
#        $val = $connect->selectrow_array("show tables like '". $tables[0]. "'") if(@tables > 0);
    }
    my $res = $connect->selectrow_array("select * from $tbl where pos = 1") if defined $val;
    return $res;
}

### Checking number of base per line in the reference fasta file  ####################################
sub _check_fasta_length{
    my $ret;
    my $fh = _getFileHandle();
#    my @refs = _readRef();
#    for (@refs) {
    while(<$fh>) {
        chomp;
        if ($_ !~ /^>/i){
            if(length($_) <= 200){
                $ret = 0;
            }
            last;
        }
    }
    close($fh);
    return $ret;
}

#### 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);
}

