#! /usr/bin/perl -w

# Script lga_out_2_newGDTTS.pl reads LGA output (CASP style) and calculates a modified GDTTS measure, 
# penalizing model CAs that are spatially close but do not correspond to target CA:
# -1 for D< 1A; -0.5 for D<2A; -0.25 for D<4A.

# LGA run (in directory LGA_package): ./runlga.mol_mol.pl <model> <target>  -3  -ie  -o1  -sda  -d:<distance:1,2,4, or 8>

if($#ARGV<5) { die "Usage: lga_out_2_newGDTTS_uniformcuts_noneighb_1.pl -m <model pdb file> -t <target pdb file> -w <penalty weight> -o <output file with GDT_TS, new measure, and total number of compared residues> -rundir <running dir>\n"; } 

# my $lgadir = "/home/syshi_3/casp/LGA_package";
my @lgads = ("1.0", "2.0", "4.0", "8.0"); # used distances for LGA 
# my %dcuts = ("1.0" => "-1.0", "2.0" => "-2.0", "4.0" => "3.5", "8.0" => "3.5"); # distance cutoffs for penalties applied to different LGA superpositions
my %pen = ("1.0" => 1.0, "2.0" => 1.0, "4.0" => 1.0); # penalties for having non-corresponding residues within distance cutoffs

my $weight = 1.0;
my $rundir = "/home/syshi_3/casp/test";
while (@ARGV and $ARGV[0] =~ /^-/) {
        $_ = shift;
        last if /^--$/;
#        if (/^-lgadir/) {$_ = shift; $lgadir = $_;}
        if (/^-m/) {$_ = shift; $model = $_;}
        if (/^-t/) {$_ = shift; $target = $_;}
        if (/^-w/) {$_ = shift; $weight = $_;}  # penalty weight
        if (/^-o/) {$_ = shift; $outfile = $_;}
        if (/^-rundir/) {$_ = shift; $rundir = $_;}
}


# read CA coordinates (x1,y1,z1){$i} for model
my %x1=();
my %y1=();
my %z1=();
my $nowmodel=$rundir.$model;
open(MOD, $nowmodel) || die "Cannot open input pdb file with model $nowmodel : $!\n";
while(my $line=<MOD>) {
	chomp($line);
	if($line=~/^ATOM/){
		
	}else{
		next;
	}		
	my $type=substr($line,12,4);

	if($type eq " CA "){ # pdb line format, e.g: ATOM      2  CA  MET     1     489.230 945.691 891.975  1.00  5.00
		my $xpos=substr($line,30,8);
		if($xpos=~/(-?\d+\.?\d+)/){
				$xpos=$1;
		}	
		my $ypos=substr($line,38,8);
		if($ypos=~/(-?\d+\.?\d+)/){
				$ypos=$1;
		}	
		my $zpos=substr($line,46,8);
		if($zpos=~/(-?\d+\.?\d+)/){
				$zpos=$1;
		}	
		
		my $n=substr($line,22,5);
		if($n=~/(-?\w+)/){
			$n=$1;
		}	
		$x1{$n} = $xpos; 
		$y1{$n} = $ypos; 
		$z1{$n} = $zpos; 
	}			
}
close(MOD); 

# read CA coordinates (x2,y2,z2){$i} for target
my %x2=();
my %y2=();
my %z2=();
my $nowtarget=$rundir.$target;
open(TARG, $nowtarget) || die "Cannot open input pdb file with target $nowtarget : $!\n";
while(my $line=<TARG>) {
	chomp($line);
	if($line=~/^ATOM/){
		
	}else{
		next;
	}	
	my $type=substr($line,12,4);
	if($type eq " CA "){ # pdb line format, e.g: ATOM      2  CA  MET     1     489.230 945.691 891.975  1.00  5.00
		my $xpos=substr($line,30,8);
		if($xpos=~/(-?\d+\.?\d+)/){
				$xpos=$1;
		}	
		my $ypos=substr($line,38,8);
		if($ypos=~/(-?\d+\.?\d+)/){
				$ypos=$1;
		}	
		my $zpos=substr($line,46,8);
		if($zpos=~/(-?\d+\.?\d+)/){
				$zpos=$1;
		}	
		my $n=substr($line,22,5);
		if($n=~/(-?\w+)/){
			$n=$1;
		}	
		$x2{$n} = $xpos; 
		$y2{$n} = $ypos; 
		$z2{$n} = $zpos; 

	}			
}
close(TARG); 



#chomp($rundir);
my $outlga = "$model\.$target\.res";


my %res_gc = ();

my $checkgdtts=-1.00;
my $firstflag=0;

my %res_penalty_mod=();
my %res_penalty_targ=();

for $lgad(@lgads) {

	# Run lga on model and target
#	my $cmd = "./runlga.mol_mol.pl $rundir"."$model $rundir"."$target  -3  -ie  -o1  -sda  -d\:$lgad";
# 
#	print STDERR "running LGA in $lgadir: $cmd\n";
#	
#	system("cd $lgadir; $cmd");

	$lga_file{$lgad} = $outlga;
	#print $outlga."\t".$rundir."\n";
	$lga_file{$lgad} =~ s/\.res$/\.d$lgad\.res/;
	#print $lga_file{$lgad}."\n";
	$lga_file{$lgad} =$rundir.$lga_file{$lgad};
#print $lga_file{$lgad}."\n";
#exit;
	
#	system("mv $lgadir\/RESULTS\/$outlga $lga_file{$lgad}");
	#print "$lga_file{$lgad}\n";

	# read LGA file; get (i) residue correspondence; (ii) GDT_TS; (iii) transformation coefficients for 1=>2;  
	open(FLGA, $lga_file{$lgad}) || die "Cannot open lgafile $lga_file{$lgad} : $!\n"; 

	my $rmsd = -1.0;
	my $gdtts = -1.0;
	my @res1 = ();
	my @res2 = ();
	my @cx = ();
	my @cy = ();
	my @cz = ();
	my $n_lcslines = 0;
	while(<FLGA>) {
		if(/^LCS_GDT /) { # lines for corresponding residue pairs
			next if ++$n_lcslines < 4; # header
			@fs = split;
			push(@res1, $fs[2]);
			push(@res2, $fs[4]);
		} elsif( /^SUMMARY\(GDT\) / ) { # rmsd, gdt_ts, etc.
			@fs = split;
			$rmsd = $fs[5];
			$gdtts = $fs[6];
			if($firstflag==0){
				$firstflag=1;
				$checkgdtts=$gdtts;
			}else{
				if($gdtts!=$checkgdtts){
					print "warning $model $target $gdtts $checkgdtts not identical\n";
					#exit;
				}	
			}		

		# coefficients to transform pdb1 (model) into pdb2 (target):
		} elsif(/^  X_new \=  (.*?) \* X  \+  (.*?) \* Y  \+  (.*?) \* Z  \+ (.*?)$/) {
			$cx[0]=$1; $cx[1]=$2; $cx[2]=$3; $cx[3]=$4;
		} elsif(/^  Y_new \=  (.*?) \* X  \+  (.*?) \* Y  \+  (.*?) \* Z  \+ (.*?)$/) {
			$cy[0]=$1; $cy[1]=$2; $cy[2]=$3; $cy[3]=$4;
		} elsif(/^  Z_new \=  (.*?) \* X  \+  (.*?) \* Y  \+  (.*?) \* Z  \+ (.*?)$/) {
			$cz[0]=$1; $cz[1]=$2; $cz[2]=$3; $cz[3]=$4;
		}

	}
	close(FLGA);
	$npairs = scalar(@res1);
  #$ntpairs = scalar (@res2);
	# Transform coord of 1st structure (model) using LGA's coefficients for superposition
	#print $cx[0]."\t".$cx[1]."\t".$cx[2]."\t".$cx[3]."\n";
	#print $cy[0]."\t".$cy[1]."\t".$cy[2]."\t".$cy[3]."\n";
	#print $cz[0]."\t".$cz[1]."\t".$cz[2]."\t".$cz[3]."\n";
	
	my %xt1=();
	my %yt1=();
	my %zt1=();
	for $n(sort {$a<=>$b} keys %x1) {
		my $transflag=1;
		if(defined $cx[0] and defined $cx[1] and defined $cx[2] and defined $cx[3] and defined $cy[0] and defined $cy[1] and defined $cy[2] and defined $cy[3] and defined $cz[0] and defined $cz[1] and defined $cz[2] and defined $cz[3]){
			
		}else{
			$transflag=0;
			#print "transform not defined good: $model\t$target\n";
		}
		if(defined $x1{$n} and defined $y1{$n} and defined $z1{$n}){
			
		}else{
			$transflag=0;
			#print "model coordinates error: $model\t$target\n";
		}		
		if($transflag==1){
			$xt1{$n} = $cx[0]*$x1{$n} + $cx[1]*$y1{$n} + $cx[2]*$z1{$n} + $cx[3];	
			$yt1{$n} = $cy[0]*$x1{$n} + $cy[1]*$y1{$n} + $cy[2]*$z1{$n} + $cy[3];	
			$zt1{$n} = $cz[0]*$x1{$n} + $cz[1]*$y1{$n} + $cz[2]*$z1{$n} + $cz[3];
		}	
	}

  # As a control, calculate GDT_TS for pairs of corresponding atoms; compare to that reported by LGA.
	my @rewardres1 =();
	my @rewardres2 =();
	for $i(0..$#res1) {
		my $n1 = $res1[$i];
		my $n2 = $res2[$i];

		$dx = $xt1{$n1} - $x2{$n2};
		$dy = $yt1{$n1} - $y2{$n2};
		$dz = $zt1{$n1} - $z2{$n2};

		$d = sqrt($dx*$dx + $dy*$dy + $dz*$dz);
		if($d < $lgad) { # within radius cutoff
			unless( defined $res_gc{$n1}{$n2} ) { $res_gc{$n1}{$n2}=0.0; }
			$res_gc{$n1}{$n2} += 0.25;
			push(@rewardres1,$n1);
			push(@rewardres2,$n2);
		}
	}

	# For each pair of non-corresponding atoms, compare distance to threshold and penalize accordingly;
	for $i(0..$#rewardres1) {
		$n1 = $rewardres1[$i];
		unless(exists($res_penalty_mod{$lgad}{$n1}) ) { $res_penalty_mod{$lgad}{$n1} = 0.0; }
		# Find spatially close non-correspoding CAs (in the whole target) for CA1
		for $m2(keys %x2) {

			next if abs($m2 - $rewardres2[$i]) <= 1; # skip aligned residue pairs and their neighbors 
# collect penalties for all three cutoffs keys %pen; average them
			$dx = $xt1{$n1} - $x2{$m2};
			$dy = $yt1{$n1} - $y2{$m2};
			$dz = $zt1{$n1} - $z2{$m2};
			$d = sqrt($dx*$dx + $dy*$dy + $dz*$dz);

		#	print STDERR "LGAD=$lgad : $n1 $m2  d=$d\n";

			for $dcut(keys %pen) {
				if($d < $dcut) { # within radius cutoff
					$res_penalty_mod{$lgad}{$n1} += $pen{$dcut};
				}
			} 
		}


		$n2 = $rewardres2[$i];
		unless(exists($res_penalty_targ{$lgad}{$n2}) ) { $res_penalty_targ{$lgad}{$n2} = 0.0; }
		# Find spatially close non-correspoding CAs (in the whole model) for CA2
		for $m1(keys %x1) {

			next if abs($m1 - $rewardres1[$i]) <= 1; # skip aligned residue pairs and their neighbors 
# collect penalties for all three cutoffs; average them
			$dx = $xt1{$m1} - $x2{$n2};
			$dy = $yt1{$m1} - $y2{$n2};
			$dz = $zt1{$m1} - $z2{$n2};
			$d = sqrt($dx*$dx + $dy*$dy + $dz*$dz);

		#	print STDERR "LGAD=$lgad : $m1 $n2  d=$d\n";

			for $dcut(keys %pen) {
				if($d < $dcut) { # within radius cutoff
					$res_penalty_targ{$lgad}{$n2} += $pen{$dcut};
				}
			} 
		}
	}

}

# Calculate total penalized GDTTS for the current LGA distance
my %gcpen = ();
for $lgad(@lgads) { $gcpen{$lgad} = 0.0; }
my $gc=0.0;
for $n1(sort {$a<=>$b} keys %res_gc) {
	for $n2(sort {$a<=>$b} keys %{$res_gc{$n1}}) {

#		next unless exists($res_gc{$n1}{$n2}); # take only residue pairs aligned by at least one LGA superposition
		$gc += $res_gc{$n1}{$n2};

		for $lgad(@lgads) { 		
			$res_penalty_mod{$lgad}{$n1} = 0.0 unless exists($res_penalty_mod{$lgad}{$n1});
			$res_penalty_targ{$lgad}{$n2} = 0.0 unless exists($res_penalty_targ{$lgad}{$n2});
				
			my $res_penalty = 0.5*($res_penalty_mod{$lgad}{$n1} + $res_penalty_targ{$lgad}{$n2});	
			$res_penalty *= 1.0/scalar(keys %pen);	
			$res_gcpen = $res_gc{$n1}{$n2} - $weight * $res_penalty;
			$res_gcpen = 0.0 if $res_gcpen < 0.0;

		#	print STDERR "$n1\t$n2\t$res_gc{$n1}{$n2}\t$res_penalty\t$res_gcpen\n";

			$gcpen{$lgad} += $res_gcpen;	
		}
	}
}


my $pdblen=keys %x2;
my $newgc=$gc* (1.0/$pdblen)*100;
$newgc=sprintf("%.3f",$newgc);
if($newgc!=$checkgdtts){
		if($checkgdtts==-1.0){
			
		}else{	
			print "warning $model $target $newgc $checkgdtts $npairs $pdblen\n";
		}	
		#exit;
}	else{
		#print "$newgc $checkgdtts $npairs $pdblen\n";
}


open(FOUT, ">> $outfile") || die "Cannot open $outfile: $!\n"; 
print FOUT "Model\tTarget\tPenalty_W\tNrespairs\tPDBlen\t";
for $lgad(@lgads) { 	
	print FOUT "PenGDTTS\_$lgad\t";
}
print FOUT "GDTTS\n";
print FOUT "$model\t$target\t$weight\t$npairs\t$pdblen\t";
for $lgad(@lgads) { 	
	$gcpen{$lgad} *= (1.0/$pdblen)*100;
	$gcpen{$lgad}=sprintf("%.3f",$gcpen{$lgad});
	printf(FOUT "%.3f\t", $gcpen{$lgad});
}
# printf(FOUT "%.3f\n", $newgc);
printf(FOUT "%.3f\n", $checkgdtts);
close(FOUT);