Beamline_Scan_File_Generator.py

from pyscicat.client import ScicatClient
import json
import sys
import os
import re
from pathlib import Path

# Currently configured to run on local host
# separate to another file
ENERGY_SCAN_TEMPLATE = {
    "XAS-4.0.2": {
        "header": "Flying Beamline Energy (763, 813, 0.1, 0)",
        "subheader": ["Hx", "EPU Polarization", "Z", "Theta", "Temp B"],
        "defaults": {
            "element_energy_ranges_prefix": "flying",
            "hx": "0",
            "epu": ["100", "190"],
        },
        "element_energy_ranges": {
            "O": "(533.1, 553.1, 0.1, 0)",
            "Al": "(1549.6, 1569.6, 0.1, 0)",
            "Si": "(1829, 1839, 0.1, 0)",
            "Ti": "(433.8, 480.2, 0.1, 0)",
            "V": "(492.1, 539.8, 0.1, 0)",
            "Cr": "(554.1, 593.8, 0.1, 0)",
            "Mn": "(618.7, 669.9, 0.1, 0)",
            "Fe": "(686.8, 739.9, 0.1, 0)",
            "Co": "(758.1, 823.2, 0.1, 0)",
            "Ni": "(832.7, 890.0, 0.1, 0)",
            "Cu": "(912.7, 972.3, 0.1, 0)",
        }
    },
    "XLD-4.0.2": {
        "header": "Flying Beamline Energy (763, 813, 0.1, 0)",
        "subheader": ["Hx", "EPU Polarization", "Z", "Theta", "Temp B"],
        "defaults": {
            "element_energy_ranges_prefix": "flying",
            "hx": "0",
            "epu": ["100", "190"],
        },
        "element_energy_ranges": {
            "O": "(533.1, 553.1, 0.1, 0)",
            "Al": "(1549.6, 1569.6, 0.1, 0)",
            "Si": "(1829, 1839, 0.1, 0)",
            "Ti": "(433.8, 480.2, 0.1, 0)",
            "V": "(492.1, 539.8, 0.1, 0)",
            "Cr": "(554.1, 593.8, 0.1, 0)",
            "Mn": "(618.7, 669.9, 0.1, 0)",
            "Fe": "(686.8, 739.9, 0.1, 0)",
            "Co": "(758.1, 823.2, 0.1, 0)",
            "Ni": "(832.7, 890.0, 0.1, 0)",
            "Cu": "(912.7, 972.3, 0.1, 0)",
        }
    },
    "XMCD-4.0.2": {
        "header": "Flying Beamline Energy (763, 813, 0.1, 0)",
        "subheader": ["Hx", "EPU Polarizatoin", "Z", "Theta", "Temp B"],
        "defaults": {
            "element_energy_ranges_prefix": "flying",
            "epu": ["-0.9", "+0.9"],
        },
        "element_energy_ranges": {
            "O": "(533.1, 553.1, 0.1, 0)",
            "Al": "(1549.6, 1569.6, 0.1, 0)",
            "Si": "(1829, 1839, 0.1, 0)",
            "Ti": "(433.8, 480.2, 0.1, 0)",
            "V": "(492.1, 539.8, 0.1, 0)",
            "Cr": "(554.1, 593.8, 0.1, 0)",
            "Mn": "(618.7, 669.9, 0.1, 0)",
            "Fe": "(686.8, 739.9, 0.1, 0)",
            "Co": "(758.1, 823.2, 0.1, 0)",
            "Ni": "(832.7, 890.0, 0.1, 0)",
            "Cu": "(912.7, 972.3, 0.1, 0)",
        },
    },
    "XAS-6.3.1": {
        "header": "Flying Energy (763, 813, 0.1, 0)",
        "subheader": ["Theta", "Y_samples", "Z_samples", "Magnet Field", "Polarization"],
        "defaults": {
            "magnetic_field_631": "0",
            "element_energy_ranges_prefix": "flying",
            "y_samples": "0",
            "polarization": "-1",
        },
        "element_energy_ranges": {
            "O": "(533.1, 553.1, 0.1, 1)",
            "Al": "(1549.6, 1569.6, 0.1, 1)",
            "Si": "(1829, 1839, 0.1, 1)",
            "Ti": "(433.8, 480.2, 0.1, 1)",
            "V": "(492.1, 539.8, 0.1, 1)",
            "Cr": "(554.1, 593.8, 0.1, 1)",
            "Mn": "(618.7, 669.9, 0.1, 1)",
            "Fe": "(686.8, 739.9, 0.1, 1)",
            "Co": "(758.1, 823.2, 0.1, 1)",
            "Ni": "(832.7, 890.0, 0.1, 1)",
            "Cu": "(912.7, 972.3, 0.1, 1)",
        }
    },
    "XLD-6.3.1": {
        "header": "Flying Energy (500, 600, 0.1, 1)",
        "subheader": ["Theta", "Y_samples", "Z_samples", "Magnet Field", "Polarization"],
        "defaults": {
            "element_energy_ranges_prefix": "flying",
            "theta": ["0", "-60"],
            "y_samples": "0",
            "magnetic_field_631": "0",
            "polarization": "0",
        },
        "element_energy_ranges": {
            "O": "(533.1, 553.1, 0.1, 1)",
            "Al": "(1549.6, 1569.6, 0.1, 1)",
            "Si": "(1829, 1839, 0.1, 1)",
            "Ti": "(433.8, 480.2, 0.1, 1)",
            "V": "(492.1, 539.8, 0.1, 1)",
            "Cr": "(554.1, 593.8, 0.1, 1)",
            "Mn": "(618.7, 669.9, 0.1, 1)",
            "Fe": "(686.8, 739.9, 0.1, 1)",
            "Co": "(758.1, 823.2, 0.1, 1)",
            "Ni": "(832.7, 890.0, 0.1, 1)",
            "Cu": "(912.7, 972.3, 0.1, 1)",
        }
    },
    "XMCD-6.3.1": {
        "header": "Flying Energy (763, 813, 0.1, 1)",
        "subheader": ["Theta", "Y_samples", "Z_samples", "Magnet Field", "Polarization"],
        "defaults": {
            "element_energy_ranges_prefix": "flying",
            "y_samples": "0",
            "polarization": "-1",
        },
        "element_energy_ranges": {
            "O": "(533.1, 553.1, 0.1, 1)",
            "Al": "(1549.6, 1569.6, 0.1, 1)",
            "Si": "(1829, 1839, 0.1, 1)",
            "Ti": "(433.8, 480.2, 0.1, 1)",
            "V": "(492.1, 539.8, 0.1, 1)",
            "Cr": "(554.1, 593.8, 0.1, 1)",
            "Mn": "(618.7, 669.9, 0.1, 1)",
            "Fe": "(686.8, 739.9, 0.1, 1)",
            "Co": "(758.1, 823.2, 0.1, 1)",
            "Ni": "(832.7, 890.0, 0.1, 1)",
            "Cu": "(912.7, 972.3, 0.1, 1)",
        }
    },
    "constants": {
        "zero_input": "0",
        "neg_one": "-1",
        "save": "file", 
    }
}


'''
change txt -> scn 

Make a class for the bottom functions 
'''

def build_scan_body(lines, repetitions):
    scan_body = []
    save = ENERGY_SCAN_TEMPLATE["constants"]["save"]
    repetitions = int(int(repetitions) / 2)
    for line1, line2 in lines:
        for _ in range(repetitions):
            scan_body.append('\t'.join(line1))
            scan_body.append('\t'.join(line2))
            scan_body.append(f"{save}\n")

            scan_body.append('\t'.join(line2))
            scan_body.append('\t'.join(line1))
            scan_body.append(f"{save}\n")
    return "".join(scan_body)

def build_XAS_402_body(scan_params, element_energy_tag, repetitions):
    scan_type_defaults = ENERGY_SCAN_TEMPLATE["XAS-4.0.2"]["defaults"]
    hx = f"{scan_type_defaults["hx"]}"
    epu = scan_type_defaults["epu"]
    z = f"{scan_params["z_position"]}"
    theta = f"{scan_params["incident_angle_theta"]}"
    temp = f"{scan_params["temperature"]}"
    line1 = [hx, epu[0], z, theta, temp, element_energy_tag]
    line2 = [hx, epu[1], z,  theta, temp, element_energy_tag]
    return build_scan_body([(line1, line2)], repetitions)

def build_XLD_402_body(scan_params, element_energy_tag, repetitions):
    scan_type_defaults = ENERGY_SCAN_TEMPLATE["XLD-4.0.2"]["defaults"]
    hx = f"{scan_type_defaults["hx"]}"
    hx = f"{ENERGY_SCAN_TEMPLATE["XLD-4.0.2"]["defaults"]["hx"]}"
    epu = scan_type_defaults["epu"]
    z = f"{scan_params["z_position"]}"
    theta = f"{scan_params["incident_angle_theta"]}"
    temp = f"{scan_params["temperature"]}"
    line1 = [hx, epu[0], z, theta, temp, element_energy_tag]
    line2 = [hx, epu[1], z, theta, temp, element_energy_tag]
    return build_scan_body([(line1, line2)], repetitions)

def build_XMCD_402_body(scan_params, element_energy_tag, repetitions):
    hx = [f"+{scan_params["magnetic_field_402"]}", f"-{scan_params["magnetic_field_402"]}"]
    epu = ENERGY_SCAN_TEMPLATE["XMCD-4.0.2"]["defaults"]["epu"]
    z = f"{scan_params["z_position"]}"
    theta = f"{scan_params["incident_angle_theta"]}"
    temp = f"{scan_params["temperature"]}"
    save = ENERGY_SCAN_TEMPLATE["constants"]["save"]
    line1 = [hx[0], epu[0], z, theta, temp, element_energy_tag] # + mag_field | + polarization
    line2 = [hx[1], epu[0], z, theta, temp, element_energy_tag] # - mag_field | + polarization
    line3 = [hx[0], epu[1], z, theta, temp, element_energy_tag] # + mag_field | - polarization
    line4 = [hx[1], epu[1], z, theta, temp, element_energy_tag] # - mag_field | - polarization
    
    scan_body = []
    repetitions = int(int(repetitions) / 2 )
    for _ in range(repetitions):
        scan_body.append('\t'.join(line1))
        scan_body.append('\t'.join(line2))
        scan_body.append(f"{save}\n")

        scan_body.append('\t'.join(line3))
        scan_body.append('\t'.join(line4))
        scan_body.append(f"{save}\n")
    
    return "".join(scan_body)
    # return build_scan_body([(line1, line2)], repetitions)


def build_XAS_631_body(scan_params, element_energy_tag, repetitions):
    scan_type_defaults = ENERGY_SCAN_TEMPLATE["XAS-6.3.1"]["defaults"]
    theta = f"{scan_params["incident_angle_theta"]}"
    y = f"{scan_type_defaults["y_samples"]}"
    z = f"{scan_params["z_position"]}"
    mag_field = f"{scan_type_defaults["magnetic_field_631"]}"
    polarization = f"{scan_type_defaults["polarization"]}"
    line1 = [theta, y, z, mag_field, polarization, element_energy_tag]
    line2 = [theta, y, z, mag_field, polarization, element_energy_tag]
    return build_scan_body([(line1, line2)], repetitions)

def build_XLD_631_body(scan_params, element_energy_tag, repetitions):
    scan_type_defaults = ENERGY_SCAN_TEMPLATE["XLD-6.3.1"]["defaults"]
    theta = scan_type_defaults["theta"]
    y = f"{scan_type_defaults["y_samples"]}"
    z = f"{scan_params["z_position"]}"
    mag_field = f"{scan_type_defaults["magnetic_field_631"]}"
    polarization = f"{scan_type_defaults["polarization"]}"
    line1 = [theta[0], y, z, mag_field, polarization, element_energy_tag]
    line2 = [theta[1], y, z, mag_field, polarization, element_energy_tag]
    return build_scan_body([(line1, line2)], repetitions)

def build_XMCD_631_body(scan_params, element_energy_tag, repetitions):
    scan_type_defaults = ENERGY_SCAN_TEMPLATE["XMCD-6.3.1"]["defaults"]
    theta = f"{scan_params["incident_angle_theta"]}"
    y = f"{scan_type_defaults["y_samples"]}"
    z = f"{scan_params["z_position"]}"
    mag_field = [f"+{scan_params["magnetic_field_631"]}", f"-{scan_params["magnetic_field_631"]}"]
    polarization = f"{scan_type_defaults["polarization"]}"
    # print("mag_field: ", mag_field[0])
    line1 = [theta, y, z, mag_field[0], polarization, element_energy_tag]
    line2 = [theta, y, z, mag_field[1], polarization, element_energy_tag]
    return build_scan_body([(line1, line2)], repetitions)

SCAN_BODY_BULIDERS = {
    "XAS-4.0.2": build_XAS_402_body,
    "XLD-4.0.2": build_XLD_402_body,
    "XMCD-4.0.2": build_XMCD_402_body,
    "XAS-6.3.1": build_XAS_631_body,
    "XLD-6.3.1": build_XLD_631_body,
    "XMCD-6.3.1": build_XMCD_631_body,

}
# Prompt to direct users on how to use program
def display_prompt():
    print("\nWelcome to the LabView Scan File Generator\n")

# Collects user credentials to login when ORCID is initialized into process later down the road
# def get_user_credentials():
#     proposal_id = input("Enter your proposal-ID: ")
#     beamline = input("Enter the beamline you are at: ")
#     username = input("Enter your Scicat username: ")
#     password = input("Enter your Scicat password: ")
#     user = {
#         "proposal_id": proposal_id,
#         "beamline": beamline,
#         "username": username,
#         "password": password,
#         }
#     return user

# Retrieve bar / sample configurations from Scicat Server
def get_bars_from_scicat_server(proposal_id_input=None, username_input=None, password_input=None): 
    if not proposal_id_input or not username_input or not password_input:
        print("Error: Missing required credentials.")
        return None
    
    try:
        # Create a client instance
        # For live hosting
        # client = ScicatClient(
        #     scicat_url = "https://dataportal-staging.als.lbl.gov/api/v3/",
        #     scicat_username = "admin",
        #     scicat_password = "876h876ddms@",
        #     )
        
        # For local hosting
        client = ScicatClient(
            base_url="http://localhost/api/v3", 
            username="admin", 
            password="2jf70TPNZsS", 
            auto_login=False
            )
        
        # For user authentication
        # client = ScicatClient(
        #     base_url="http://localhost/api/v3", 
        #     username=username_input, 
        #     password=password_input, 
        #     auto_login=False
        #     )
        print(f"\nAttempting to login as {client._username}...")

        # Temporarily redirect stderr to suppress library output
        old_stderr = sys.stderr
        sys.stderr = open(os.devnull, 'w')
        client._headers["Host"] = "backend.localhost"
        client.login()

        # Restor stderr
        sys.stderr.close()
        sys.stderr = old_stderr
        print("\nLogin successful.")
    except Exception:
        print(f"Login failed due to incorrect credentials.")
        return None
    
    # When we deploy we should use this line of code so that users have to input a valid input
    # proposal_id = proposal_id_input
    proposal_id = "admin"
    # filter_group = {"ownerGroup": proposal_id, "sampleCharacteristics:als_sample_tracking:group_id":"402"} # hard code each version at the beamline? ask christoph 
    filter_group = {"ownerGroup": proposal_id}
    try:
        scicat_results = client.samples_get_many(filter_fields=filter_group)
        print("Downloading data...\n")
        # print(json.dumps(scicat_results, indent=2))
        return(scicat_results)
    except Exception:
        print(f"Failed to retrieve data:")
        return None

# Create Dictionary of {"bar_id": {"bar_parameters": data, "samples": samples}}
def filter_scicat_bars(proposal_id, scicat_results):
    # print("Scicat_results:\n", json.dumps(scicat_results, indent=2))
    # dictionary for filters dictionary {key : data}, key = bar ID, data = bar parameters and nested sample list
    bars_by_id = {}
    samples_to_process = []
    for result in scicat_results:
        is_id = result.get("_id", {})
        is_owner = result.get("ownerGroup", {})
        is_valid = result.get("sampleCharacteristics", {}).get("als_sample_tracking", {}).get("valid", {}) is True
        is_set = result.get("sampleCharacteristics", {}).get("als_sample_tracking", {}).get("type", {}) == "set"
        is_sample = result.get("sampleCharacteristics", {}).get("als_sample_tracking", {}).get("type", {}) == "configuration"
        
        # Filter for valid sets
        if (is_id != "" and is_owner == proposal_id and is_valid and is_set):
            scicat_bar = {
                "bar_id": result["_id"],
                "bar_name": result["description"],
                "proposal_id": result["ownerGroup"],
                "valid": result["sampleCharacteristics"]["als_sample_tracking"]["valid"],
                "sample_configs": []
            }
            bars_by_id[scicat_bar["bar_id"]] = scicat_bar
            # print("Current bar:\n", json.dumps(barsById[scicat_bar["bar_id"]]))
        elif (is_id != "" and is_owner == proposal_id and is_valid and is_sample):
            samples_to_process.append(result)
            # sorted_samples = sorted(samples_to_process)
        # Filter for valid samples
    # print("Valid Samples: ", json.dumps(samples_to_process, indent=2))
    for sample in samples_to_process:
        shortcut = sample["sampleCharacteristics"]["als_sample_tracking"]
        set_id = shortcut["set_id"]
        sample_parameters = shortcut["parameters"]
        # Filter for samples that pertain to a valid set
        if (set_id in bars_by_id):
            # print("Samples to process:\n", json.dumps(sample, indent=2))
            # print(f"This sample is part of a valid bar: {set_id}", json.dumps(sample,indent=2))
            # print("Current sample:\n", json.dumps(result, indent=2))
            scicat_sample = {
                "sample_id": sample["_id"],
                "sample_name": sample["description"],
                "bar_id": shortcut["set_id"],
                "scan_type": shortcut["scan_type"],
                "scan_parameters": sample_parameters
            }

            # print("Adding Sample: ", json.dumps(scicat_sample,indent=2))
            bars_by_id[scicat_sample["bar_id"]]["sample_configs"].append(scicat_sample)
    # print("barsById", json.dumps(bars_by_id, indent=2))
    
    # return bars_by_id sorted
    for bar_id, bar_data in bars_by_id.items():
        bar_data["sample_configs"].sort(key=lambda sample: float(sample["scan_parameters"]["order_of_samples"]))

    # print("Sorted Samples:\n", json.dumps(bars_by_id, indent=2))
    
    return bars_by_id

def get_scan_number(scan_type): 
    match = re.search(r"-(\d+\.\d+\.\d+)", scan_type)
    if match: 
        return match.group(1).replace('.', '') # 402 or 631
    return None

def build_scan_file(filtered_bars):
    # print("Filtered Bars:\n", json.dumps(filtered_bars, indent=2))
    scan_files = {}
    start_count = 1
    end_count = 0
    header_count = 0
    
    # Iterate through the bars
    for bar_id, bar_data in filtered_bars.items():
        proposal_id = bar_data["proposal_id"]
        bar_name = bar_data["bar_name"]
        bar_id = bar_data["bar_id"]
        scan_files[bar_id] = {
            "proposal": proposal_id,
            "bar_name": bar_name,
            "bar_id": bar_id,
            "scan_files": []
        }
        # print("Scan File:\n", json.dumps(scan_files, indent=2))
        # Iterate through samples
        for config in bar_data["sample_configs"]:
            sample_name = config["sample_name"]
            sample_id = config["sample_id"]
            scan_type = config["scan_type"]
            beamline_number = get_scan_number(scan_type)
            scan_parameters = config["scan_parameters"]
            repetitions = int(scan_parameters["repetitions"])
            if (scan_type in ENERGY_SCAN_TEMPLATE and scan_type in SCAN_BODY_BULIDERS):
                end_count += repetitions
                scan_file_info = [
                    f"# Sample Info: sample name: {sample_name} | bar name: {bar_name} | scan type: {scan_type}", 
                    "# Scan Parameters: " + ", ".join(f"{k}: {v}" for k, v in scan_parameters.items()), 
                    f"# Scan Numbers: [{start_count}, {end_count}]"]
                scan_file_info = "\n".join(scan_file_info) + "\n"
                scan_file_info = "".join(scan_file_info)
                header = f"{ENERGY_SCAN_TEMPLATE[scan_type]["header"]}" + "\n"
                subheader = "\t".join(ENERGY_SCAN_TEMPLATE[scan_type]["subheader"]) + "\n"
                scan_tag = ENERGY_SCAN_TEMPLATE[scan_type]["defaults"]["element_energy_ranges_prefix"]
                element = scan_parameters["element"]
                # if scan_type in ["XAS-4.0.2", "XLD-4.0.2", "XMCD-4.0.2"]:
                if beamline_number == '402':
                    element_energy_tag = f"{scan_tag}{ENERGY_SCAN_TEMPLATE[scan_type]["element_energy_ranges"][element]}\n"
                    build_body = (SCAN_BODY_BULIDERS[scan_type](scan_parameters, element_energy_tag, repetitions))
                elif beamline_number == '631':
                    element_energy_tag = f"{scan_tag}{ENERGY_SCAN_TEMPLATE[scan_type]["element_energy_ranges"][element]}\n"
                    build_body = SCAN_BODY_BULIDERS[scan_type](scan_parameters, element_energy_tag, repetitions)
                    
                if header_count < 1:
                    build_file = f"{scan_file_info}{header}{subheader}{build_body}"
                else:
                    build_file = f"{scan_file_info}{build_body}"
                scan_files[bar_id]["scan_files"].append({
                    "sample_id": sample_id,
                    "sample_name": sample_name,
                    "scan_type": scan_type,
                    "bar_id": bar_id,
                    "bar_name": bar_name,
                    "file": build_file,
                })
                start_count = start_count + int(repetitions)
            else: 
                print("This Sample Scan Type is not avaiable.")
            header_count += 1
    # print(json.dumps(scan_files,indent=2))
    return scan_files

def save_files(scan_files, proposal_id):
    # print("Scan_Files: ", json.dumps(scan_files, indent=2))
    main_dir = f"{proposal_id}"
    os.makedirs(main_dir, exist_ok=True)
    print(f"Main dir:\n\t{main_dir}")
    for bar_id, bar_data in scan_files.items():
        bar_name = bar_data["bar_name"]
        bar_dir = os.path.join(main_dir, bar_name)
        os.makedirs(bar_dir, exist_ok=True)
        combined_scan_file = ""
        for sample in bar_data["scan_files"]:
            combined_scan_file += sample["file"]
        combined_scan_file = combined_scan_file.rstrip('\n')
        scan_file_name = os.path.join(bar_dir, f"{bar_name}.txt")
        with open(scan_file_name, 'w') as f:
            f.write(combined_scan_file)
        print(f"Saved: {scan_file_name}")


def run_file_generator():
    # user_credentials = get_user_credentials()
    user_credentials = {
        "proposal_id": "admin",
        "beamline": "",
        "username": "admin",
        "password": "2jf70TPNZsS",
        }
    scicat_results = get_bars_from_scicat_server(proposal_id_input=user_credentials["proposal_id"], username_input=user_credentials["username"], password_input=user_credentials["password"])
    if scicat_results is None:
        print("Failed to retrieve data from Scicat. Please try again.\n")
        return
    # print("Scicat Results:\n", json.dumps(scicat_results, indent=2))
    filtered_bars = filter_scicat_bars(user_credentials["proposal_id"], scicat_results)
    scan_files = build_scan_file(filtered_bars)
    save_files(scan_files, user_credentials["proposal_id"])
    print("Files saved to computer.")

def main():
    display_prompt()
    repeat = "y"
    run_file_generator()
    # while repeat[0] != 'n':
    #     repeat = input("Would you like to generate any more scan files? (Y/N): ").lower() # quits the loop
    #     print("\n")
    #     if repeat[0] == 'y':
    #         run_file_generator()
    #     elif repeat[0] == 'n':
    #         print("Quitting program...")
    #         exit
    #     else:
    #         print("You made an invalid choice:")
    

if __name__ == "__main__":
    main()


    '''
    How can users log in and retrieve data (user log in)
    Make 1 profile to get Alpha and Christoph to be able to test it so they have an account and proposal to use on both the interface and python scicat scripts. 
    Orcid is the way the logging and verification will handled, get users to sign up with orcid. 
    Christoph needs to sign with orcid
    near future deploying authentication server where pyscicat will connect to so we can use orcid in conjunction 
    currently can only use standard active predefined accounts for scicat log in, will soon be improved



    If users don't have an account make a guest profile possible. 
    Notes about each samples paramters (need a file counter to distinguish sample scans, use starting number of 1 for counter)

    Notes about build scan body & body building function create class to encapsulate most of these functions and their work. 
    Summary at the top of this file 
    1 - 2 lines of comment for each function
    '''