#!/usr/bin/python 
#Program:
#	This program works in the parent directory of subdirectories containing VASP calculations for magnetic configurations.
#	The subdirectories are named in a 44-- fashion, corresponding to a magnetic configuration of 4,-4,4,-4 in POSCAR_M, which only contains magnetic ions. 
#	Files and directories with letters a-z,A-Z in their names are ignored.
#
#	Based on the Heisenberg Hamiltonian, construct array of equations like 3J_1+2J_2+0J_3+E_0=E
#	and outputs its rank.
#	Prints the solved J's, a pair of atoms whose magnetic coupling are of the J's type, and the distances.
#
#	Different magnetic interactions are classified according to interatomic distances. An error of 1E-5 is allowed, so trim your POSCAR_M if necessary!
#	However, for pairs with same distances but different types of couplings, the distances are manually increased by 0.4 in all outputs (and in the program itself).
#	Also, distances are calculated between atoms in a unit cell, and its 27 nearest mirrors (for obvious reasons).
#	
#Usage:
#	Execute the python script directly in the parent directory. No parameters are needed.
#
#Tunable parameters:
#	threshold_dist = maximum number of distances
#	exclude_dist = interatomic distances corresponding to the J's. (remember to add 0.4 if necessary)
#	exclude_pair = pair indices corresponding to the J's
#	The IF statement...
#
#Dependencies:
#	A 'grepen' executable in $PATH, that, when run in a directory with OUTCAR, outputs the energy, which
#	 should be something like: grep "energy without" OUTCAR | tail -1 | awk '{printf "%12.6f \n", $5 }' 
#History:
#	2015/08/13 Xiang First release
#	2015/08/15 Xiang Modified to display matrix rank, and use least square method.
#	2015/08/23 Xiang Modified to read POSCAR_M, to include exclude_pair feature, and to include +0.4 mechanism.
#	2015/08/28 Xiang Modified to conform to VBird Shebang standard.
#
#I.Get the types of interactions (classified by bond length) and atom pairs as a list-of-interactions(list-of-2-tuples(tuples))
#imports and options.
import os
import numpy as np
from subprocess import call
import re
np.set_printoptions(threshold=np.nan,suppress=True)
#----------------------------------User-adjustable parameters-----------------------------------
threshold_dist=3
exclude_pair=[[6,155],[123,456]]
exclude_dist=[3.5957,4.0072,5.3177,5.7177,6.0671,6.6585,6.9914,7.3914,6.1673]
#----------------------------------End adjustable parameters------------------------------------
#1.get system parameter
f=open("POSCAR_M","r")
lines=f.readlines()
base=[lines[i].split()[0:3] for i in range(8,len(lines))]
cell=[lines[i].split() for i in range(2,5)]
base=np.float64(base)
cell=np.float64(cell)
#2.image to supercell: for i,j,k,baseindex: create image array 1*3 ; append image array to endpoint
#toappend=1darray:base[index]+1darray:[i,j,k]*2darray:cell(a,\\b,\\c)
endpoint=[]
for i in [1,0,-1]:
 for j in [1,0,-1]:
  for k in [1,0,-1]:
   toadd=np.float64([i,j,k])
   for baseindex in range(0,len(base)):
    toappend=toadd+base[baseindex] 
    toappend=np.dot(toappend,cell)
    endpoint.append(toappend)
endpoint=np.float64(endpoint)
base=np.dot(base,cell)
#3.calculate distances. do this with an array:distance
#calculate distance
distance=[]
changefromto=[]
for i in range(0,len(base)):
 for j in range(0,len(base)):
  for k in [j+l*len(base) for l in range(0,27)]:
   dist=np.linalg.norm(base[i]-endpoint[k])
   if abs(dist)<0.1:
    continue
   #THE 'IF' JUDGEMENT THAT ADD 0.4 TO DISTANCES WHEN NECESSARY. i IS THE INDEX OF FIRST ATOM, STARTING FROM 0. j IS THAT OF THE SECOND. k IS THE INDEX OF IMAGED SECOND ATOM.
   if i>1 and j>1:   
    dist=dist+0.4
   distance.append([i,j,dist])
distance=np.float64(distance)
distance=distance[np.argsort(distance[:,2])]
#get the excluded distances by the way: using list exclude_list
for pair in exclude_pair:
 pair=np.float64(pair)
 exclude_tris=[tridist for tridist in distance if tridist[0]==pair[0] and tridist[1]==pair[1]]
 if exclude_tris!=[]:
  exclude_dist.append(exclude_tris[0][2])
exclude_dist=np.float64(exclude_dist)
#4.get non-repetitive distance list and truncate distance with a list-of-array:truncated_distance and list nonrep_distance
nonrep_distance=[]
truncated_distance=[]
for tridist in distance:
 dist=tridist[2]
 if (abs(exclude_dist-dist)<1e-3).any():   #problem here is ACCURACY
  continue
 if len(nonrep_distance)<=threshold_dist:
   truncated_distance.append(tridist)
 else:
  break
 if len(nonrep_distance)==0 or (abs(nonrep_distance-dist)>1e-3).all():
  nonrep_distance.append(dist)
truncated_distance.pop(-1)
nonrep_distance.pop(-1)
truncated_distance=np.float64(truncated_distance)
#get indices. do this with a list of list of 2-element-lists:indices
indices=[]
for dist in nonrep_distance:
 indices.append([[x[0],x[1]] for x in truncated_distance if abs(x[2]-dist)<1e-3])
#II.In root folder with 11111111-------- format, get each folder's magmomlist:mag_list
print '============================================================================================================================================'
print '\tEquation'
print '============================================================================================================================================'
matrices=[]
for i in range(0,threshold_dist):
 tmpmat=np.zeros((len(base),len(base)))
 for j in indices[i]:
  tmpmat[tuple(j)]=tmpmat[tuple(j)]+0.5
 matrices.append(tmpmat)
#iterate through subdirs
A=[]
b=[]
x=[]
for subdir in os.popen('ls').read().split():
 if re.search('[a-zA-Z]',subdir):
  continue
 #get magnetic moments list
 mag_list=[]
 mag_list=' '.join(subdir).split()
 if 1==3:
  for i in range(0,len(mag_list)):
   if mag_list[i]=='-' and i>=8:
    mag_list[i]=-0.5
   elif mag_list[i]=='-' and i<8:
    mag_list[i]=-2
   elif mag_list[i]=='1' and i>=8:
    mag_list[i]=0.5
   elif mag_list[i]=='0':
    mag_list[i]=0
   else: 
    mag_list[i]=2
 elif 1==2:
  for i in range(0,len(mag_list)):
   if mag_list[i]=='-' and i>=2:
    mag_list[i]=-0.5
   elif mag_list[i]=='-' and i<2:
    mag_list[i]=-2
   elif mag_list[i]=='1' and i>=2:
    mag_list[i]=0.5
   elif mag_list[i]=='0':
    mag_list[i]=0
   else: 
    mag_list[i]=2
 else:
  for i in range(0,len(mag_list)):
   if mag_list[i]=='-':
    mag_list[i]=-1  
   else:                         
    mag_list[i]=1   
 #A_line=(product1,product2,...), set matrix mat first
 A_line=[]
 for i in range(0,threshold_dist):
  product_x=np.matrix(mag_list)
  product_xT=product_x.getT() 
  product=product_x*matrices[i]*product_xT
  A_line.append(product.item(0))
 A_line.append(1)
 A.append(A_line)
 os.chdir(subdir)
 energy=os.popen('grepen').read()
 os.chdir('..')
 b.append(energy)
#specific assignments
#A.append([1,0,0,0])
#b.append(3.6718/8065.54429)
#III. Compute energy-coefficient of various J's. THE LAST J IS NON-MAG BASIC ENERGY.
#solve linear equation by least-squre
solution,residuals,rank,singular=np.linalg.lstsq(A,b,rcond=-1)
for i in range(0,len(A)):
 print '\t'+' \t'.join(format(A[i][idx],'5.1f')+'J_'+str(idx+1)+'+' for idx in range(0,len(A[i])))+'\t&=&\t'+str(float(b[i]))+' eV\\\\ '
print '============================================================================================================================================'
print '\tSolution'
print '============================================================================================================================================'
for i in range(0,len(solution)):
 if i==len(solution)-1:
  pairsname="Non-magnetic energy"
 else:
  pairsname="Distance:"+"{:.4f}".format(nonrep_distance[i])+'\tpairs:'+str(np.int_(np.array(indices[i][0])+1))+'\t'+str(len(indices[i])/2)+'pairs/unitcell'
 print '\t'+'J'+str(i+1)+'\t=\t'+"{:8.4f}".format(solution[i])+' eV\t# '+pairsname
print '\tNorm of residual='+str(residuals*8065.54429)+' cm-1\tEquation rank='+str(rank)
tmpE=abs(solution[0]*6+max(solution[1],solution[2])*12)*11601
print solution[1],solution[2],solution[0]
print max(tmpE-1000,50),tmpE+1000,100

#for i in indices:
# for j in i:
#  print j