#!/usr/bin/perl -w

#
# tasuke_tsv_db.pl
# http://tasuke.dna.affrc.go.jp/
#
# Author: Ryutaro Itoh
#
# Copyright (c) 2013 National Institute of Agrobiological Sciences   
#
# This software is released under the MIT License.
#
# http://opensource.org/licenses/mit-license.php
#

use strict;
use DBI;
use Getopt::Long;
use File::Basename;
use Cwd;

my $NAME   = '';
my $RESTORE= 0;
my $LIST   = 0;
my $VCF    = '';
my $DB     = '';
my $HOST   = 'localhost';
my $USER   = '';
my $PASS   = '';
my $READS = 0;
my $CUSTOMTRACK = 0;
my $RDBMS = '';
my $TABLESPACE = '';
my $CONDITIONID = '';

################################################################
# Checking parameters and database connection
################################################################
GetOptions( 'n|name=s'      => \$NAME,
            'r|restore'     => \$RESTORE,
            'l|list'        => \$LIST,
            'f|file=s'      => \$VCF,
            'c|custom'      => \$CUSTOMTRACK,
            'i|id=s'      => \$CONDITIONID,
            '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 eq ''||($RESTORE&&!$NAME)||(!$RESTORE&&!-e$VCF)) {
  _error_method();  
  exit;
}

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(1);
    }
}else{
    $connect= DBI->connect('DBI:mysql:database='.$DB.';host='.$HOST,$USER,$PASS)
    or die "Can't connect to database: $DBI::errstr\n";
}

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

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

###########################################################
# Checking accession id
###########################################################
my $TABLE =_check_db();
if ( ! defined $TABLE ) {
    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);
}
if($CUSTOMTRACK){
    my $index = _get_condtable($CONDITIONID);
    if(! defined $index &&  ! defined $CONDITIONID){
        print "Condition:$CONDITIONID was not found.\n";
        _exit();
    }
    if(defined $index){
        $TABLE.="_accession_cstm_".$index;
    }else{
        $TABLE.="_accession_cstm";
    }
}else{
    $TABLE.="_accession";
}


###########################################################
# Restoring depth table
###########################################################
if ($RESTORE) {
    _dbexe("TRUNCATE TABLE $TABLE");
    if($CUSTOMTRACK){
        print "Custom track $NAME table is restored.\n";
    }else{
        print "DEPTH $NAME table is restored.\n";
    }
    _exit();
}

###########################################################
#  Checking DEPTH file
###########################################################
 my $num=0;
 _print_time("File checking");
 my $ext = _checkExt();
 if ($ext eq "gz") {
     open IN, "gzip -dc $VCF |" or die "file open error: $!";
 } elsif ($ext eq "bz2") {
     open IN, "bzip2 -dc $VCF |" or die "file open error: $!";
 } else {
     open IN, '<', $VCF or die "file open error: $!";
 }
 #open IN, '<', $VCF or die "file open error: $!";
 while (<IN>) {
    if($num==0){
        $num++;
        my @line = split(/\t/, $_);
        my @column = split(/=/, $line[0]);
        if($CUSTOMTRACK){
            if(!defined $column[1]){
                print " Tasuke_depth_db.pl can not get number of reads from this TSV file.\n Try to make TSV file using TASUKE release 1.3.0 or later.\n Exit. \n";
                _exit(1);
            }
            if($column[1] =~ /\d+/){
                $READS=$column[1];
            }else{
                $READS=0;
            }
        }
    }else{
        $num++;
        my @line = split(/\t/, $_);
        foreach my $x(@line){
            if($x eq ""){
                print "The TSV file is probably truncated. near:$num \n";
                print $_."\n";
                print "Exit.\n";
                _exit(1);
            }
        }
    }
 }
 $num--;
 my $chk_max = _max();
 if($chk_max != $num){
    print "The TSV file is probably truncated. near:$num\n";
    print "Exit.\n";
    _exit(1);
 }
 
###########################################################
# Creating depth table
###########################################################
# Set chromosomes 
my (@chrs) = _chrs();
my ($max) = _max();
  
### Adding a column for number of reads when you use old table(~1.2.0) ########
if($RDBMS ne "oracle"){
    if(!defined _check_exists_cstm()){
        _dbexe("ALTER TABLE sys_accession add column num_reads bigint DEFAULT 0");
    }
    if($CUSTOMTRACK){
        my @tables = _get_tables();
        foreach my $x(@tables){
            if (!defined _check_exists_table($x."_cstm")){
                _dbexe("CREATE TABLE ".$x."_cstm like ".$x);
            }
        }
    }
}

### TSV LOADER ###
_print_time(" [".$NAME." tsv] begin create table");
if($RDBMS eq "oracle"){
    _LoadDepthOracle($VCF, $TABLE, $NAME, $READS);
}else{
    _dbexe("LOCK TABLES $TABLE WRITE, sys_accession WRITE, sys_chromosome READ;");
    _LoadDepth($VCF, $TABLE, $NAME, $READS);
    _dbexe("UNLOCK TABLES;"); 
}
_print_time(" [".$NAME." tsv] complete create table");
_exit();

### check vcf extention ##########################################
sub _checkExt {
    my @reffile = split(/\./,$VCF);
    my $ext = $reffile[$#reffile];
    return($ext);
}
            
### 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;
}

### Checking for number of reads column ####################################
sub _check_exists_cstm {
    my $val = $connect->selectrow_array("show columns from sys_accession where Field = 'num_reads'");
    return $val;
}

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

sub _error_method {
    print "\nCreating depth table\n";
    print "tasuke_tsv_db.pl -n [id] -f [tsv file] -db [database] -u [user] -p [password] \n( For oracle database: -b oracle -s [tablespace])\n\n";
    print "Creating general-purpose table\n";
    print "tasuke_tsv_db.pl -c -n [id] -f [tsv file] -db [database] -u [user] -p [password]\n\n";
    print "Restoring table\n";
    print "tasuke_tsv_db.pl -r -n [name] -db [database] -u [user] -p [password]\n\n";
    print "Listing accession\n";
    print "tasuke_tsv_db.pl -l -db [database] -u [user] -p [password]\n\n";
}

######## LOAD DEPTH DATA ########
sub _LoadDepth{
    my $DEPTH = shift;
    my $TABLE = shift;
    my $NAME = shift;
    my $READS = shift;
    if($CUSTOMTRACK){
        _dbexe("UPDATE sys_accession SET num_reads=".$READS." where accession_id='".$NAME."'");
    }
    my $ext = _checkExt();
    if ($ext eq "gz" or $ext eq "bz2") {
	zipImport();
=pod
	my $dbquery;
        open IN, "gzip -dc $VCF |" or die "file open error: $!";
        my $query = "select id from sys_chromosome";
        my $query_handle = $connect->prepare($query);
        $query_handle->execute();
        my $chrs = $query_handle->fetchall_arrayref;
        $dbquery = "INSERT INTO $TABLE(pos";
        for my $chr (@chrs) {
            $dbquery .= ",".$chr."_b,".$chr."_k";
        }
        $dbquery .= ")";
        $dbquery .= " values";

	while(<IN>) {
	    my $line = $_;
	    $line =~ s/[\r\n]//g;
	    my @lines = split(/\t/,$_);
	    $dbquery .= "(";
	    my $count = 0;
	    for my $data (@lines) {
		my $mod = $count % 2;
		$data =~ s/[\r\n]//g;
		if ($mod eq 0) {
			if ($data eq '\N') {
				$data = "NULL";
			}
		}
		
		if ($count < $#lines) {
			if($data eq 'NULL') {
				$dbquery .= $data.",";
				$count++;
			} else {
				$dbquery .= "'".$data."',";
				$count++;
			}
		} else {
			if($data eq 'NULL') {
                                $dbquery .= $data."),";
			} else {
				$dbquery .= "'".$data."'),";
			}
		}
	    }
	}
	$dbquery =~ s/,$//g;
	$dbquery .= ";";
        _dbexe($dbquery);
=cut
    } else {
    	_dbexe("LOAD DATA LOCAL INFILE "."\'".$DEPTH."\'"." replace INTO TABLE ".$TABLE." IGNORE 1 LINES ;");
    }
}

sub zipImport {
        my $dbquery;
        my @reffile = split(/\./,$VCF);
        my $ext = $reffile[$#reffile];
	if ($ext eq "gz") {
        	open IN, "gzip -dc $VCF |" or die "file open error: $!";
	} elsif ($ext eq "bz2") {
		open IN, "bzip2 -dc $VCF |" or die "file open error: $!";
	}
        my $query = "select id from sys_chromosome";
        my $query_handle = $connect->prepare($query);
        $query_handle->execute();
        my $chrs = $query_handle->fetchall_arrayref;
        $dbquery = "INSERT INTO $TABLE(pos";
        for my $chr (@chrs) {
            $dbquery .= ",".$chr."_b,".$chr."_k";
        }
        $dbquery .= ")";
        $dbquery .= " values";

        while(<IN>) {
            my $line = $_;
            $line =~ s/[\r\n]//g;
            my @lines = split(/\t/,$_);
            $dbquery .= "(";
            my $count = 0;
            for my $data (@lines) {
                my $mod = $count % 2;
                $data =~ s/[\r\n]//g;
                if ($mod eq 0) {
                        if ($data eq '\N') {
                                $data = "NULL";
                        }
                }

                if ($count < $#lines) {
                        if($data eq 'NULL') {
                                $dbquery .= $data.",";
                                $count++;
                        } else {
                                $dbquery .= "'".$data."',";
                                $count++;
                        }
                } else {
                        if($data eq 'NULL') {
                                $dbquery .= $data."),";
                        } else {
                                $dbquery .= "'".$data."'),";
                        }
                }
            }
        }
        $dbquery =~ s/,$//g;
        $dbquery .= ";";
        _dbexe($dbquery);
}	
######## LOAD DEPTH DATA FOR ORACLE ########
sub _LoadDepthOracle{
    my $DEPTH = shift;
    my $TABLE = shift;
    my $NAME = shift;
    my $READS = shift;
    my $crtdir = Cwd::getcwd();
    my $tempdir = $DEPTH.".temp";
    my $temp_base_name = basename($DEPTH);
    
    if($CUSTOMTRACK){
        _dbexe("UPDATE sys_accession SET num_reads=".$READS." where accession_id='".$NAME."'");
    }
    _mkTmpDir($tempdir);
    
    my $control_file = $tempdir."/".$temp_base_name.".ctl";
    my $logfile = $tempdir."/log.log";
    my $badfile = $tempdir."/bad.bad";
    my $num_chrs = @chrs;
    my $ctlstr = "";
    my @files_for_SQLLoader;
    
    $ctlstr .= "OPTIONS( DIRECT = TRUE, MULTITHREADING = TRUE, SKIP = 1) ";
    $ctlstr .= "LOAD DATA INFILE '$DEPTH' BADFILE '$tempdir\/bad' DISCARDFILE '$tempdir\/disc' ";
    $ctlstr .= "APPEND INTO TABLE $TABLE 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]."_b NULLIF ".$chr_name[$#chr_name]."_b=\'\\\\N\',";
            $ctlstr .= $chr_name[$#chr_name]."_k NULLIF ".$chr_name[$#chr_name]."_k=\'\\\\N\'";
        }else{
            $ctlstr .= $chr_name[$#chr_name]."_b NULLIF ".$chr_name[$#chr_name]."_b=\'\\\\N\',";
            $ctlstr .= $chr_name[$#chr_name]."_k NULLIF ".$chr_name[$#chr_name]."_k=\'\\\\N\',";
        }
    }
    $ctlstr .= ")";
    
    open OUTCTL, '>', $control_file or die "file open error: $!";
    print OUTCTL $ctlstr;
    close (OUTCTL);
    
    my $cmd2 = "sqlldr userid=\"$USER/$PASS\@$DB\" control=\'$control_file\' log='$tempdir\/log' errors=0";
    if (system($cmd2)){
        print  "Error: Loading the TSV file is failed.\n";
        _exit(1);
    }
    push(@files_for_SQLLoader,$logfile);
    push(@files_for_SQLLoader,$badfile);
    push(@files_for_SQLLoader,$control_file);
    clean_dir($crtdir, \@files_for_SQLLoader, $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";
}

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

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

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

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

### 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 = '$TABLESPACE' 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;
}

sub _get_tables{
    my $val="";
    my @id;
    my $query = "show tables like 'd%_accession'";
    my $query_handle = $connect->prepare($query);
    $query_handle->execute();
    $query_handle->bind_columns(undef, \$val);
    while($query_handle->fetch()) {
        push @id, $val;
    }
    return @id;
}

sub _get_condtable{
    my ($t) = @_;
    my $val;
    if($RDBMS eq "oracle"){
        $t = uc($t);
        $val = $connect->selectrow_array("select table_id from sys_condition where condition_id = '$t' and ROWNUM < 2");
    }else{
        $val = $connect->selectrow_array("select table_id from sys_condition where condition_id = '$t' limit 1");
    }
    return $val;
}

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