#####################################################
#                                                   #
# This program enables a user to generate a RIXS    #
# plane macro that moves all needed DAVES motors    #
# for each desired emission energy point. The       #
# incident energy scan parameters are defined at    #
# the beginning and then run at each emission       #
# energy.  For the DAVES motor positions I copied   #
# wholesale the code from fscanMaker.py v 0.2.1;    #
# eventually this codes should be unified into a    #
# single scanMaker script.  For now, the user must  #
# manually insert sample position changes; this     #
# functionality will hopefully be added later.      #
# I'm sure this code is ugly and will need to be    #
# rewritten later, but for now it works.            #
#                                                   #
# v 0.4 - 2023.04.06                                #
#   - simplified script name for easier future      #
#   management                                      #
# v 0.3 - 2022.02.24                                #
#   - added the option to read scan parameters from #
#   a file populated by the BL scientist,           #
#   streamling the creation of scans for the same   #
#   experiment                                      #
#   - moved some code around such that this script  #
#   should now essentially mirror fscanMaker v 0.5  #
# v 0.2 - 2022.02.22                                #
#   - changed the scale factor to be adjustable     #
#   by the user                                     #
# v 0.1 - 2020.03.30                                #
#   - initial iteration; most code stolen from      #
#   fscanMaker.py                                   #
#                                                   #
# CJP, CHESS                                        #
#####################################################

import numpy as np

def main():
    # The following code will eventually be used in a unified scanMaker script for allowing the user
    # to choose what kind of scan file they want to create.  With any luck I'll get to that soon...

    #print("\nWhat kind of scan file do you want to generate?\n")
    #print("[1] XES fscan")
    #print("[2] RIXS macro")
    #print("[3] Exit\n")

    #choice = "0"
    #while choice != "1" and choice != "2":
    #    choice = input("Choose one:  ")
    #    if choice == "1":
    #        print("\nXES fscan selected")
    #    elif choice == "2":
    #        print("\nRIXS macro selected")
    #    else:
    #        exit()

    # Beginning of script where the code either reads in pre-saved values from the _VALUES file or
    # solicits them from the users

    print("#################################################################################\n")
    print("This is a script that will generate a RXES scan for the ID2A beamline at CHESS.")
    print("It is likely that most of the alignment data you need has been loaded into the \n'_VALUES' file and so you can answer 'yes' to the first question and only need \nto input your energy range(s) and number(s) of steps.\n")
    print("v 0.3, 2022.02.24")
    print("\n#################################################################################\n")

    readValues = input("Do you want to read saved paramaters? (y/n)  ")
    if readValues == "y":
        valueDict = {}
        print("\nValues read from file:\n")
        with open("_VALUES.txt") as g:
            valueLines = g.readlines()
            for a in valueLines:
                vLine = a.split("=")
                valueDict[vLine[0]] = vLine[1].strip()
        scan = valueDict['ESCAN'] + "\n"
        xtal = valueDict['XTAL']
        h = int(valueDict['H'])
        k = int(valueDict['K'])
        l = int(valueDict['L'])
        radius = int(valueDict['R'])
        centerE = float(valueDict['E_ALIGN'])
        centerAnaAngle = float(valueDict['ANA'])
        centerDetAngle = float(valueDict['DET'])
        updown = valueDict['UPDN']
        scale = float(valueDict['SCALE'])
        print("Crystal analyzers = " + xtal + "(" + str(h) + str(k) + str(l) + ")")
        print("Crystal radius = " + str(radius))
        print("\nSpectrometer aligned at " + str(centerE) + " eV with the " + updown + " analyzers at " + str(centerAnaAngle) + " and detector at " + str(centerDetAngle))
        print("Theta / 2-theta scaling factor = " + str(scale) +"\n")
        print("The defined energy scan is: " + scan)

    else:
        # Here we make sure the user doesn't try anything funny like making the high energy < the low energy
        energyCheck = 0
        while energyCheck == 0:
            lowE = float(input("Beginning of XAS scan in eV:  "))
            highE = float(input("End of XAS scan in eV:  "))
            if lowE >= highE:
                print("The high energy must be > the low energy value!  Try again!\n")
            else:
                energyCheck = 1

        # Assume the user wants a normal Escan unless they specify otherwise
        moveUnd = input("\nMove undulator along with mono (y/n)?  ")
        if moveUnd == "y":
            scanType = "EscanU"
            print("\nEscanU selected")
        else:
            scanType = "Escan"  
            print("\nEscan selected")
        
        points = int(input("\nHow many points should the scan have?  "))
        time = int(input("How many seconds per point?  "))

        scan = scanType + " " + str(lowE/1000) + " " + str(highE/1000) + " " + str(points) + " " + str(time) + "\n"
        print("\nThe scan you created is:  " + scan)

    #-------------------   Here is where the fscanMaker code starts   -------------------#

    if readValues != "y":
        xtal = "nothing"
        while (xtal != "Si" and xtal != "Ge"):               # While loop ensures user inputs a valid material
            xtal = input("Crystal material (Si or Ge):  ")   # Input for type of analyzer crystal
            if xtal == "Si":
                lat = 5.43102
                print("Silicon; a = " + str(lat) + "\n")
            elif xtal == "Ge":
                lat = 5.65782
                print("Germanium; a = " + str(lat) + "\n")
            else:
                print("Invalid material! Must be either Si or Ge!\n")
                xtal = "nothing"
    else:
        if xtal == "Si":
            lat = 5.43102
        elif xtal == "Ge":
            lat = 5.65782

    if readValues != "y":
        reflection = 0
        while (reflection == 0):                             # This block makes sure the inputs make sense and the reflection is allowed
            hString = input("h:  ")                          # Inputs for h,k,l values
            kString = input("k:  ")
            lString = input("l:  ")

            if (hString == "" or kString == "" or lString == ""):   # If the user doesn't input a value, send them back
                print("Must fill in every index!\n")
                continue

            h = int(hString)                                # Convert h,k,l to integers
            k = int(kString)
            l = int(lString)

            m = (h+k+l-2) % 4                               # Check for h + k + l =/= 4n + 2
            a = (h % 2) + (k % 2) + (l % 2)                 # Check for h,k,l being all even or odd

            if m == 0:                                      # If reflection isn't allowed, sends user back to the beginning
                reflection = 0
                print("Forbidden reflection!  h + k + l = 4m + 2\n")
            elif a == 1:
                reflection = 0
                print("Forbidden reflection!  h,k,l must be all even or odd!\n")
            elif a == 2:
                reflection = 0
                print("Forbidden reflection!  h,k,l must be all even or odd!\n")
            elif (h == 0 and k == 0 and l == 0):
                reflection = 0
                print("h = k = l = 0 is not valid!\n")
            else:
                reflection = 1
                print("Reflection ok\n")

    # Print out 2d regardless of the manner of input for hkl
    d = lat / np.sqrt(h*h + k*k + l*l)
    x = np.round_(d, decimals=6)                    # Need to convert d to a decimal so it can be rounded to a reasonable number of sig figs
    print("2d = ",2*x," A")

    if readValues != "y":
        radius = 0
        while (radius != 850 or radius != 1000):        # Get the bend radius from the user and, if it's not a standard value, make sure they really want it
            radiusString = input("\nBend radius in mm (850 or 1000):  ")    
            radius = int(radiusString)
            if (radius == 850 or radius == 1000):
                break
            else:
                radCheck = input("The value you entered isn't a starndard radius, are you sure you want to continue (y/n)?  ")
                if radCheck == "y":
                    break
                else:
                    radius = 0

        # Check whether user is using the up or downstream spectrometers
        # This is modified slight from fscanMaker since no header is needed here
        updown = "nothing"
        while updown == "nothing":
            updown = input("Are you using the upstream or downstream spectrometers? (up/dn)  ")
            if updown == "up" or updown =="dn":
                print("\nYou have selected the " + updown + " analyzers")
            else:
                updown = "nothing"
                print('Invalid choice!  Must input either "up" or "down"!\n')

        # Can include a scaling factor between the analyzer / detector arms to modify the theta / 2-theta relationship
        scale = 2
        scaleQ = input("Is there a scaling factor needed for the analyzer / detector movement (y/n)?  ")
        if scaleQ == "y":
            scale = float(input("\nEnter the theta/2-theta scaling factor (default is 2):  "))
        else:
            pass
    else:
        if updown == "up":
            head = '#M xesuana xesudet up1rot up2rot up4rot up5rot'
        elif updown == "dn":
            head = '#M xesdana xesddet dn1rot dn2rot dn4rot dn5rot'

    if readValues != "y":
        # Get from the user energy they have aligned to and check that this makes sense given the reflection
        wavelengthOverTwod = 0
        while (wavelengthOverTwod > 1 or wavelengthOverTwod < 0.9396):
            centerEstring = input("What energy, in eV, are you aligning to?  ")
            centerE = float(centerEstring)
            centerLambda = 12398.4 / centerE
            wavelengthOverTwod = centerLambda / (2*d)
            if wavelengthOverTwod > 1:
                print("This energy is too low to be accessible with this reflection!\n")
            elif wavelengthOverTwod < 0.9396:
                print("This energy requires a Bragg angle < 70 degrees, so this is a poor reflection to use.\n")
            else:
                thetaBragg = np.degrees(np.arcsin(centerLambda / (2*d)))

        # Now get the DAVES analyzer and detector angle that this energy corresponds to
        centerAnaAngleString = input("What angle is the analyzer array at for this energy?  ")
        centerAnaAngle = float(centerAnaAngleString)
        centerDetAngleString = input("What angle is the detector arm at for this energy?  ")
        centerDetAngle = float(centerDetAngleString)

        # Regurgitate these values as a sanity check
        print("\nOk, so an energy of " + str(centerE) + " eV corresponds to DAVES positions of ana = " + str(centerAnaAngle) + " and det = " + str(centerDetAngle) + ".")
    else:
        centerLambda = 12398.4 / centerE
        wavelengthOverTwod = centerLambda / (2*d)
        thetaBragg = np.degrees(np.arcsin(centerLambda / (2*d)))
    
    print("This energy corresponds to a Bragg angle of " + str(np.round_(thetaBragg,4)) + " degrees.\n")

    # Check if the user wants to do things in degrees or energy
    units = "nothing"
    while (units != "a" and units != "e"):
        units = str(input("Do you want to define a scan in terms of energy (e) or angle (a)?  "))
        if units == "e":
            print("\nScan will be defined in terms of ENERGY and take equally-spaced ENERGY steps\n")
        elif units == "a":
            print("\nScan will be defined in terms of ANGLE and take equally-spaced ANGLE steps\n")
        else:
            print('\nInvalid choice!  Must enter "e" for energy or "a" for angle.  Try again.\n')
            units = "nothing"

    # Solicit number of regions the user wants to use
    regions = 0
    while regions < 1:
        regions = int(input("How many regions do you want to use?  "))
        currentRegion = 1
        if regions < 1:
            print("Number of regions must be > 0!\n")
        else:
            continue

    # This is the main loop that will solicit desired energy / angle ranges, check their validity, solicit number of points, and then assemble scan file
    positionArray = []                  # Need to define this early otherwise it gets overwritten with the final region
    while currentRegion <= regions:
        print("For region " + str(currentRegion) + ": \n")
        if units == "a":
            lowEnergyBraggAngle = 0
            highEnergyBraggAngle = 0
            # The following 2 while loops make sure that both the upper and lower desired angles are between 70 - 90 degrees
            while (lowEnergyBraggAngle > 90 or lowEnergyBraggAngle < 70):
                highAnaAngleString = input("What is the highest analyzer angle you want to scan in this region?  ")
                highAnaAngle = float(highAnaAngleString)
                lowEnergyBraggAngle = thetaBragg + 0.5*(highAnaAngle - centerAnaAngle)
                if lowEnergyBraggAngle > 90:
                    print("The desired high angle is too high!  Choose a smaller upper angle.\n")
                elif lowEnergyBraggAngle < 70:
                    print("The desired high angle is < 70 and thus this is a poor reflection!  Choose a bigger upper angle.\n")
                else:
                    continue
            while (highEnergyBraggAngle > 90 or highEnergyBraggAngle < 70):
                lowAnaAngleString = input("What is the lowest analyzer angle you want to scan in this region?  ")
                lowAnaAngle = float(lowAnaAngleString)
                highEnergyBraggAngle = thetaBragg + 0.5*(lowAnaAngle - centerAnaAngle)
                if highEnergyBraggAngle > 90:
                    print("The desired low is too high!  Choose a smaller lower angle.\n")
                elif highEnergyBraggAngle < 70:
                    print("The desired low angle is < 70 and thus this is a poor reflection!  Choose a bigger lower angle.\n")
                else:
                    continue
            
            # Here we find out how many steps the user wants and convert that into angle increments
            stepsString = input("How many steps do you want to take in this region?  ")
            steps = int(stepsString)
            angleRange = highAnaAngle - lowAnaAngle
            phi = angleRange / steps
            print("\nThis corresponds to angular steps of " + str(np.round_(phi,3)) + " degrees per step.\n")

            currentAnaAngle = highAnaAngle
            currentDetAngle = centerDetAngle + 2*(highAnaAngle - centerAnaAngle)
            currentBraggAngle = thetaBragg + 0.5*(highAnaAngle - centerAnaAngle)

            index = 0
            for index in range(0,steps):
                motorRow = []
                currentA1 = np.degrees(np.arcsin(126/(radius*(np.sin(np.radians(currentBraggAngle))**2))))
                currentA2 = np.degrees(np.arcsin(250/(radius*(np.sin(np.radians(currentBraggAngle))**2))))
                currentD1 = (radius*np.sin(np.radians(currentBraggAngle))**2)*(1-np.sqrt(1-(126/(radius*np.sin(np.radians(currentBraggAngle))**2))**2))
                currentD2 = (radius*np.sin(np.radians(currentBraggAngle))**2)*(1-np.sqrt(1-(250/(radius*np.sin(np.radians(currentBraggAngle))**2))**2))
                #print(str(currentBraggAngle))
                #print(str(currentAnaAngle))
                motorRow.append(currentAnaAngle)
                motorRow.append(currentDetAngle * (scale / 2))
                motorRow.append(currentA2)
                motorRow.append(currentA1)
                motorRow.append(currentA1)
                motorRow.append(currentA2)
                #motorRow.append(-currentD2)
                #motorRow.append(-currentD1)
                #motorRow.append(-currentD1)
                #motorRow.append(-currentD2)
                positionArray.append(motorRow)
                currentBraggAngle = currentBraggAngle - 0.5*phi
                currentAnaAngle = currentAnaAngle - phi
                currentDetAngle = currentDetAngle - 2*phi
                index = index + 1

            currentRegion = currentRegion + 1

        elif str(units) == "e":
            lowEOver2d = 0
            highEOver2d = 0
            # The following 2 while loops make sure that both the lower and upper desired energies have Bragg angles between 70 - 90 degrees
            while (lowEOver2d > 1 or lowEOver2d < 0.9396):
                lowEnergyInputString = input("What is the lowest energy you want to scan in this region?  ")
                lowEnergyInput = float(lowEnergyInputString)
                lowEOver2d = 12398.4/(2*d*lowEnergyInput)
                if lowEOver2d > 1:
                    print("Desired low energy is too low to be reached with this reflection!  Choose a higher energy.\n")
                elif lowEOver2d < 0.9396:
                    print("Desired low energy needs a Bragg angle < 70 degrees, so this is a poor reflection.  Choose a lower energy.\n")
                else:
                    lowEnergyBraggAngle = np.degrees(np.arcsin(12398.4/(2*d*lowEnergyInput)))
                    highAnaAngle = centerAnaAngle + 2*(lowEnergyBraggAngle - thetaBragg)
            while (highEOver2d > 1 or highEOver2d < 0.9396):
                highEnergyInputString = input("What is the highest energy you want to scan in this region?  ")
                highEnergyInput = float(highEnergyInputString)
                highEOver2d = 12398.4/(2*d*highEnergyInput)
                if highEOver2d > 1:
                    print("Desired high energy is too low to be reached with this reflection!  Choose a higher energy.\n")
                elif highEOver2d < 0.9396:
                    print("Desired high energy needs a Bragg angle < 70 degrees, so this is a poor reflection.  Choose a lower energy.\n")
                else:
                    highEnergyBraggAngle = np.degrees(np.arcsin(12398.4/(2*d*highEnergyInput)))
                    lowAnaAngle = centerAnaAngle + 2*(highEnergyBraggAngle - thetaBragg)

            # Here we find out how many steps the user wants and convert that into energy increments
            stepsString = input("How many steps do you want to take in this region?  ")
            steps = int(stepsString)
            energyRange = highEnergyInput - lowEnergyInput
            dE = energyRange / steps
            print("\nThis corresponds to angular steps of " + str(np.round_(dE,3)) + " eV per step.\n")

            currentEnergy = lowEnergyInput
            currentAnaAngle = highAnaAngle
            currentDetAngle = centerDetAngle + 2*(highAnaAngle - centerAnaAngle)
            currentBraggAngle = thetaBragg + 0.5*(highAnaAngle - centerAnaAngle)

            index = 0
            for index in range(0,steps):
                motorRow = []
                currentA1 = np.degrees(np.arcsin(126/(radius*(np.sin(np.radians(currentBraggAngle))**2))))
                currentA2 = np.degrees(np.arcsin(250/(radius*(np.sin(np.radians(currentBraggAngle))**2))))
                currentD1 = (radius*np.sin(np.radians(currentBraggAngle))**2)*(1-np.sqrt(1-(126/(radius*np.sin(np.radians(currentBraggAngle))**2))**2))
                currentD2 = (radius*np.sin(np.radians(currentBraggAngle))**2)*(1-np.sqrt(1-(250/(radius*np.sin(np.radians(currentBraggAngle))**2))**2))
                #print(str(currentBraggAngle))
                #print(str(currentAnaAngle))
                motorRow.append(currentAnaAngle)
                motorRow.append(currentDetAngle * (scale / 2))
                motorRow.append(currentA2)
                motorRow.append(currentA1)
                motorRow.append(currentA1)
                motorRow.append(currentA2)
                #motorRow.append(-currentD2)
                #motorRow.append(-currentD1)
                #motorRow.append(-currentD1)
                #motorRow.append(-currentD2)
                positionArray.append(motorRow)

                #need to convert dE into dTheta
                currentEnergy = currentEnergy + dE
                currentEnergyBraggAngle = np.degrees(np.arcsin(12398.4/(2*d*currentEnergy)))
                dThetaBragg = currentBraggAngle - currentEnergyBraggAngle

                currentBraggAngle = currentEnergyBraggAngle
                currentAnaAngle = currentAnaAngle - 2*dThetaBragg
                currentDetAngle = currentDetAngle - 4*dThetaBragg
                index = index + 1

            currentRegion = currentRegion + 1

        else:
            print("Invalid choice!")
    
    position = np.array(positionArray)
    positions = np.around(position, decimals = 4)
    rows, cols = positions.shape
    lineNum = 0

    filename = input("Name for file:  ") + ".txt"
    if filename == ".txt":
        filename = "filename.txt"

    # Write everything out to a file, specifying whether the up/dn analyzers are being used
    with open(filename,'w') as f:
        while lineNum < rows:
            if updown == "up":
                line = str("umv xesuana " + str(positions[lineNum,0]) + " xesudet " + str(positions[lineNum,1]) + " up1rot " + str(positions[lineNum,2]) + " up2rot " + str(positions[lineNum,3]) + " up4rot " + str(positions[lineNum,4]) + " up5rot " + str(positions[lineNum,5]) + "\n")
                print(line)
                f.write(line)
                print(scan)
                f.write(scan)
                lineNum += 1
            elif updown == "dn":
                line = str("umv xesdana " + str(positions[lineNum,0]) + " xesddet " + str(positions[lineNum,1]) + " dn1rot " + str(positions[lineNum,2]) + " dn2rot " + str(positions[lineNum,3]) + " dn4rot " + str(positions[lineNum,4]) + " dn5rot " + str(positions[lineNum,5]) + "\n")
                print(line)
                f.write(line)
                print(scan)
                f.write(scan)
                lineNum += 1
            else:
                f.write("This shouldn't be possible, but file creation failed!")
    
    print("\nFile created, but need to MANUALLY insert sample position changes!\n")

    input("Press enter to exit...")

if __name__ == "__main__":
    main()