from ase.io import read
from ase.io.vasp import read_vasp
from ase.geometry import cell_to_cellpar
import numpy as np
import matplotlib as mpl
import sys
import re
import os
import json
import spglib
import fileinput
import random
import time

poscar = sys.argv[1]
settings_file = sys.argv[2]
config_file = sys.argv[3]

def spinflipper():
    
    start_time = time.time()
    
    with open(poscar,'r') as file:
        data = file.readlines()
        
    with open(settings_file,"r") as f:
        jsondata = json.load(f)
        
    with open(config_file,'r') as g:
        config_data = json.load(g)
        
        
    data[7] = 'Cartesian\n'
    
    with open(poscar, 'w') as file:
        file.writelines(data)
    
    #################################
    ## Load in and read the POSCAR ##
    #################################
    
    structure = read(poscar) # Reads POSCAR using the ase read() package
    
    ## Extract the cell lattice params, and the positions and identities of all atoms
    
    lpas = structure.get_cell() 
    lpas_f34 = np.transpose(lpas) # Need the transpose to change the way CRYSTAL defines lpas compared to VASP
    coords = structure.get_scaled_positions() # x,y,z in columns; each row corresponds to a single atom
    Z = structure.get_atomic_numbers() # Atomic number of each site
    
    alat=np.sqrt(np.sum(lpas[0,:]**2))
    blat=np.sqrt(np.sum(lpas[1,:]**2))
    clat=np.sqrt(np.sum(lpas[2,:]**2))
    
    alphalat = np.arccos(np.dot(lpas[1,:],lpas[2,:])/blat/clat)*360/2/np.pi
    betalat  = np.arccos(np.dot(lpas[0,:],lpas[2,:])/alat/clat)*360/2/np.pi
    gammalat = np.arccos(np.dot(lpas[0,:],lpas[1,:])/alat/blat)*360/2/np.pi
        
    # Establish number of atoms in cell
    
    n_atoms = len(Z)    
    
    # Identify each atomic number with the element symbol in the periodic table 
    
    Z_list = np.linspace(1,104,104)
    element_list = ['H','He','Li','Be','B','C','N','O','F','Ne','Na','Mg','Al','Si','P','S','Cl','Ar',
                    'K','Ca','Sc','Ti','V','Cr','Mn','Fe','Co','Ni','Cu','Zn','Ga','Ge','As','Se','Br',
                    'Kr','Rb','Sr','Y','Zr','Nb','Mo','Tc','Ru','Rh','Pd','Ag','Cd','In','Sn','Sb','Te',
                    'I','Xe','Cs','Ba','La','Ce','Pr','Nd','Pm','Sm','Eu','Gd','Tb','Dy','Ho','Er','Tm',
                    'Yb','Lu','Hf','Ta','W','Re','Os','Ir','Pt','Au','Hg','Tl','Pb','Bi','Po','At','Rn',
                    'Fr','Ra','Ac','Th','Pa','U','Np','Pu','Am','Cm','Bk','Cf','Es','Fm','Md','No','Lr',
                    'Rf']
    
    atom_ids = [] # Atom identities
    atom_ids2 = [] # Atom identities
    
    for i in range(0,n_atoms):
        idx = (np.abs(Z_list - Z[i])).argmin()
        atom_ids.append(element_list[idx])
        atom_ids2.append(element_list[idx])
                
    # initialize a null list
    unique_list = []
    
    # traverse for all elements
    for x in atom_ids:
        # check if exists in unique_list or not
        if x not in unique_list:
            unique_list.append(x)
    
    ox_states = []
    new_unique_species_list = []
    loc_diff_ox_states = []
    
    for i in range(0,len(unique_list)):
        
        #sing_ox_state = input("Does "+str(unique_list[i])+" have a single oxidation state in your structure? [y/n] ")
        
        if jsondata["Single oxidation states"][i]=='y':
            ox_state_i = float(jsondata["Oxidation states"][i])
            ox_states.append(ox_state_i)
            new_unique_species_list.append(unique_list[i]+str(int(ox_state_i)))
            
        if jsondata["Single oxidation states"][i]=='n':
            ox_i_list = []
            
            for j in range(0,int(jsondata["Number of different oxidation states per unique chemical species"])):
                ox_i_ = input("Enter oxidation state no. "+str(j+1) +" of "+str(unique_list[i])+" as a number (not Roman numeral). Only include sign if negative ox state: ")
                ox_i_list.append(ox_i_)
                new_unique_species_list.append(str(unique_list[i])+str(ox_i_))
                ox_states.append(float(ox_i_))
                loc_ox_i = input("Enter site numbers, as a comma-separated list, for species "+str(unique_list[i])+str(ox_i_)+": ")
                loc_diff_ox_states.extend([loc_ox_i])
                
    # If a species has multiple oxidation states, need to give them different atomic numbers, 
    # so that CRYSTAL reads them in correctly. In this case, multiples of 200 are added to the 
    # atomic number, starting from the true atomic number of the element
    
    
    for i in range(0,len(loc_diff_ox_states)):
        b = np.ndarray.tolist(np.asarray(loc_diff_ox_states[i]))
        c = b.split(",")
        d = []
        for j in range(0,len(c)):
            d.append(int(c[j]))
        for j in range(0,len(d)):
            Z[d[j]-1]=(((i)*200)+Z[d[j]-1])
            
            
    ###################################################
    ## Begin creating array of lines for INPUT files ##
    ###################################################
    
    print("Beginning INPUT file construction...")
    
    ## Header ##
    
    sys_name = jsondata["Name"]
    
    
    ## Basis sets ##
    
    elements_basis = ['H1', 'Li1', 'B3', 'C4', 'N-3', 'N3', 'N5', 'O-2', 'O-1', 'F-1', 'Na1', 'Mg2', 'Al3',
                 'Si4', 'P3', 'P5', 'S-2', 'S4', 'S6', 'Cl-1', 'K1', 'Ti2', 'Ti3', 'Ti4', 'V2', 'V3', 'V4', 'V5',
                 'Cr2', 'Cr3', 'Cr4', 'Cr5', 'Cr6', 'Mn2', 'Mn3', 'Mn4', 'Mn5', 'Mn6', 'Mn7', 'Fe2', 'Fe3', 'Fe4',
                 'Fe6', 'Co2', 'Co3', 'Co4', 'Ni2', 'Ni3', 'Ni4', 'Cu1', 'Cu2', 'Zn2', 'Ru3', 'Ru4', 'Ru5', 'Ru7',
                 'Ru8', 'Ir3', 'Ir4', 'Ir5', 'W4', 'W5', 'W6', 'end']

    Z_basis = [1,3,5,6,7,307,407,8,308,9,11,12,13,14,15,315,16,316,416,17,19,322,422,22,23,323,423,523,
               324,24,424,524,624,325,425,25,525,625,725,326,26,426,526,327,27,427,28,328,428,329,29,30,
               244,344,444,544,644,377,277,477,577,274,374,474,0]
        
    new_unique_species_basis_map = []
    
    for i in range(0,len(new_unique_species_list)):
        for j in range(0,len(elements_basis)):
            if elements_basis[j]==new_unique_species_list[i]:
                new_unique_species_basis_map.append(j)
    
    baslines = []
        
    with open('/home/euannb/scripts/basis_H_to_Zn.dat') as basfile:
        content = basfile.read()
            
    for i in range(0,len(new_unique_species_basis_map)):
        a = elements_basis[new_unique_species_basis_map[i]]+"\n"
        b = elements_basis[new_unique_species_basis_map[i]+1]
        
        # Find and validate before-part.
        pos_a = content.find(a)
        if pos_a == -1: 
            ""
        # Find and validate after-part.
        pos_b = content.rfind(b)
        if pos_b == -1:  
            ""
        # Return middle part.
        adjusted_pos_a = pos_a + len(a)
        if adjusted_pos_a >= pos_b: 
            ""
        baslines.append(content[adjusted_pos_a:pos_b])
        
    new_atom_ids = []

    newnewnew_unique_species = []
    for i in range(0,len(new_unique_species_list)):
        if len(new_unique_species_list[i])==2:
            newnewnew_unique_species.append(new_unique_species_list[i][0])
        if len(new_unique_species_list[i])==3 and new_unique_species_list[i][1]!='-':
            newnewnew_unique_species.append(new_unique_species_list[i][0:2])
        if len(new_unique_species_list[i])==3 and new_unique_species_list[i][1]=='-':
            newnewnew_unique_species.append(new_unique_species_list[i][0])
    
    
    for i in range(0,len(atom_ids)):
        for j in range(0,len(newnewnew_unique_species)):
            if atom_ids[i]==newnewnew_unique_species[j]:
                atom_ids[i]=new_unique_species_list[j]
                
    Z_ox_adjusted = []
                
    for i in range(0,len(atom_ids)):
        for j in range(0,len(elements_basis)):
            if elements_basis[j]==atom_ids[i]:
                Z_ox_adjusted.append(Z_basis[j])
                
    number_of_unique_species = []
    
    for i in range(0,int(len(new_unique_species_list))):
        number_of_unique_species.append(atom_ids.count(new_unique_species_list[i]))
                        
    # Obtain the Z, x,y,z matrix for all atoms
    
    Zxyz = []
    
    for i in range(0,int(n_atoms)):
        Zxyz.append("   "+str(Z_ox_adjusted[i])+"    "+str("{:.8f}".format(coords[i][0]))+"    "+str("{:.8f}".format(coords[i][1]))+"    "+str("{:.8f}".format(coords[i][2]))+"\n")

            
    ######################################################
    ## Define lines for everything after the basis sets ##
    ######################################################
    
    ## "Mid-footer" ##
    
    # Define level of HF for hybrid calc
    hyb = jsondata["Hybrid level"]
    
    
    # Determine, from the sizes of lattice params, the k-mesh density to achieve sampling < 0.05 A^-1
    magnitude_a = ((lpas[0][0]**2)+(lpas[0][1]**2)+(lpas[0][2]**2))**0.5
    magnitude_b = ((lpas[1][0]**2)+(lpas[1][1]**2)+(lpas[1][2]**2))**0.5
    magnitude_c = ((lpas[2][0]**2)+(lpas[2][1]**2)+(lpas[2][2]**2))**0.5
    
    n_a_exact = 1.0/(0.05*magnitude_a)
    n_b_exact = 1.0/(0.05*magnitude_b)
    n_c_exact = 1.0/(0.05*magnitude_c)
    
    n_a_round = int(np.ceil(n_a_exact)) # k-density along a
    n_b_round = int(np.ceil(n_b_exact)) # k-density along b
    n_c_round = int(np.ceil(n_c_exact)) # k-density along c
    
    # Obtain information about spins in the system
    
    dofs = np.asarray(config_data["dof"]["local_dofs"]["Cunitmagspin"]["values"])
    dofs_ = []

    for i in range(0,len(dofs)):
        dofs_.append(dofs[i][0])
    
    S_i = []
    ion_tot = 0

    for i in range(0,len(number_of_unique_species)-1):
        S_i.append(sum(dofs_[ion_tot:number_of_unique_species[i]+ion_tot]))
        ion_tot = int(ion_tot+number_of_unique_species[i])
    
    S_i.append(sum(dofs_[ion_tot:]))
    Stot = int(np.dot(S_i,jsondata["Spin of spin species"]))
    print(Stot)    
    
    
    flipped_spin_indices = []
    
    for i in range(0,int(len(config_data["dof"]["local_dofs"]["Cunitmagspin"]["values"]))):
        if float(config_data["dof"]["local_dofs"]["Cunitmagspin"]["values"][i][0])==-1.0:
            flipped_spin_indices.append(i+1)
            
    n_flipped_spins = len(flipped_spin_indices)
    flipped_spinions_for_inp = ' '.join(str(v) for v in flipped_spin_indices)
    
    ##
    
    inp_head = [
        sys_name+"\n",
        "CRYSTAL\n",
        "0 0 0\n",
        "1\n",
        str(alat)+" "+str(blat)+" "+str(clat)+" "+str(alphalat)+" "+str(betalat)+" "+str(gammalat)+"\n",
        str(int(n_atoms))+"\n"
    ]
    
    
    ##
    
    mid_DFT = [
        "99  0\n",
        "PRINT\n",
        "END\n",
        "DFT\n",
        "SPIN\n",
        "EXCHANGE\n",
        "BECKE\n",
        "CORRELAT\n",
        "LYP\n",
        "NONLOCAL\n",
        "0.9 0.81\n",
        "HYBRID\n",
        str(hyb)+"\n",
        "XLGRID\n",
        "END\n",
        "TOLINTEG\n",
        "7 7 7 7 14\n",
        "SCFDIR\n",
        "SHRINK\n",
        "0 8\n",
        str(n_a_round)+" "+str(n_b_round)+" "+str(n_c_round)+"\n"
        "FMIXING\n",
        "70\n",
        "SPINLOCK\n",
        str(Stot)+" 10\n",
        "SPINEDIT\n",
        str(n_flipped_spins)+"\n",
        str(flipped_spinions_for_inp)+"\n"        
    ]
    
    
    ## Footer ##
    
    footer = [
        "EXCHSIZE\n",
        "40751000\n",
        "BIPOSIZE\n",
        "20000000\n",
        "MAXCYCLE\n",
        "400\n",
        "TOLDEE\n",
        "7\n",
        "SAVEWF\n",
        "GUESSP\n",
        "END"
    ]
    
    ## Properties input ##
    
    props = [
        "SETPRINT\n",
        "1\n",
        "18 1\n",
        "ANISOTRO\n",
        "ALL\n",
        "ISOTROPIC\n",
        "ALL\n",
        "POTC\n",
        "0\n",
        "0\n",
        "0\n",
        "PPAN\n",
        "NEWK\n",
        "0 8\n",
        "2 2 2\n",
        "1 0\n",
        "ECH3\n",
        "100\n",
        "POT3\n",
        "50\n",
        "5\n"
        "END",
    ]
            
    with open('INPUT', 'w') as inp_file:
        
        for line in inp_head:
            inp_file.write(line)
            
        for line in Zxyz:
            inp_file.write(line)
        
        inp_file.write("ENDG\n")
            
        for line in baslines:
            inp_file.write(line)
            
        for line in mid_DFT:
            inp_file.write(line)
            
        for line in footer:
            inp_file.write(line)
            
    with open('properties.input', 'w') as prop_inp_file:
        
        for line in props:
            prop_inp_file.write(line)
            
    print("INPUT and properties.input successfully written to.")
    
    data[7] = 'Direct\n'
    
    with open(poscar, 'w') as file:
        file.writelines(data)
        
    print("--- %s seconds ---" % (time.time() - start_time))
                   
    return
    
x = spinflipper()
