#!/usr/bin/env python3
import os
import shutil
import json
import numpy as np

def make_calypso_input(nameofatoms,numberofatoms,
                           numberofformula,volume,
                           distanceofion,psoratio,popsize,
                           maxstep,icode,split,vsc,
                           maxnumatom,ctrlrange,pstress,fmax):
    ret = "################################ The Basic Parameters of CALYPSO ################################\n"
    ret+= "# A string of one or several words contain a descriptive name of the system (max. 40 characters).\n"
    assert (nameofatoms != None )
    ret+= "SystemName = %s\n"%(''.join(nameofatoms))
    ret+= "# Number of different atomic species in the simulation.\n"
    ret+= "NumberOfSpecies = %d\n"%(len(nameofatoms))
    ret+= "# Element symbols of the different chemical species.\n"
    ret+= "NameOfAtoms = %s\n"%(' '.join(nameofatoms))
    ret+= "# Number of atoms for each chemical species in one formula unit. \n"
    assert numberofatoms != None  or len(numberofatoms) == len(nameofatoms)
    ret+= "NumberOfAtoms = %s\n"%(' '.join(list(map(str,numberofatoms))))
    ret+= "# The range of formula unit per cell in your simulation. \n"
    assert numberofformula != None  or len(numberofformula) == 2 or type(numberofformula) == type([1,2])
    ret+= "NumberOfFormula = %s\n"%(' '.join(list(map(str,numberofformula))))
    ret+= "# The volume per formula unit. Unit is in angstrom^3.\n"
    ret+= "Volume = %s\n"%(volume[0])
    ret+= "# Minimal distance between atoms of each chemical species. Unit is in angstrom.\n"
    assert len(distanceofion) == len(nameofatoms) #"check distance of ions and the number of atoms"
    assert len(distanceofion[0]) == len(nameofatoms)
    ret+= "@DistanceOfIon \n"
    for temp in distanceofion:
        ret+="%4s \n"%(' '.join(list(map(str,temp))))
    ret+= "@End\n"
    ret+= "# It determines which algorithm should be adopted in the simulation.\n"
    ret+= "Ialgo = 2\n"
    ret+= "# Ialgo = 1 for Global PSO\n"
    ret+= "# Ialgo = 2 for Local PSO (default value)\n"
    ret+= "# The proportion of the structures generated by PSO.\n"
    assert (0 <=  psoratio[0] <= 1  )
    ret+= "PsoRatio = %s\n"%(psoratio[0])
    ret+= "# The population size. Normally, it has a larger number for larger systems.\n"
    assert popsize[0] != None  or popsize!=None or type(popsize) == type([0,1]) or type(popsize[0]) == type(0)
    ret+= "PopSize = %d\n"%(popsize[0])
    assert maxstep[0] != None  or maxstep!=None or type(maxstep) == type([0,1]) or type(maxstep[0]) == type(0)
    ret+= "# The Max step for iteration\n"
    ret+= "MaxStep = %d\n"%(maxstep[0])
    ret+= "#It determines which method should be adopted in generation the random structure. \n"                             
    ret+= "GenType= 1 \n"
    ret+= "# 1 under symmetric constraints\n"
    ret+= "# 2 grid method for large system\n"
    ret+= "# 3 and 4 core grow method \n"
    ret+= "# 0 combination of all method\n"
    ret+= "# If GenType=3 or 4, it determined the small unit to grow the whole structure\n"
    ret+= "# It determines which local optimization method should be interfaced in the simulation.\n"
    assert icode != None  or type(icode) == type([0,1]) or type(icode[0]) == type(0)
    ret+= "ICode= %d\n"%(icode[0])
    ret+= "# ICode= 1 interfaced with VASP\n"
    ret+= "# ICode= 2 interfaced with SIESTA\n"
    ret+= "# ICode= 3 interfaced with GULP\n"
    ret+= "# The number of lbest for local PSO\n"
    ret+= "NumberOfLbest=4\n"
    ret+= "# The Number of local optimization for each structure.\n"
    ret+= "NumberOfLocalOptim= 3\n"
    ret+= "# The command to perform local optimiztion calculation (e.g., VASP, SIESTA) on your computer.\n"
    ret+= "Command = sh submit.sh\n"
    ret+= "MaxTime = 9000 \n"
    ret+= "# If True, a previous calculation will be continued.\n"
    ret+= "PickUp = F\n"
    ret+= "# At which step will the previous calculation be picked up.\n"
    ret+= "PickStep = 1\n"
    ret+= "# If True, the local optimizations performed by parallel\n"
    ret+= "Parallel = F\n"
    ret+= "# The number node for parallel \n"
    ret+= "NumberOfParallel = 4\n"
    assert split != None 
    ret+= "Split = %s\n"%(split)
    assert pstress != None  or type(pstress) == type([200])
    ret+= "PSTRESS = %f\n"%(pstress[0])
    assert fmax != None  or type(fmax) == type([200])
    ret+= "fmax = %f\n"%(fmax[0])
    ret+= "################################ End of The Basic Parameters of CALYPSO #######################\n"
    if vsc == 'T':
        assert len(ctrlrange) == len(nameofatoms) #'check distance of ions and the number of atoms'
        ret+= "##### The Parameters For Variational Stoichiometry  ##############\n"
        ret+= "## If True, Variational Stoichiometry structure prediction is performed\n"
        ret+= "VSC = %s\n"%(vsc)
        ret+= "# The Max Number of Atoms in unit cell\n"
        ret+= "MaxNumAtom = %s\n"%(maxnumatom[0])
        ret+= "# The Variation Range for each type atom \n"
        ret+= "@CtrlRange\n"                                                                                                      
        for ttemp in ctrlrange:
            ret+="%4s \n"%(' '.join(list(map(str,ttemp))))
        ret+= "@end\n"
        ret+= "###################End Parameters for VSC ##########################\n"
    return ret

def _make_model_devi_buffet(jdata,calypso_run_opt_path):

    calypso_input_path = jdata.get('calypso_input_path')
    shutil.copyfile(os.path.join(calypso_input_path,'input.dat'),os.path.join(calypso_run_opt_path, 'input.dat'))

def _make_model_devi_native_calypso(iter_index,model_devi_jobs, calypso_run_opt_path):

    for iiidx, jobbs in enumerate(model_devi_jobs):
        if iter_index in jobbs.get('times'):
            cur_job = model_devi_jobs[iiidx]

    work_path = os.path.dirname(calypso_run_opt_path)
    with open(os.path.join(work_path, 'cur_job.json'), 'w') as outfile:
        json.dump(cur_job, outfile, indent = 4)

    # Crystal Parameters
    nameofatoms = cur_job.get('NameOfAtoms')
    numberofatoms = cur_job.get('NumberOfAtoms')
    numberofformula = cur_job.get('NumberOfFormula',[1,1])
    volume = cur_job.get('Volume')
    distanceofion = cur_job.get('DistanceOfIon')
    psoratio = cur_job.get('PsoRatio')
    popsize = cur_job.get('PopSize')
    maxstep = cur_job.get('MaxStep')
    icode = cur_job.get('ICode',[1])
    split = cur_job.get('Split','T')
    # VSC Control
    maxnumatom = None
    ctrlrange = None
    vsc = cur_job.get('VSC','F')
    if vsc == 'T':
        maxnumatom = cur_job.get('MaxNumAtom')
        ctrlrange = cur_job.get('CtrlRange')

    # Optimization
    pstress = cur_job.get('PSTRESS',[0.001])
    fmax = cur_job.get('fmax',[0.01])

    # Cluster

    # 2D

    file_c = make_calypso_input(nameofatoms,numberofatoms,
                               numberofformula,volume,
                               distanceofion,psoratio,popsize,
                               maxstep,icode,split,vsc,
                               maxnumatom,ctrlrange,pstress,fmax)
    with open(os.path.join(calypso_run_opt_path, 'input.dat'), 'w') as cin :
        cin.write(file_c)

def write_model_devi_out(devi, fname):
    assert devi.shape[1] == 8
    #assert devi.shape[1] == 7
    header = '%5s' % 'step'
    for item in 'vf':
        header += '%16s%16s%16s' % (f'max_devi_{item}', f'min_devi_{item}',f'avg_devi_{item}')
    header += '%16s'%str('min_dis')
    np.savetxt(fname,
               devi,
               fmt=['%5d'] + ['%17.6e' for _ in range(7)],
               delimiter='',
               header=header)
    return devi

