#!/usr/bin/perl

use strict;
use warnings;

use lib qw(myClasses);

use PDBParser;

###########################################################
#
# The script runs FlexE  
#
###########################################################

my $RSYNC = 1;  # the parameter-switcher whether rsync data from/to $SERVER
		# 1 - for normal functioning of the program
		# 0 - for cases when network connection was too slow

my $CASP = "CASP13";
my $TARGETS_DIR = "/local/$CASP/TARGETS/";
my $MODELS_DIR  = "/local/$CASP/MODELS/";
my $RESULTS_DIR = "/local/$CASP/RESULTS/FlexE/";
my $TMP_DIR = "/tmp/FlexE_.$$/";
my $SERVER = "predictioncenter.org";

my $program = "python /usr/local/FlexE/run_FlexE.py ";

my @TARGETS = ();

&rsyncTargets($TARGETS_DIR);
if ( defined($ARGV[0]) ){
   push @TARGETS, $ARGV[0];
} else {
   &parseTargetsDir();
}


foreach my $target (@TARGETS){
  #print $target."\n"; next;  
  if (!-f "$TARGETS_DIR/$target.pdb"){
     warn("File $TARGETS_DIR/$target.pdb doesn't exist\n");
     next;
  }
  my @models = &getModels($target);
  if (! -d "$MODELS_DIR/$target" ){
     warn("Directory $MODELS_DIR/$target doesn't exist");
     next;
  } 
  #print "$target\n"; next;
  if (-d "$RESULTS_DIR/$target"){
     system("rm -rf $RESULTS_DIR/$target");
  }
  system("mkdir -p $RESULTS_DIR/$target");
  system("mkdir -p $TMP_DIR");
  foreach my $model (@models){
    my $target_path = sprintf("%s/%s.pdb", $TARGETS_DIR, $target);
    my $model_path = sprintf("%s/%s/%s", $MODELS_DIR, $target, $model);
    &adjust_pdb_structures($target_path, $model_path);
    my $command = sprintf ("%s --ref %s/%s.%s --pdb %s/%s > %s/%s/%s.flexe", $program, $TMP_DIR, $target, $model, $TMP_DIR, $model, $RESULTS_DIR, $target, $model);
    system("$command");
  #   print "$command\n";
  }

  my $resDir = sprintf("%s/%s/", $RESULTS_DIR, $target);
  &uploadResultFile($resDir);

  # run starting model for refinement target TR
  if (substr($target, 1,1) eq 'R'){
	my $start_model_path = "$TARGETS_DIR/templates/$target.pdb.txt";
	my $start_model = "$target.pdb.txt";
	my $target_path = sprintf("%s/%s.pdb", $TARGETS_DIR, $target);
	&adjust_pdb_structures($target_path, $start_model_path);
	system("mkdir -p $RESULTS_DIR/TR/$target");
	my $command = sprintf ("%s --ref %s/%s.%s --pdb %s/%s > %s/TR/%s/%s.flexe", $program, $TMP_DIR, $target, $start_model, $TMP_DIR, $start_model, $RESULTS_DIR, $target, $start_model);
	system("$command");
	my $resFileDir = sprintf("%s/TR/%s/", $RESULTS_DIR, $target, $start_model);
	&uploadResultFile($resFileDir);
  }
  
  system("rm -rf $TMP_DIR");
  sleep(2);
}

#----------------------------
# SUBROUTINES
#----------------------------

sub getModels{
    my $t = shift;
    my @models = ();
    &rsyncModels("$MODELS_DIR/$t");
    if (! -d "$MODELS_DIR/$t") {return @models;}
    if (opendir D, "$MODELS_DIR/$t"){
	while(defined(my $m = readdir(D))){
		if ($m !~ m/^T[0-9Racspx]/){next;}
		push @models, $m;
	}
        closedir D; 
        @models = sort @models;
    }
    return @models;
}

sub usage{
   my $text = shift;
   print "\nUSAGE: $0 --target <TARGET>\n\n";
   if (defined($text)){
	warn($text);
   }
   exit;
}

sub adjust_pdb_structures {
    my ($target_path, $model_path) = @_;
    my $target = $target_path; $target =~ s/.*\///g; $target =~ s/\.pdb$//;
    my $model = $model_path; $model =~ s/.*\///g; 
    # remove chains and renumerate target and model structures
    my $tmp_target = sprintf("%s/tmp.%s.pdb", $TMP_DIR, $target);
    my $tmp_model = sprintf("%s/tmp.%s", $TMP_DIR, $model);

    my $NEW_FILE;
    if (! -e $tmp_target){
       if (open($NEW_FILE, sprintf("> %s", $tmp_target))) {
           my $pdb_parser = new PDBParser();
           $pdb_parser->open($target_path);
           $pdb_parser->transform_chains_into_1000($NEW_FILE);
           close($NEW_FILE);
           $pdb_parser->close();
       }
    }
   
    if (open($NEW_FILE, sprintf("> %s", $tmp_model))) {
           my $pdb_parser = new PDBParser();
           $pdb_parser->open($model_path);
           $pdb_parser->transform_chains_into_1000($NEW_FILE);
           close($NEW_FILE);
           $pdb_parser->close();
    }
    # read target structure
    my $target_parser = new PDBParser();
    my $target_file_in = sprintf("%s",  $tmp_target);
    $target_parser->open($target_file_in);
    my %target_str = $target_parser->parseStructure();
    $target_parser->close();

    #read model structure
    my $model_parser = new PDBParser();
    my $model_file_in = sprintf("%s", $tmp_model);
    $model_parser->open($model_file_in);
    my %model_str = $model_parser->parseStructure();
    $model_parser->close();

  # write simultaneously into two files
    my $target_file_out = sprintf("%s/%s.%s", $TMP_DIR, $target, $model);
    my $model_file_out = sprintf("%s/%s", $TMP_DIR, $model);

    open FILE1, "> $target_file_out";
    open FILE2, "> $model_file_out";
    foreach my $res_no ( sort {$a<=>$b} keys %target_str) {
        if (defined($target_str{$res_no}) && defined($model_str{$res_no})){ # check if residues in target and in model with a certain residue number are defined  
                my $target_res = $target_str{$res_no};
                my $model_res = $model_str{$res_no};
                # arrays of atoms
                my @target_atoms = @{$target_res->{_atoms}} if defined $target_res->{_atoms};
                my @model_atoms = @{$model_res->{_atoms}} if defined $model_res->{_atoms};
                        foreach my $atom (@target_atoms){
                                if ($atom->{_name} eq 'N' || $atom->{_name} eq 'C' || $atom->{_name} eq 'CA' || $atom->{_name} eq 'O' || $atom->{_name} eq 'CB'){
                                        print FILE1 $atom->{_line};
                                }
                        }
                        foreach my $atom (@model_atoms){
                                if ($atom->{_name} eq 'N' || $atom->{_name} eq 'C' || $atom->{_name} eq 'CA' || $atom->{_name} eq 'O' || $atom->{_name} eq 'CB'){
                                        print FILE2 $atom->{_line};
                                }
                        }
         }
    }
    close FILE1;
    close FILE2;
}

sub uploadResultFile{
    my ($fullPathResFile) = shift;
    unless ($RSYNC) {return ;} # if $RSYNC = 0 then do nothing
    my $command = sprintf("rsync -az -e ssh $fullPathResFile  $SERVER:/$fullPathResFile ");
    system("$command");
#    print("$command\n");

}


sub rsyncModels{
    my ($fullPathModelDir) = @_;
    unless ($RSYNC) {return ;} # if $RSYNC = 0 then do nothing
    my $command = sprintf("rsync -az -e ssh $SERVER:/%s/ /%s/ --delete ", $fullPathModelDir, $fullPathModelDir);
    system("$command");
#    print("$command\n");
}

sub rsyncTargets{
    my ($fullPathTargetDir) = @_;
    unless ($RSYNC) {return ;} # if $RSYNC = 0 then do nothing
    my $command = sprintf("rsync -az -e ssh $SERVER:/%s/ /%s/ --delete", $fullPathTargetDir, $fullPathTargetDir);
    system("$command");
    $command = sprintf("rsync -az -e ssh $SERVER:/%s/templates/ /%s/templates/ --delete", $fullPathTargetDir, $fullPathTargetDir);
    system("$command");
    #print("$command\n");
}


sub parseTargetsDir {
    opendir D, $TARGETS_DIR;
    while(defined(my $f = readdir(D))){
        next unless $f =~ m/^(T|R).*.pdb$/;
        # skip targets -D1 that have been copied from -D0
        if (($f =~ m/-D1\.pdb/) && (-l "$TARGETS_DIR/$f")) {next;}
        # skip combined domains -D12  
        if ($f =~ m/-D[1-9]{2}/) {next;}
        # skip test domain
        if ($f =~ m/-D[7-9]{1}/) {next;} # skip test domains
        $f =~ s/\.pdb$//;
        if (&checkIfResultsExist($f)){
                push @TARGETS, $f;
        }
    }
    closedir D;
}

sub checkIfResultsExist {
    my $target = shift;
    my $targetFile = sprintf("%s/%s.pdb", $TARGETS_DIR, $target);
    my $resultDir = sprintf("%s/%s/", $RESULTS_DIR, $target);
    if (! -e $resultDir) {
        return 1;
    }
    # check if target file has been modified AFTER result dir has been created
    if ((stat($targetFile))[9] > (stat($resultDir))[9]){
        return 1;
    }
    return 0;
}