#!/usr/bin/perl

use Math::VectorReal;
use Math::MatrixReal;

#=== This program is described on page 107 of simulations book 200-300

# --- Program defaults are here!
$maxDistAngstroms = 10;
# Nearest neighbor cutoff distance in Angstroms

#---------------------------------------------------------------------------------------------


$eyeFile = $ARGV[0];
$xyzFile = $ARGV[1];
@supercellEdge = (0, 0, 0);
@cellsPerSupercell = (5, 5, 5);

print "\nThis is adp.pl. Finding atomic displacement factors. Hit enter at a prompt to use default values.\n";


open (READIT, $eyeFile) || die("** Can't read file! ** \nScript use: perl adp.pl [eyeinput] [xyzinput].cfg\n");
@eyeIn = <READIT>;
close (READIT);


print "Done reading in file $eyeFile \n";


print "\nEnter the number of unit cells in each direction (default @cellsPerSupercell): ";

$userInput = <STDIN>;
chomp($userInput);
if ($userInput ne "") {
	@cellsPerSupercell = split(" ",$userInput);
}


if ($#cellsPerSupercell != 2) {exit;}



$i=0;
foreach (@eyeIn) {
	chomp();
	if (/Number of particles \=/) {
		$totalAtoms = trim($');
	}

	if (/H0\(1,1\) = /) {
		$supercellEdge[0] = trim($');		
	}
	
	if (/H0\(2,2\) = /) {
		$supercellEdge[1] = trim($');		
	}

	if (/H0\(3,3\) = /) {
		$supercellEdge[2] = trim($');		
	}

	if (!/^\w/) {
		push(@atomIn, $eyeIn[$i]);
	}

	$i++;
	
}

foreach (@supercellEdge) {
	chop(); chop;
}


print "I found $totalAtoms atoms\n";
print "Supercell dimensions are @supercellEdge Angstroms\n\n";


print "\nEnter the maximum distance in A [this is often irrelevant] (default: $maxDistAngstroms): ";

$userInput = <STDIN>;
chomp($userInput);
if ($userInput ne "") {
	$maxDistAngstroms = $userInput;
}

print "\nUsing max distance = $maxDistAngstroms A\n";



#assuming cubic here.
$maxDist = $maxDistAngstroms/$supercellEdge[0]*$cellsPerSupercell[0];
$maxDist2 = $maxDist*$maxDist;

print "Squared distances is $maxDist2 \n";


open (READIT, $xyzFile) || die("** Can't read file! ** \nScript use: perl adp.pl [eyeinput] [xyzinput].cfg\n");
@xyzIn = <READIT>;
close (READIT);

# read in ideal positions. index 0=label, 1=x, 2=y, 3=z
# each xyz needs to be divided by the supercell dimensions
$i=0;
$numSites = $#xyzIn -1;
foreach (@xyzIn) {
	if ($i > 1) {
		chomp();
		@idealIn = split(/\s+/, $xyzIn[$i]);
		$idealPosition[$i-2][0] = $idealIn[0];
		for ($j = 1; $j < 4; $j++) {
			$idealPosition[$i-2][$j] = $idealIn[$j]; #/$supercellEdge[$j-1];
		}
	}

	$i++;
}

print "Total number of ideal sites is $numSites\n";

$lastAtomType = "XX";
$i=0;
$j=0;
$k=0;
foreach (@atomIn) {
	my @atomLine = split(/\s+/, $atomIn[$i]);
	if ($atomLine[2] eq $lastAtomType) {
		$atom[$i][0] = $k;
		$j++;
	} else {
		if ($j > 0 ) {
			push(@totalAtomsOfType, $j);
			print "Total number of $lastAtomType atoms is $j\n";
			$k++;
		}
		print "Found new atom type: $atomLine[2]\n";
		$atom[$i][0] = $k;
		$lastAtomType = $atomLine[2];
		$j=1;
	}
	$atom[$i][0] = $atomLine[2];
	$atom[$i][1] = mod($atomLine[3]*$cellsPerSupercell[0],1);
	$atom[$i][2] = mod($atomLine[4]*$cellsPerSupercell[1],1);
	$atom[$i][3] = mod($atomLine[5]*$cellsPerSupercell[2],1);

	$i++;
}
print "Total number of $lastAtomType atoms is $j\n\n";

print "Atomic positions are loaded into memory!\n\nFinding nearest ideal sites...\n";


# search over all atoms and find the nearest matching ideal site
for ($i=0; $i < $totalAtoms; $i++) {

	$shortestDistance = 999;
	
	# search over all while correcting for PBC
	for ($j=0; $j < $numSites; $j++) {
		if ($atom[$i][0] ne $idealPosition[$j][0]) {
			next;
			#this requires the element/label sites to be equivalent. sometimes this should be disabled.
		}

		$xDist = $idealPosition[$j][1] - $atom[$i][1];
		$yDist = $idealPosition[$j][2] - $atom[$i][2];
		$zDist = $idealPosition[$j][3] - $atom[$i][3];


		if ($xDist > 0.5) {
			$xDist = $xDist - 1;
		} elsif ($xDist < -0.5) {
			$xDist = $xDist + 1;
		}

		if (abs($xDist) > $maxDist) {
			next;
		}

		if ($yDist > 0.5) {
			$yDist = $yDist - 1;
		} elsif ($yDist < -0.5) {
			$yDist = $yDist + 1;
		}

		if (abs($yDist) > $maxDist) {
			next;
		}

		if ($zDist > 0.5) {
			$zDist = $zDist - 1;
		} elsif ($zDist < -0.5) {
			$zDist = $zDist + 1;
		}

		if (abs($zDist) > $maxDist) {
			next;
		}

		#this assumes that the 3 directions are on the same scale
		#check bvs and angle scripts!

		$vectDist2 = $xDist*$xDist + $yDist*$yDist + $zDist*$zDist;

		if (($vectDist2 < $maxDist2) && ($vectDist2 < $shortestDistance)) {
			$shortestDistance = $vectDist2;
			$displacement[0] = $xDist;
			$displacement[1] = $yDist;
			$displacement[2] = $zDist;
			
		}

	}
	$averageDisplacement[0] += $displacement[0];
	$averageDisplacement[1] += $displacement[1];
	$averageDisplacement[2] += $displacement[2];

	# calculate bond valence sums for each polyhedron, for each oxidation state [n/a here]
	
	# calculate bond angles for the NN pairs [n/a here]


}

$averageDisplacement[0] /= $totalAtoms;
$averageDisplacement[1] /= $totalAtoms;
$averageDisplacement[2] /= $totalAtoms;
print "Average displacement is @averageDisplacement\nShifting supercell and calculating ADPs...\n\n";

$finalAverageDisplacement = vector(0, 0, 0);

#shift supercell so any drifting lies back on the ideal position (on average)
for ($i=0; $i < $totalAtoms; $i++) {

	for ($j=0; $j < 3; $j++) {
		$atom[$i][$j+1] += $averageDisplacement[$j];
	}

	$shortestDistance = 999;
	
	# search over all while correcting for PBC
	for ($j=0; $j < $numSites; $j++) {
		if ($atom[$i][0] ne $idealPosition[$j][0]) {
			next;
			#this requires the element/label sites to be equivalent. sometimes this should be disabled.
		}

		$xDist = $idealPosition[$j][1] - $atom[$i][1];
		$yDist = $idealPosition[$j][2] - $atom[$i][2];
		$zDist = $idealPosition[$j][3] - $atom[$i][3];


		if ($xDist > 0.5) {
			$xDist = $xDist - 1;
		} elsif ($xDist < -0.5) {
			$xDist = $xDist + 1;
		}

		if (abs($xDist) > $maxDist) {
			next;
		}

		if ($yDist > 0.5) {
			$yDist = $yDist - 1;
		} elsif ($yDist < -0.5) {
			$yDist = $yDist + 1;
		}

		if (abs($yDist) > $maxDist) {
			next;
		}

		if ($zDist > 0.5) {
			$zDist = $zDist - 1;
		} elsif ($zDist < -0.5) {
			$zDist = $zDist + 1;
		}

		if (abs($zDist) > $maxDist) {
			next;
		}

		#this assumes that the 3 directions are on the same scale
		#check bvs and angle scripts!

		$vectDist2 = $xDist*$xDist + $yDist*$yDist + $zDist*$zDist;

		if (($vectDist2 < $maxDist2) && ($vectDist2 < $shortestDistance)) {
			$shortestDistance = $vectDist2;
			$closestSite = $j;
			$displacement[0] = $xDist;
			$displacement[1] = $yDist;
			$displacement[2] = $zDist;
			
		}

	}
	$finalAverageDisplacement[0] += $displacement[0];
	$finalAverageDisplacement[1] += $displacement[1];
	$finalAverageDisplacement[2] += $displacement[2];
	$nearestSite[$i] = $closestSite;
	$squaredDistance[$i] = $shortestDistance;
	# calculate bond valence sums for each polyhedron, for each oxidation state [n/a here]
	
	# calculate bond angles for the NN pairs [n/a here]

	

}


$finalAverageDisplacement[0] /= $totalAtoms;
$finalAverageDisplacement[1] /= $totalAtoms;
$finalAverageDisplacement[2] /= $totalAtoms;
print "final average displacement is @finalAverageDisplacement \n*** These should be really really small!\n\n";


#find adp by atom site label


for ($i=0; $i < $totalAtoms; $i++) {
	$sumSqDisp[$nearestSite[$i]] += $squaredDistance[$i];
	$numOnSite[$nearestSite[$i]]++;
}

$lastAtomType = $idealPosition[0][0];
$groupedIdealPositionName[0] = $lastAtomType;
$j=0;

for ($i=0; $i < $numSites; $i++) {
	print "Site #$i: type $idealPosition[$i][0], total atoms $numOnSite[$i],* sqdisp sum $sumSqDisp[$i]\n";
	if ($idealPosition[$i][0] ne $lastAtomType) {
		$lastAtomType = $idealPosition[$i][0];
		$j++;
		$groupedIdealPositionName[$j] = $lastAtomType;

	}
	$groupedIdealPositionSqDisp[$j] += $sumSqDisp[$i];
	$groupedIdealPositionNum[$j] += $numOnSite[$i];


}
print "\n * Check the total atoms on each site in the list above.\n If it differs from site to site something may be wrong\n";

print "\n Found " . ($#groupedIdealPositionNum + 1) . " distinct ideal positions\n\n";
print " ========================================\n  Site\tADP from ideal positions (A^2)\n ========================================\n";

for ($i=0; $i <= $#groupedIdealPositionNum; $i++) {
	$ADP[$i] = $groupedIdealPositionSqDisp[$i]/$groupedIdealPositionNum[$i]*$supercellEdge[0]*$supercellEdge[0]/$cellsPerSupercell[0]/$cellsPerSupercell[0];
	print "  $groupedIdealPositionName[$i]\t$ADP[$i] \n";
}


print "\n All finished!\n\n";


sub mod { return $_[0] - int ($_[0] / $_[1]) * $_[1]; };


sub trim($)
{
	my $foo = shift;
	$foo =~ s/^\s+//;
	$foo =~ s/\s+$//;
	return $foo;
}