#!/usr/bin/python
import time, os, LGI;
import math, sys;

def create_poscar_file(location, Z, r):
    #CREATE A POSCAR FILE
    #input: a3D [A](double) -- Supercell [int] -- layers [int] -- vacuum distance [A](int)  -- fix -- layerpositions [] [] []... -- Atom/Molecule -- u -- v -- t -- p -- Z -- r 
    #string = "./pes 3.6287 3 6 20 fix 0.00000 2.09503 4.18699 6.28615 8.37945 10.46748 Molecule %s %1.5f %1.5f > POSCAR" %(location, Z, r);  #SRPsol75-DF2 layer positions
    #string = "./pes 3.6349 3 6 20 fix 0.00000 2.09861 4.19334 6.29482 8.39086 10.48173 Molecule %s %1.5f %1.5f > POSCAR" %(location, Z, r); #SRPsol70-DF2 layer positions
    #string = "./pes 3.64427 3 6 20 fix 0.00000 2.10402 4.20374 6.31019 8.41056 10.50654 Molecule %s %1.5f %1.5f > POSCAR" %(location, Z, r); #SRPsol62.65-DF2 Layer positions
    string = "./pes 3.63999 3 6 20 fix 0.00000 2.10154 4.19923 6.30387 8.40225 10.49567 Molecule %s %1.5f %1.5f > POSCAR" %(location, Z, r); #B86SRP68
    print "string= ", string; 
    os.system(string) #this is the commandline, executing the command in varibale string --> creating a POSCAR file using the desirec input
    print 

# FUNCTION ENDS HERE
#----------------------------------------------------------------------------------------------------------

def check_finished(CURRENT_DIRECTORY, Client, functional ):
    print "Entering check_finished function...."
    print "Current date & time " + time.strftime("%c")

    LIST_FINISHED  = Client.GetJobList( Application = "vasp535", State = "finished", Limit=1024);
    N_FINISHED = LIST_FINISHED[ 'LGI' ][ 'response' ][ 'number_of_jobs' ];

    if (N_FINISHED > 1):
      #for JOB_FINISHED in LIST_FINISHED[ 'LGI' ][ 'response' ][ 'job' ][ : ][ 'job' ]:
      for job_num in range(0, N_FINISHED):
 
        #GET THE JOB ID  
        JOB_FINISHED   = LIST_FINISHED[ 'LGI' ][ 'response' ][ 'job' ][ job_num ][ 'job' ];                   
        job_id         = JOB_FINISHED[ 'job_id' ]
        job_specifics  = JOB_FINISHED[ 'job_specifics' ]
        project        = job_specifics[ 'project' ]

        if ( project == "PES" ):
          repository_url = job_specifics[ 'repository_url' ] 
          folder         = job_specifics[ 'foldername' ]
          func           = job_specifics[ 'functional' ] 

          #only do something if the job is from batch 
          if ( func == functional ):
            #TRANSFER FILES
            #CHANGE THE DIRECTORY  
            os.chdir("LGI_OUTPUT");
            #MAKE A NEW DIRECTORY
            os.mkdir("%s" %(folder));
            os.chdir("%s" %(folder));
            #TRANFER FILE
            #Client.LGI_filetransfer download repository_url OUTCAR OSZICAR POSCAR vasprun.xml ;
            print "downloading: ", repository_url
            FileTransferClient = LGI.LGI_Client( URL = repository_url );
            FileTransferClient.DownLoadFiles( [ "OUTCAR", "OSZICAR", "CHGCAR", "POSCAR", "vasprun.xml" ], "./" );
            print "deleting job nr ", job_id
            Client.DeleteJob( job_id );
            del FileTransferClient;
            #GO BACK TO WORKING DIRECTORY
            os.chdir(CURRENT_DIRECTORY);




# FUNCTION ENDS HERE
#----------------------------------------------------------------------------------------------------------

#ASK FOR THE CURRENT DIRECTORY
CURRENT_DIRECTORY = os.getcwd()

project    = "PES"
functional = "B86SRP68-DF2"
Queue      = "long"
maxwalltime= "12:00:00"
maxram     = "60"
Client     = LGI.LGI_Client();


#Create A POSCAR FILE
#location contains the location for which to do RZ cuts [ "u v theta phi" ]

#		T2H ,T2H, T2H , T2H, T2H, T2H, FCC, FCC, FCC, FCC, FCC, T2F, T2F, T2F, T2F, T2F, T2F 				
location    = [ "0 0 0 0", "0 0 90 0", "0 0 90 30", "0.5 0 0 0", "0.5 0 90 0", "0.5 0 90 60", "0.5 0 90 90", "0.66666666 0.33333333 0 0", "0.66666666 0.33333333 45 30", "0.66666666 0.33333333 45 210", "0.66666666 0.33333333 90 0", "0.66666666 0.33333333 90 30", "0.33333333 0.16666666 0 0", "0.33333333 0.16666666 45 30", "0.33333333 0.16666666 45 120", "0.33333333 0.16666666 45 210", "0.33333333 0.16666666 90 30", "0.33333333 0.16666666 90 120", "0.33333333 -0.33333333 0 0", "0.33333333 -0.33333333 45 150", "0.33333333 -0.33333333 45 330", "0.33333333 -0.33333333 90 0", "0.33333333 -0.33333333 90 330", "0.16666666 -0.16666666 0 0", "0.16666666 -0.16666666 45 150", "0.16666666 -0.16666666 45 240", "0.16666666 -0.16666666 45 330", "0.16666666 -0.16666666 90 240", "0.16666666 -0.16666666 90 330" ];
R           = [0.4, 0.5, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 1.0, 1.25, 1.5, 1.75, 2.0, 2.3];
Z           = [ 0.25, 0.5, 1.0, 1.25, 1.50, 1.75, 2.0, 2.25, 2.50, 2.75, 3.0, 3.25, 3.50, 3.75, 4.0, 4.25, 4.5, 4.75, 5.0, 5.5, 6.0, 7.0 ];


#Nloc = 1;
#NR = 0;
#NZ = 0;
Nloc = 29;
NR = 14; 
NZ = 22;

MAX_NUM_RUNNING_JOB  = 1000;
MAX_NUM_QUEUED_JOB   = 300;

#os.mkdir("LGI_OUTPUT");
COUNT = 0;
COUNTFINISH = 0

my_file = "restart_file";
#start_val = [0, 0, 0];
start_val = [0, 0, 0]
#check whether the restart_file exists
# and if it exists, then read the start value

if os.path.isfile(my_file) :
  restart_file = open("%s" %my_file, "r");
  data_liste = (restart_file.readline());
  start_val = data_liste.split(" ");
  start_val[0] = int(start_val[0]);
  start_val[1] = int(start_val[1]);
  start_val[2] = int(start_val[2]) + 1; #BECAUSE THIS IS THE NEXT GEOMETRY TO BE CALCULATED  
  restart_file.close();
  print "RESTART FILE FOUND \n";
else :
  start_val[0] = int(0);
  start_val[1] = int(0);
  start_val[2] = int(0);
  print "NO RESTART_FILE, START FROM SCRATCH\n";
print "START VALUES ARE ",  start_val[0], start_val[1], start_val[2]


#start submit loop
for loc in range(start_val[0], Nloc):
  for z in range(start_val[1], NZ):
    for r in range(start_val[2],NR):
   
      print "now at iteration:"

      #folder name we are going to give to job_specifics together with functional
      foldername = "site{0}z{1}r{2}".format(loc,z,r)
      print "functional= ", functional, " loc= ", loc, " location= ", location[loc], " z= ", z, " Z= ", Z[z], " r= ", r , " R= ", R[r], " foldername= ", foldername
      #CREATE A POSCAR FILE
      create_poscar_file(location[loc], Z[z], R[r])


      # FIRST: ASK LGI HOW MANY JOBS ARE RUNNING
      #        AND HOW MANY JOBS ARE QUEUED
      # IF THE NUMBER OF QUEUED JOBS IS LARGER THEN MAX_NUM_QUEUED_JOB
      # THEN FALL ASLEEP, OTHERWISE CONTINUE WITH SUBMITTING JOBS
      CONTINUE_VAR = 1 ;
      while CONTINUE_VAR == 1 :
        LIST_QUEUED  = Client.GetJobList( State = "queued", Limit=0 );
        LIST_RUNNING = Client.GetJobList( State = "running", Limit=0);
        N_QUEUED     = LIST_QUEUED[ 'LGI' ][ 'response' ][ 'number_of_jobs' ] ;
        N_RUNNING    = LIST_RUNNING[ 'LGI' ][ 'response' ][ 'number_of_jobs' ];
        print "queued:", N_QUEUED, "running: ", N_RUNNING
        #IF THE NUMBER OF QUEUED JOBS IS SMALLER THAN 
        #if (N_QUEUED < MAX_NUM_QUEUED_JOB and N_RUNNING < MAX_NUM_RUNNING_JOB):
        if (N_QUEUED + N_RUNNING < MAX_NUM_QUEUED_JOB + MAX_NUM_RUNNING_JOB):
           CONTINUE_VAR = 0;
        else:        
           time.sleep(600);  

      #IF THERE ARE ABORTED JOBS THEN DELETE THEM
      LIST_ABORTED  = Client.GetJobList( Application = "vasp535", State = "aborted", Limit=1024);
      N_ABORTED = LIST_ABORTED[ 'LGI' ][ 'response' ][ 'number_of_jobs' ];
      if N_ABORTED > 1:
        for job_num in range(0, N_ABORTED):
          JOB_ABORTED = LIST_ABORTED[ 'LGI' ][ 'response' ][ 'job' ][ job_num ][ 'job' ];
          job_id = JOB_ABORTED[ 'job_id' ];
          job_specifics = JOB_ABORTED[ 'job_specifics' ];
          repository_url = job_specifics[ 'repository_url' ] ;
          Client.DeleteJob( job_id );

      ## SECOND: CREATE AND SUBMIT A JOB  
      JOB = Client.SubmitJob( 
      Application     = "vasp535", 
      TargetResources = "mark@erc2.wks.gorlaeus.net,mark@erc1.wks.gorlaeus.net,mark@octo.wks.gorlaeus.net",
      #TargetResources = "mark@erc1.wks.gorlaeus.net",
      FileList        = [ "./STANDARD/INCAR", "./STANDARD/POTCAR", "./POSCAR", "./STANDARD/KPOINTS", "./STANDARD/vdw_kernel.bindat", "./STANDARD/XCFUNC" ], 
      #now job specifics -- here we can give specifiers for the job in xml format
      JobSpecifics    = "<functional> {0} </functional> <foldername> {1} </foldername> <walltime> {2} </walltime> <mem> {3} </mem> <project> {4} </project> <Queue> {5} </Queue>".format(functional, foldername, maxwalltime, maxram, project, Queue) )
      
#      print "submitted: ", JOB;

      COUNT = COUNT + 1;         
      print " Number of Jobs submitted ", COUNT

     
      #write restart file
      fout_file = open("restart_file", "w");
      fout_file.write("%d %d %d" %(loc, z, r)); 
      fout_file.close();

      #now check for finished jobs
      check_finished(CURRENT_DIRECTORY, Client, functional )

print
print "entering retrieval loop"
CONTINUE_VAR = 1 ;
while CONTINUE_VAR == 1 :
  check_finished(CURRENT_DIRECTORY, Client, functional )
  print "Going to sleep ......"
  time.sleep(1800);     


