#!/usr/bin/perl -w
# collect structure factors and pdb entry given a code

use strict;

use File::Temp qw/ tempdir tempfile /;
my $tdir = tempdir();

my $userepo = 1;

my @tools=("cif2mtz","truncate","cad","combat","sortmtz","scala");
for my $n (@tools)
  {
    my $loc = `which $n`;
    if (length($loc) == 0)
      {
	die "Can't find tool '$n'\n";
      }      
  }

my $wgloc = `which wget`; my $culoc = `which curl`;
if (length($wgloc) == 0 && length($culoc) == 0)
{
    die "Can't find either wget or curl\n";
}

my $grabber;
if (length($wgloc) == 0)
{
    $grabber = "curl -s \"XX\" > YY";
}
else
{
    $grabber = "wget -O YY \"XX\"";
}

# get the cif2mtz version line
my $cv = `echo "END" | cif2mtz HKLOUT /dev/null 2>/dev/null | grep version`;
if ($cv =~ /version 5\.0/)
  {
    die "You are using ccp4 version 5, whose cif2mtz cannot handle most PDB SF.CIF files.  Please use version 6 instead\n";
  }


# options: -inplace -noR
my @allopts = ("-inplace","-noR","-nocollect","-norep");
my @pdbs = ();
my %opts;
for my $potopt (@ARGV)
  {
    if ($potopt =~ /^-/)
      {
	if (grep {/$potopt/} @allopts)
	  {
	    $opts{$potopt} = 1;
	  }
	else
	  {
	    die "Unknown command-line option $potopt\n";
	  }
      }
    else 
      {
	if (length($potopt) == 4)
	  {
	    push @pdbs, $potopt;
	  }
	else
	  {
	    die "$potopt doesn't look like a PDB reference\n";
	  }
      }
  }
if ($#pdbs == -1)
  {
    die "No PDB reference given on command line\n";
  }

if (defined $opts{"-norep"} ) { $userepo = 0; }

for my $pdb (@pdbs)
  {
    my $gdir = $pdb;
    if (defined $opts{"-inplace"}) {$gdir = "."}
    
    my $wantRfree = 1;
    if (defined $opts{"-noR"}) {$wantRfree = 0}
    
    $pdb = lc $pdb;
    unless (-d $gdir) {system "mkdir $gdir";}
    
    if (not defined $opts{"-nocollect"})
      {
	my $success = 0;
	
	if ($userepo == 1)
	  {
	    my $mid2 = substr($pdb,1,2);
	    my $rn = "/repository/PDB/v3/data/structures/divided/pdb/$mid2/pdb$pdb.ent.gz";
	    if (-e $rn)
	      {
		system "gunzip -c $rn > $gdir/$pdb.pdb";
		$success ++;
	      }
	    $rn = "/repository/PDB/v3/data/structures/divided/structure_factors/$mid2/r$pdb"."sf.ent.gz";
	    if (-e $rn)
	      {
		system "gunzip -c $rn > $gdir/$pdb.sf.cif";
		$success ++;
	      }
	    if ($success != 2) {$success = 0;}
	    else
	      {
		print "Collected $pdb from repository\n";
	      }
	  }
	
	if ($success == 0)
	  {
	    print "Collecting $pdb from web\n";
	    my $px = lc(substr($pdb,1,2));

	    my $pdbname="ftp://ftp.wwpdb.org/pub/pdb/data/structures/divided/pdb/$px/pdb${pdb}.ent.gz";
	    print "Collecting $pdbname\n";
	    my $gs = $grabber; $gs =~ s#YY#$gdir/$pdb.pdb.gz#g; $gs =~ s#XX#$pdbname#g;
	    system $gs; system "gunzip $gdir/$pdb.pdb.gz";

	    my $sfcifname="ftp://ftp.wwpdb.org/pub/pdb/data/structures/divided/structure_factors/$px/r${pdb}sf.ent.gz";
	    print "Collecting $sfcifname\n";
	    $gs = $grabber; $gs =~ s#YY#$gdir/$pdb.sf.cif.gz#g; $gs =~ s#XX#$sfcifname#g;
	    system $gs; system "gunzip $gdir/$pdb.sf.cif.gz";
	    print "Collected $pdb from web\n";
	  }
      }
    die "Cannot obtain PDB file" unless (-e "$gdir/$pdb.pdb");
    die "Cannot obtain SF.CIF file" unless (-e "$gdir/$pdb.sf.cif");

    # Correct _diffrn.detail typo
    system "mv $gdir/$pdb.sf.cif $gdir/$pdb.sf.cif.orig";
    open A,"< $gdir/$pdb.sf.cif.orig";
    open B,"> $gdir/$pdb.sf.cif";
    my @lines = ();
    while (<A>)
      {
	s/_diffrn\.detail /_diffrn\.details /g;
	print B $_;
      }
    close A; close B;
	    
    
    # Get temporary filenames for the CIF2MTZ and de-Friedel procedures
    
    my ($fh1, $t1) = tempfile("$tdir/collect_$pdb.cmdlineXXXXXX");
    my ($fh2, $t2) = tempfile("$tdir/collect_$pdb.tmpmtz1XXXXXX");
    my ($fh3, $t3) = tempfile("$tdir/collect_$pdb.cmdoutXXXXXX");
    my ($fh4, $t4) = tempfile("$tdir/collect_$pdb.tmpmtz2XXXXXX");
    my ($fh5, $t5) = tempfile("$tdir/collect_$pdb.cmderrXXXXXX");
    my ($fh6, $t6) = tempfile("$tdir/collect_$pdb.fixsfcifXXXXXX");
    
    my $crystcard = `grep "^CRYST1" $gdir/$pdb.pdb`;
    print "CRYST1 card: $crystcard";
    
    #0000000000111111111122222222223333333333444444444455555555556666666666
    #0123456789012345678901234567890123456789012345678901234567890123456789
    #CRYST1  128.049  128.049  112.759  90.00  90.00 120.00 P 65          6
    
    # look at the .sf.cif file: if it lacks '_refln.status'
    # then there are no Rfree figures to be measured
    
    # [things that failed: 2agd 1ehz 1obp 1bay]
    
    my $lacksRfree = 0;
    my $hasstat = `grep -c "_refln.status" $gdir/$pdb.sf.cif`;
    if ($hasstat == 0)
      {
	print "Deposited structure factors have no Rfree flag\n";
	$lacksRfree = 1;
      }
    
    my $cell = substr($crystcard,7,47);
    my $cif2mtzcmd;# = "cell ".$cell."\n";

    # get symmetry from sf.cif file
    open A,"< $gdir/$pdb.sf.cif";
    my $sname_cif;
    while (<A>)
      {
	if (/symmetry.space_group_name_H-M +'(.*)'/)
	  {
	    $sname_cif = $1;
	  }
      }
    close A;

    my $sname = substr($crystcard,55,11);
    $sname =~ s/ //g;
    if (defined $sname_cif)
      {
	print "CIF file defines cell and space group; using that information\n";
	$sname_cif =~ s/ //g;
	if ($sname_cif ne $sname)
	  {
	    print "symmetry in CIF is different from PDB: $sname_cif vs $sname\n";
	    $sname = $sname_cif;
	  }
      }
    else
      {
	print "No space-group named in the CIF file; using the cell and symmetry from PDB\n";
	$cif2mtzcmd="cell $cell\nsymm $sname\n";
      }
    unless ($hasstat == 0) {$cif2mtzcmd .= "labout FREE=FreeR_flag\n";}
    $cif2mtzcmd .= "name project $pdb\n";
    $cif2mtzcmd .= "end\n";
    
    open A,"> $t1";
    print A $cif2mtzcmd;
    close A;

    print STDERR "Calling cif2mtz ";
    my $nfirstruns = 0; my $firstrunsuccess = 0;
    while ($nfirstruns < 5 and $firstrunsuccess == 0)
      {
	my $first_cif2mtz_result = system "cif2mtz hklin $gdir/$pdb.sf.cif hklout $gdir/$pdb.mtz< $t1 >> $t3 2>> $t5";
	if ($first_cif2mtz_result == 0) {$firstrunsuccess=1;}
	$nfirstruns++;
      }
    if ($nfirstruns > 4)
      {
	system "cat $t3 $t5";
	die "First CIF2MTZ call fell over repeatedly\n";
      }
    
    my $sfcif_confused = 0;
    open B,"< $t5";
    while (<B>)
      {
	if (/Unexpected context type/)
	  {
	    print "\nCIF2MTZ can't convert sf.cif; it contains unknown columns\n";
	    $sfcif_confused = 1;
	  }
      }
    close B;
    
    # code to trim out terms of which CIF2MTZ hasn't heard
    
    my @knownterms=("_refln.wavelength_id", "_refln.crystal_id", "_refln.scale_group_code", 
		    "_refln.index_h", "_refln.index_k", "_refln.index_l", 
		    "_refln.status", 
		    "_refln.F_meas_au", "_refln.F_meas_sigma_au",
		    "_refln.intensity_meas", 
		    "_refln.intensity_sigma",
		    "_refln.F_calc", "_refln.phase_calc", "_refln.fom");
    
    my %kti = (); for my $j (@knownterms) {$kti{$j}=1;}
    
    if ($sfcif_confused == 1)
      {
	open C,"< $gdir/$pdb.sf.cif";
	open D,"> $t6";
	my ($gathering_loop, $reling_loop, $nrpx) = (0,0) ; my (@looph, @xterm, @reflnh);
	while (<C>)
	  {
	    chomp; $_ =~ s/ *$//g;
	    $_ =~ s/intensity_meas_au/intensity_meas/g;
	    $_ =~ s/intensity_sigma_au/intensity_sigma/g;
	    if ($gathering_loop == 1)
	      {
		if ($_ =~ /^_/)
		  {
		    push @looph, $_;
		  }
		else
		  {
		    $gathering_loop = 0;
		    if ($looph[0] =~ /^_refln\./)
		      {
			$reling_loop = 1; 
			@reflnh = @looph; 
			$nrpx = $#looph;
			for my $u (0..$nrpx)
			  {
			    if (defined $kti{$reflnh[$u]}) {print D $reflnh[$u],"\n"};
			  }
		      }
		    else
		      {
			print D join "\n",@looph;
		      }
		    @looph=();
		  }
	      }
	    if ($reling_loop == 1)
	      {
		if (/^#/)
		  {
		    $reling_loop = 0;
		  }
		else
		  {
		    my @terms = split " ",$_;
		    push @xterm, @terms;
		    if ($#xterm == $nrpx)
		      {
			my @nterm = ();
			for my $u (0..$nrpx)
			  {
			    if (defined $kti{$reflnh[$u]}) {push @nterm, $xterm[$u]}
			  }
			print D join " ",@nterm,"\n";
			@xterm = ();
		      }
		  }
	      }
	    if ($reling_loop != 1 && $gathering_loop != 1)
	      {
		unless (/\_diffrn.detail/)
		  {
		    print D $_,"\n";
		  }
	      }
	    if ($_ eq "loop_")
	      {
		$gathering_loop = 1;
	      }
	  }
	close C;
	system "mv $gdir/$pdb.sf.cif $gdir/$pdb.sf.cif.1";
	system "mv $t6 $gdir/$pdb.sf.cif";
	print STDERR "Pruned sf.cif file; retrying cif2mtz\n";
	my $trial2 = system "cif2mtz hklin $gdir/$pdb.sf.cif hklout $gdir/$pdb.mtz< $t1 > $t3 2> $t5";    
	if ($trial2 != 0)
	  {
	    system "cat $t5";
	    die "Second cif2mtz run returned $trial2\n";
	  }
	else
	  {
	    print STDERR "Success: running ";
	  }
      }
    
    # now we have an MTZ file
    my $mtzn = "$gdir/$pdb.mtz";
    print STDERR "\nAnalysing MTZ file\n";

    my $resolution = -1; my $proprfact = -1; my $ofsline = -1; my $ofxline = -1;

    my ($offset,$checkl);
    open A,"< $gdir/$pdb.mtz"; read A,$checkl,4; die unless $checkl eq "MTZ ";
    read A,$offset,4; $offset=unpack("l",$offset); $offset=4*$offset-4; seek A,$offset,0;
    
    my ($line,$colno,$firstline,$dummy)=("START",0,1);
    my ($freeRcol,$ncols,$nrows);
    my @dodgycolumns = ();
    
    while (substr($line,0,4) ne "END ")
      {
#	print "debug $line\n";
	read A,$line,80;
	if ($firstline == 1)
	  {
	    $firstline = 0;
	    unless ($line =~ /^VERS MTZ/)
	      {
		die "MTZ file seems to be corrupt; expected a version line and got $line\n";
	      }
	  }
	if ($line =~ /^NCOL/)
	  {
	    ($dummy,$ncols,$nrows) = split " ",$line;
	  }
	if ($line =~ /^RESO/)
	  {
	    my ($dummy,$rlo,$rhi) = split " ",$line;
	    $rlo=1/sqrt($rlo); $rhi=1/sqrt($rhi); $resolution = $rhi;
	  }
	if ($line =~ /^COLUMN/)
	  {
	    $colno++;
	    my ($dummy,$cname,$ctype,$cmin,$cmax) = split " ",$line;
	    if ($cname eq "FreeR_flag") { $freeRcol = $colno; }
	    if ($cmin == $cmax)
	      {
		print STDERR $cname," is an uninformative column\n";
		push @dodgycolumns,$cname;
	      }
	  }
      }

    # now that the headers are parsed, run through the data itself
    # to compute the average of freeR
    seek A,80,0;
    if (defined $freeRcol)
      {
	my ($f,$sumfree);
	for my $row (0..$nrows-1)
	  {
	    read A,$dummy,4*($freeRcol-1);
	    read A,$f,4;
	    $sumfree+=unpack("f",$f);
	    read A,$dummy,4*($ncols-$freeRcol);
	  }
	$proprfact=$sumfree/$nrows;
      }

    if ($#dodgycolumns != -1)
      {
	my $zapccols = "read \"$gdir/$pdb.mtz\"\nno\n";
	for my $u (@dodgycolumns)
	  {
	    $zapccols .= "delete column $u\n";
	  }
	$zapccols .= "write \"$gdir/$pdb.mtz\"\nyes\nquit\n";
	open A,"> $gdir/zdcscript";
	print A $zapccols;
	close A;
	system "sftools < $gdir/zdcscript > $gdir/$pdb.zdclog";
      }
    
    if ($resolution == -1) {die "Can't parse out resolution range from MTZ file"};
    if ($proprfact != -1)
      {
	if ($proprfact == 1 or $proprfact == 0)
	  {
	    if ($proprfact == 0) {print "All "} else {print "No "};
	    print "reflections marked as free; throwing away free-reflections column\n";
	    $cif2mtzcmd =~ s/labout FREE=FreeR_flag//g;
	    open A,"> $t1";
	    print A $cif2mtzcmd;
	    close A;
	    
	    my $cif2mtz_r = system "cif2mtz hklin $gdir/$pdb.sf.cif hklout $gdir/$pdb.mtz< $t1 > $t3";
	    die "cif2mtz returned $cif2mtz_r" unless ($cif2mtz_r == 0);
	    $lacksRfree = 1;
	  }
	if ($proprfact != 0 && $proprfact < 0.5)
	  {
	    print "\n More than half the reflections are marked as free; assuming sense is wrong\n";
	    my $fliprsense = "read \"$gdir/$pdb.mtz\"\nno\ncalc col FreeR_flag = 1 col FreeR_flag -\nwrite \"$gdir/$pdb.mtz\"\nyes\nquit\n";
	    open A,"> $t1";
	    print A $fliprsense;
	    close A;
	    system "sftools < $t1 > $gdir/$pdb.sflog";
	  }
      }
    
    my $hasForFP = `grep -c "_refln.F_meas" $gdir/$pdb.sf.cif`;
    my $hasI=`grep -c "_refln.intensity_meas" $gdir/$pdb.sf.cif`;
    if ($hasForFP == 0 && $hasI == 1)
      {
	print "MTZ file contains intensities not amplitudes; truncating\n\n";
	open A,"> $t1";
	print A "labin IMEAN=I SIGIMEAN=SIGI\nEND\n";
	close A;
	system "truncate HKLIN $gdir/$pdb.mtz HKLOUT $gdir/$pdb.tmtz < $t1 > $gdir/$pdb.tlog";
	system "mv $gdir/$pdb.mtz $gdir/$pdb.noamp.mtz";
	system "mv $gdir/$pdb.tmtz $gdir/$pdb.mtz";
      }
    my $friedel_done_already = 0;
    if ($hasForFP == 0 && $hasI == 0)
      {
	my $hasFpm = `grep -c "_refln.pdbx_F_plus" $gdir/$pdb.sf.cif` * `grep -c "_refln.pdbx_F_minus" $gdir/$pdb.sf.cif`;
	if ($hasFpm != 0)
	  {
	    print "MTZ file has F(+) and F(-): calling \n";
	    open A,"> $t1";
	    print A "INPUT MTZF\nLABI F(+)=F(+) SIGF(+)=SIGF(+) F(-)=F(-) SIGF(-)=SIGF(-)\nBATCH 1\n";
	    close A;
	    print "combat ";
	    system "combat HKLIN $gdir/$pdb.mtz HKLOUT $gdir/$pdb.t1mtz < $t1 > $gdir/$pdb.clog1";
	    open A,"> $t1";
	    print A "H K L M/ISYM BATCH\n$gdir/$pdb.t1mtz\n";
	    close A;
	    print "sortmtz ";
	    system "sortmtz HKLOUT $gdir/$pdb.t2mtz < $t1 > $gdir/$pdb.clog2";
	    print "scala ";
	    system "echo ONLYMERGE | scala HKLIN $gdir/$pdb.t2mtz HKLOUT $gdir/$pdb.t3mtz > $gdir/$pdb.clog3";
	    print "truncate ";
	    system "echo 'ANOM YES' | truncate HKLIN $gdir/$pdb.t3mtz HKLOUT $gdir/$pdb.mtz > $gdir/$pdb.clog4";
	    print "\n\n";
	    $friedel_done_already = 1;
	  }
      }
    
    # Try to sort out cases with anomalous differences in the data
    if ($friedel_done_already == 0)
      {
	print "\nChecking for Friedel pairs\n";
	
	#print "Temp files $t1 $t2 $t3 $t4\n";
	
	my $mtzan = "$gdir/$pdb"."_anom.mtz";
	
	open A,"> $t1";
	print A "LABI FILE 1 ALL\n OUTL SPAC $sname\nEND\n";
	close A;
    
	my ($ur,$sr);
	system "cad hklin1 $mtzn hklout $t2.mtz < $t1 > $t3";
	open U,"< $t3";
	while (<U>)
	  {
	    if (/Unique records [^0-9]*([1-9][0-9]+)/) {$ur=$1;}
	    if (/Sorting Routine [^0-9]*([1-9][0-9]+)/) {$sr=$1;}
	  }
	print $ur," unique records from ",$sr," (".sprintf("%.1f",100*$ur/$sr)."%)\n";
	
	if ($ur/$sr < 0.8)
	  {
	    print "Merging Friedel pairs\n";
	    open A,"> $t1";
	    if ($hasForFP != 0)
	      {
		print A "INPUT MTZF\nLABI F=FP SIGF=SIGFP\nBATCH 1\n";
	      }
	    else
	      {
		print A "INPUT MTZF\nLABI F=F SIGF=SIGF\nBATCH 1\n";
	      }
	    close A;
	    my $combat_r = system "combat hklin $mtzn hklout $t2.mtz < $t1 > $t3";
	    unless ($combat_r == 0) {die "combat returns $combat_r (input $mtzn output $t2.mtz commands in $t1 log in $t3\n"};
	    
	    open A,"> $t1";
	    print A "H K L M/ISYM BATCH\n$t2.mtz\n";
	    close A;
	    system "sortmtz hklout $t4.mtz < $t1 >> $t3";
	    
	    my $scala_r = system "echo ONLYMERGE | scala hklin $t4.mtz hklout $t2.mtz >> $t3";
	    die "scala returns $scala_r (input MTZ $t4.mtz, writing $t2.mtz, log in $t3)\n" unless ($scala_r == 0);
	    unlink "ROGUES"; unlink "NORMPLOT"; unlink "CORRELPLOT";
	    system "mv $mtzn $mtzan";
	    my $truncate_r = system "echo 'ANOM YES' | truncate hklin $t2.mtz hklout $mtzn >> $t3";
	    die "truncate returns $truncate_r (input MTZ $t2.mtz; log in $t3)\n" unless ($truncate_r == 0);
	    unless (defined $ENV{"KEEPTMP"}) 
	      { system "rm $t1 $t2.mtz $t3 $t4.mtz $t5"; }
	  }
	else
	  {
	    # use the CAD output file because it has been sorted
	    system "mv $t2.mtz $mtzn";
	  }
      }
    
    # Add an Rfree if we need to
    if ($lacksRfree && $wantRfree)
      {
	# Run unique first
	open A,"> $t1";
	print A "SYMM $sname\nRESOL $resolution\nCELL $cell\n";
	close A;
	my $unique_r = system "unique HKLOUT $t3.mtz < $t1 >/dev/null";
	die "Failed to run unique" unless ($unique_r == 0);
	print STDERR "Adding 5% Rfree set\n";
	open A,"> $t1";
	print A "MODE BATCH\nREAD $t3.mtz\nRFREE 0.05\nWRITE $t2.mtz\nEND\n";
	close A;
	my $sftools_r = system "sftools < $t1 > $t3";
	die "Problem adding Rfree set to unique output" unless ($sftools_r == 0);
	open A,"> $t1";
	print A "LABI FILE 1 E1=RFREE\nLABO FILE 1 E1=FreeR_flag\nLABI FILE 2 ALL";
	close A;
	my $cad_r = system "cad HKLIN1 $t2.mtz HKLIN2 $mtzn HKLOUT $t3.mtz < $t1 > $t3";
	die "Problem merging Rfree set" unless ($cad_r == 0);
	system "mv $t3.mtz $mtzn";
      }
    
    unless (defined $ENV{"KEEPTMP"}) 
      {
	for my $A ($t1,$t2,"$t2.mtz",$t3,$t4,$t5,$t6)
	  {
	    unlink($A);
	  }
      }
  }
