import os
import ast
import numpy as np

# This function takes the folder where the txt files are and the filenumber i as an argument to read in the file from "folder/weak_pnr_100000_i.txt"
def load_file_data(folder, i):
    file_path = os.path.join(folder, f"weak_pnr_100000_{i}.txt")
    # It returns the data that is extracted using the read_data function
    return read_data(file_path)

# Reading in Data from data_path
def read_data(data_path):
    # Prepares lists for 
    # trigger_pattern: heralding, 
    # signal_pattern: interferometer output, 
    # counts: number of occurrences,
    # pulsepicker_check: EOM state
    trigger_pattern = []
    signal_pattern = []
    counts = []
    pulsepicker_check = []

    # Opens the txt file at data_path
    with open(data_path, 'r') as file:
        # Reads in line by line
        for line in file:
            
            # Splits the content of the line, appends it to trigger, signal patterns and counts
            # Note: The unique identifier is not used/read in, which would be parts[0]
            parts = line.strip().split()
            trigger_pattern.append(tuple(ast.literal_eval(parts[1])))
            signal_pattern.append(tuple(ast.literal_eval(parts[2])))
            counts.append(int(parts[4]))

            # Checks for EOM state
            try:
                pp_state = ast.literal_eval(parts[3])
                pulsepicker_check.append(pp_state[0] if pp_state else None)
            except:
                pulsepicker_check.append(None)

    # Converts pulsepicker state into numpy array
    pulsepicker_check = np.array(pulsepicker_check, dtype = object)
    
    # Checks for EOM state in the whole txt file
    if len(pulsepicker_check) == 0:
        # Check for empty file, in which case no data is available
        check = 2
        trigger_pattern, signal_pattern, counts = None, None, None

    elif np.all(pulsepicker_check == 0):
        # All lines are EOM state 0, i.e. GBS
        check = 0
    
    elif np.all(pulsepicker_check == 1):
        # All lines are EOM state 1, i.e. SBS/TBS
        check = 1
    
    else:
        # There are components from both cases mixed in a file or something else went wrong
        # In that case the data is discarded
        check = 2
        trigger_pattern, signal_pattern, counts = None, None, None

    # Returns the 
    # heralding/trigger pattern, converted to numpy array, 
    # signal pattern/interferometer output, converted to numpy array, 
    # counts/number of occurrences, converted to numpy array,
    # indicator whether it is GBS (0), SBS/TBS (1) or empty/mixed file (2) 
    return np.array(trigger_pattern), np.array(signal_pattern), np.array(counts), check 

if __name__ == "__main__":
    # Example to read in txt file 170 from Run1 (Squeezing of r = 0.2)
    # Note that you should extract the Run1.tar.gz first
    Heralding, Signal, Counts, EOM_Ind = load_file_data("Run1", 170)

    print("Configuration of the EOM: ", EOM_Ind)

    if EOM_Ind != 2:
        # Print out the data line by line
        print("Heralding, Interferometer, Counts")
        for i in range(len(Counts)):
            print(Heralding[i], Signal[i], Counts[i])
            A = Heralding[i] * Signal[i]