#!/usr/bin/env python3

# written by Thomas Barends
# v. minor edits by J. Beale + conversion to python3

"""
# aim
visualise preferential orientations of crystals processed by CrystFEL

# usage
python update-geom-from-lab6.py -f <path-to-stream-file>, -r <reindex-or-not>, -s step-size

# output
creates 3 plots showing the orientations of the crystals along the 3 axes
"""

# modules
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import argparse
import os

def ProjectionPlot(X,Y,Z,C,title):
    f=plt.figure()
    ax1=f.add_subplot(221)
    ax1.scatter(X,Z,s=0.5,c=C)
    ax1.set_title(title+',XZ plane')
    ax1.set_xlabel('X ')
    ax1.set_ylabel('Z (along beam)')
    ax1.set_aspect('equal', 'box')
    
    ax2=f.add_subplot(222,projection='3d')
    ax2.scatter(X,Z,-Y,s=0.5,c=C)
    ax2.set_xlabel('X ')
    ax2.set_ylabel('Z (along beam)')
    ax2.set_zlabel('Y (along jet)')
    ax2.set_aspect('equal', 'box')
    
    ax3=f.add_subplot(223)
    ax3.scatter(X,-Y,c=C,s=0.5)
    ax3.set_title(title+',XY plane')
    ax3.set_xlabel('X ')
    ax3.set_ylabel('Y (along jet)')
    ax3.set_aspect('equal', 'box')
    
    ax4=f.add_subplot(224)
    ax4.scatter(Z,-Y,c=C,s=0.5)
    ax4.set_title(title+',ZY plane')
    ax4.set_xlabel('Z (along beam)')
    ax4.set_ylabel('Y (along jet)')
    ax4.set_aspect('equal', 'box')

    f.canvas.manager.set_window_title(title)
    f.set_size_inches(9,7)

    plt.tight_layout()
    plt.savefig(title+'.png')
    plt.show()

def main( stream_file, reindex=False, step=25 ):
    # declare lists to read reciprocal vectors
    astar=[]
    bstar=[]
    cstar=[]

    # open stream file, read all lines and close again
    infile=open(stream_file,'r')
    inlines=infile.readlines()
    infile.close()

    # loop over lines and extract reciprocal cell axes 
    for line in inlines:
        if 'astar' in line:
            elements=line.split()
            vector=[float(elements[2]),float(elements[3]),float(elements[4])]
            astar.append(vector)
        
        if 'bstar' in line:
            elements=line.split()
            vector=[float(elements[2]),float(elements[3]),float(elements[4])]
            bstar.append(vector)

        if 'cstar' in line:
            elements=line.split()
            vector=[float(elements[2]),float(elements[3]),float(elements[4])]
            cstar.append(vector)

    # turn all into numpy arrays
    astararray=np.asarray(astar)
    bstararray=np.asarray(bstar)
    cstararray=np.asarray(cstar)

    # get number of cells from shape of array
    s=np.shape(astararray)
    numcells=s[0]

    # declare lists for real-space cell vector directions (i.e. normalized vectors)
    adir=[]
    bdir=[]
    cdir=[]

    # declare lists for vector lengths
    alen=[]
    blen=[]
    clen=[]

    # loop over all cells
    for n in range(numcells):
        astar=astararray[n,:]
        bstar=bstararray[n,:]
        cstar=cstararray[n,:]
        Vstar=np.dot( astar , np.cross( bstar , cstar )  )    # reciprocal cell volume
        a=(np.cross(bstar,cstar) / Vstar )                    # calculate real-space a
        b=(np.cross(cstar,astar) / Vstar )                    # calculate real-space b
        c=(np.cross(astar,bstar) / Vstar )                    # calculate real-space c

        al=np.linalg.norm(a)                                  # calculate lengths 
        bl=np.linalg.norm(b)
        cl=np.linalg.norm(c)

        if al<bl and reindex==True:
            newa=-1.0*b
            newb=a
            newc=c
            
            newal=bl
            newbl=al
            newcl=cl
            
            a=newa
            b=newb
            c=newc
            al=newal
            bl=newbl
            cl=newcl 

        adir.append(a/al)                                     # put normalized vectors in lists
        bdir.append(b/bl)
        cdir.append(c/cl)
        
        alen.append(al)                                       # put lengths in lists
        blen.append(bl)
        clen.append(cl)
        

    adir=np.asarray(adir)                                     # turn lists into numpy arrays
    bdir=np.asarray(bdir)
    cdir=np.asarray(cdir)

    alen=np.asarray(alen)
    blen=np.asarray(blen)
    clen=np.asarray(clen)

    f,((ax1,ax2,ax3))=plt.subplots(1,3)

    ax1.hist(10.*alen,int(5*np.sqrt(len(alen))),density=True,facecolor='green',edgecolor='green')
    atitle=r'$\mathit{a}$'
    ax1.set_title(atitle)

    ax2.hist(10.*blen,int(5*np.sqrt(len(blen))),density=True,facecolor='green',edgecolor='green')
    btitle=r'$\mathit{b}$'
    ax2.set_title(btitle)

    ax3.hist(10.*clen,int(5*np.sqrt(len(clen))),density=True,facecolor='green',edgecolor='green')
    ctitle=r'$\mathit{c}$'
    ax3.set_title(ctitle)

    f.canvas.manager.set_window_title('Unit Cell Parameter Histograms')
    f.set_size_inches(10,5)
    plt.tight_layout()
    plt.show()

    # plot a, coloured by length of a (normalized to maximum length observed)
    print('I will plot every',step,'th point,',len(adir[::step,0]),'points in total...')

    ProjectionPlot(adir[::step,0],adir[::step,1],adir[::step,2],(alen[::step]/np.max(alen)),'a_axis_direction')

    ProjectionPlot(bdir[::step,0],bdir[::step,1],bdir[::step,2],(blen[::step]/np.max(blen)),'b_axis_direction')

    ProjectionPlot(cdir[::step,0],cdir[::step,1],cdir[::step,2],(clen[::step]/np.max(clen)),'c_axis_direction')

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-f",
        "--stream_file",
        help="stream file of interest for assessment of crystal orientations",
        type=os.path.abspath
    )
    parser.add_argument(
        "-r",
        "--reindex",
        help="do you need to adjust the cell parameters because CrystFEL has set a > b",
        type=bool,
        default=False
    )
    parser.add_argument(
        "-s",
        "--step",
        help="how many points should i plot - every 1/10, 1/25?",
        type=int,
        default=25
    )
    args = parser.parse_args()
    # run main
    main( args.stream_file, args.reindex, args.step )