#!/usr/bin/env python3

import sys
import os
import re
import math
import cmath
import numpy as np
import xyz_utils
pi=cmath.pi
fourpi=4*pi
tiny=1.e-9 # for checking equality of norm of G points


def atomicFormFactor(q,parVec): # this functional form is claimed to represent a decent fit until q=25 Angstrom^-1
# see http://lampx.tugraz.at/~hadley/ss1/crystaldiffraction/atomicformfactors/formfactors.php
    su=parVec[9] # the parameter called "c"
    minusqover4pisq=-1.*math.pow(q/fourpi,2)
    for i in range(4):
        ai=parVec[1+i*2]
        bi=parVec[2+i*2]
        su+=ai*math.exp(bi*minusqover4pisq)
#        print("quooo",su,ai,bi,minusqover4pisq)
    return su

def abs2(x):
    return x.real**2 + x.imag**2



def xyz_structure_factor(args):

    commandname=args.prog
    global debug
    debug = args.debug
    paramsfile=os.environ['HOME']+"/bin/data/atomic_form_factors_parameters"

    pararray=[]
    f = open(paramsfile, 'r')
    for line in f:
        line=re.sub(r"^\s+","",line).rstrip()
        if not re.match(r"^#",line): # ignore commented lines
            nums=re.split(r'\s+',line)
            for i in range(1,10):
                nums[i]=float(nums[i])
            pararray.append(nums)
    f.close()

#    for i in range(100):  # test of the atomicFormFactor:
#        q=0.1*i
#        print("quaaa",q,atomicFormFactor(q,pararray[22]))

    storea=[]

    filen=args.filenames[0] # xyz file
    count=0
    f = open(filen, 'r')

    fr=xyz_utils.xyz_read_one_frame(f) # read first frame from xyz file
    while True:
        fr1=xyz_utils.xyz_read_one_frame(f) # try to read more frames from xyz file
        if fr1.atoms == None:
            break
        else:
            fr=fr1
#                       eventually keep the last frame only!!
    storea=[]
    i=0
    for at in range(len(fr.atoms)):
        atom=fr.atoms[at]
        for i in range(len(pararray)): # search atom name in list
            if pararray[i][0]==atom:
                break
        if i>=len(pararray):
            print(commandname,"ERROR: atomname",atom,"not found in available list!\nSee",paramsfile,"\nEnding here, sorry", file=sys.stderr)
            sys.exit(3)
        here=[i,fr.coords[at,0],fr.coords[at,1],fr.coords[at,2]]
        storea.append(here)
    f.close()

    filen=args.filenames[1] # process file with q points
    if filen=="-":
        f=sys.stdin
    else:
        f = open(filen, 'r')

    storeq=[]
    stored=[]
    for line in f:
        line=re.sub(r"^\s+","",line).rstrip()
        nums=re.split(r'\s+',line)
        if re.match(r"^#",line) or len(nums)<2: # reprint commented or short lines without further processing
            print(line)
        else:
            q=[float(nums[0]),float(nums[1]),float(nums[2])]
            if len(nums)>4:
                iMiller="( "+nums[3]+" "+nums[4]+" "+nums[5]+" )"
            else:
                iMiller=""
            sf=0
            for at in storea:
                R=at[1:]
                expfact=np.exp(-1j*np.dot(q,R))
                qnorm=np.linalg.norm(q)
#                print("quii",q,qnorm,R,expfact)
#                print("quaa",at[0],pararray[at[0]])
                sf+=expfact*atomicFormFactor(qnorm,pararray[at[0]])
#                print("quxx",R,q,expfact,sf)
            storeq.append(q+[abs2(sf),sf,iMiller])
    if args.powder: # powder case: sum up intensities at given |q|
        for lin in storeq:
            qnorm=np.linalg.norm(lin[:3])
            newdat=1
            for k in range(len(stored)):
#                print("quiii",k,qnorm)
                if abs(stored[k][0]-qnorm)<tiny:
                    newdat=0
                    stored[k][1]+=lin[3] # sum intensities: no coherence among different rotated crystals
            if newdat==1:
                stored.append([qnorm,lin[3],lin[4],lin[5]])
# note that in stored[k][2] there remains the complex S(q) for the first G vector found
    else: # single crystal case, no averaging:
        stored=storeq
    return stored


    
# the following function is only exectuted when the code is run as a script, and its purposes is to parse
# the command line and to generate a meaningful parsed args list to the actual function doing the job:
if __name__ == "__main__":
    import argparse

    desc="""compute S(q) for a cell provided by the last frame of a xyz file
             at the q points listed in a second file, or from stdin
OUTPUT: a list of                     qx qy qz |S(q)|^2 S(q)
        or (powder sample case) a list of  |q| |S(q)|^2 S(q)
requires a file ~/bin/data/atomic_form_factors_parameters with the
parameters needed to compute the atomic form factors  
you can download this file from
http://materia.fisica.unimi.it/manini/scripts/atomic_form_factors_parameters"""

    epil="""\t\t\tv. 1.3 by Nicola Manini, 14/04/2020"""

    ##  Argument Parser definition
    parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter
                                    , description=desc, epilog=epil)

    parser.add_argument( 'filenames', nargs='*', default=['-'],
                         help='Files to be processed. If not given, stdin is used')

    parser.add_argument( '-d', '--debug', action='store_true',
                         dest='debug',
                         help='activate debug mode -- WARNING! output is affected/spoiled!' )

    parser.add_argument( '-p', action='store_true',
                         dest='powder',
                         help='powder sample (i.e. average over rotations)')
    ## End arg parser definition
#    args = argparse.parser()
    args=parser.parse_args(sys.argv[1:])
    d = vars(args)	# adding prog in args, for unknown reasons it's not there...
    d['prog']=parser.prog
    if len(args.filenames)==0:
        print(args.prog,"requires at least one xyz file", file=sys.stderr)
        exit()
    if len(args.filenames)==1:
        if args.filenames[0]=='-':
            print(args.prog,"requires at least one xyz file", file=sys.stderr)
            exit()
        else:
            args.filenames.append('-')
            print(args.prog,"warning: expecting q points from stdin", file=sys.stderr)
    
#   here the actual function doing the job is called:
    stored=xyz_structure_factor(args)
# and the results are printed out:
    for i in stored:
        print(' '.join(map(str,i)))