magic_script.py

import xml.etree.ElementTree as ET
import pandas as pd
from pathlib import Path


class PROTOCOLS:

    HFE = 'HFE'
    B27 = 'B27'
    LCT = 'LCT'
    GLU = 'GLU'
    RHDF = 'RHDF'

class ASSAY: 

    def __init__(self, protocol: str):
        self.protocol = protocol

    def get_assay_name(self) -> str: 

        match self.protocol:
            case PROTOCOLS.HFE:
                return "1013_HFE_96"
            case PROTOCOLS.B27:
                return "HLA_B27_ViennaLab"
            case PROTOCOLS.LCT:
                return "1006_Lactose_96"
            case PROTOCOLS.GLU:
                return "Unknown"
            case PROTOCOLS.RHDF:
                return "Unknown"
            case _:
                return "Unknown"
            
    def get_run_profile_name(self) -> str:

        match self.protocol:
            case PROTOCOLS.HFE:
                return "SimpleProbe 530-640"
            case PROTOCOLS.B27:
                return "HLA_B27_ViennaLab"
            case PROTOCOLS.LCT:
                return "Simple Probe 530"
            case PROTOCOLS.GLU:
                return "Unknown"
            case PROTOCOLS.RHDF:
                return "Unknown"
            case _:
                return "Unknown"

    def get_additional_rows(self) -> pd.DataFrame:
        
        match self.protocol: 
            case PROTOCOLS.HFE:
                return pd.DataFrame({
                    "SampleId": ["Standard"] * 6 + ["1013_HFE_NTC"],
                    "WellPosition": ["B12", "C12", "D12", "E12", "F12", "G12", "H12"],
                    "Targetname": [
                        "H63 S65 / H63 65C|282Y / 282Y",
                        "63D S65 / 63D S65|C282 / C282",
                        "H63 65C / H63 65C",
                        "H63 S65 / H63 65C",
                        "H63 S65 / 63D S65|C282 / 282Y",
                        "63D S65 / H63 65C",
                        None
                    ],
                    "SampleRole": ["MSR"] * 6 + ["NTC"],
                })
            
            case PROTOCOLS.B27:
                return pd.DataFrame({
                    "SampleId": ["HLA-B27 POS", "HLA-B27 NEG", "NTC PCR"],
                    "WellPosition": ["F12", "G12", "H12"],
                    "Targetname": [None] * 3,
                    "SampleRole": ["PTC", "NPC", "NTC"],
                })
            
            case PROTOCOLS.LCT:
                return pd.DataFrame({
                    "SampleId": ["Lactose_WT", "Lactose_HT", "Lactose_MT", "1006_Lactose_NTC"],
                    "WellPosition": ["E12", "F12", "G12", "H12"],
                    "Targetname": ["Lactose_WT", "Lactose_HT", "Lactose_MT", None],
                    "SampleRole": ["MS", "MS", "MS", "NTC"],
                })

            case _: 
                return pd.DataFrame({
                    "SampleId": [],
                    "WellPosition": [],
                    "Targetname": [],
                    "SampleRole": [],
                })

def generate_qia_df_without_controls(df: pd.DataFrame) -> pd.DataFrame:

    sample_id = df["SampleId"]
    well_position = [elt.split(":")[0] + elt.split(":")[1] for elt in list(df["PositionName"])]
    target_name = [None] * len(df)
    sample_role = ["Unknown"] * len(df)

    return pd.DataFrame({
        "SampleId": sample_id,
        "WellPosition": well_position,
        "Targetname": target_name,
        "SampleRole": sample_role,
    })


def extract_protocol(rack_id: str) -> str:
    return rack_id.split('_')[1]


def read_rack_xml(filepath: str) -> pd.DataFrame:
    tree = ET.parse(filepath)
    root = tree.getroot()

    rack_id = root.findtext('RackId')

    df = pd.read_xml(filepath, xpath='.//RackPosition')
    df.insert(0, 'RackId', rack_id)

    return df

def generate_final_csv(filepath: str) -> pd.DataFrame:

    # Read the XML file and convert it to a DataFrame
    df_1 = read_rack_xml(filepath)

    # Extract protocol from RackId
    protocol = extract_protocol(df_1["RackId"].iloc[0])

    # Remove all empty rows
    df_2 = df_1[df_1["SampleType"] == "Sample"]

    # Get temporary df without controls
    df_3 = generate_qia_df_without_controls(df_2)

    # Concatenate the temporary df with the additional rows based on the protocol
    assay = ASSAY(protocol)
    additional_rows = assay.get_additional_rows()
    df_4 = pd.concat([df_3, additional_rows], ignore_index=True)

    df_4['PlateId'] = None
    df_4['PlateType'] = 96
    df_4['Dilution'] = None
    df_4['MasterMixName'] = None
    df_4['LotNumber'] = None
    df_4['RunProfileName'] = assay.get_run_profile_name()
    df_4['AssayName'] = assay.get_assay_name()
    df_4['PlatesetupName'] = df_1["RackId"].iloc[0]

    columns = [
        "PlateId", "PlatesetupName", "PlateType", "SampleId", "WellPosition",
        "AssayName", "Targetname", "Dilution", "MasterMixName", "LotNumber",
        "RunProfileName", "SampleRole",
    ]
    return df_4[columns]


def process_all(input_dir: str = "inputs", output_dir: str = "outputs") -> None:
    Path(output_dir).mkdir(exist_ok=True)

    for xml_file in Path(input_dir).glob("*.xml"):
        df = generate_final_csv(str(xml_file))
        rack_id = df["PlatesetupName"].iloc[0]
        df.to_csv(Path(output_dir) / f"{rack_id}.csv", index=False)


if __name__ == "__main__":
    process_all()