import xml.etree.ElementTree as ET
import os
import re
import csv

# --- Configuration ---
XML_FILE_PATH = 'Data.xml'
# Updated output file name for the wide table format
CSV_OUTPUT_PATH = 'stream_data_wide_format.csv'

# --- Internal States for Filtering Logic ---
STATE_GENERAL = 0
STATE_IN_UNWANTED_BLOCK = 1

# --- Component Properties List ---
# This list is used to identify simple, repeating component names that must be skipped
# when they appear under a phase-specific block (like Molar Fraction (Solid)).
# This set must contain all component names from your DWSIM simulation.
SIMPLE_COMPONENT_PROPERTIES = {
    'Water', 'Sulfuric acid', 'Cellulose_GlucoseChemeo',
    'Hemicellulose_XyloseChemeo', 'Lignin_VanillinChemeo',
    'SolidCellulose_GlucoseChemeo', 'SolidHemicellulose_XyloseChemeo',
    'SolidLignin_VanillinChemeo', 'Others', 'Furfural', 'GammaValerolactone', 'Toluene'
}

# --- Phase Keywords for Unwanted Blocks ---
# These keywords signal the start of a phase-specific block that is not "Mixture".
PHASE_KEYWORDS = ['(Vapor)', '(Liquid 1)', '(Liquid 2)', '(Aqueous)', '(Solid)', '(Overall Liquid)']

# --- Additional Properties to Exclude Explicitly ---
# These properties are simple (no slash or phase name) but still represent phase-specific
# bulk properties that should be omitted (e.g., Aqueous/Solid fractions).
EXPLICIT_EXCLUSION_PROPERTIES = {
    'Aqueous Mass Fraction',
    'Aqueous Volumetric Fraction'
}


def clean_float_value(value_str):
    """
    Standardizes number formats from the XML file for calculations.

    Process simulation software can export numbers in different formats. This function
    is designed to handle two common cases:
    1.  Decimal commas: Converts strings like '46,6221' into a standard Python
        float like 46.6221.
    2.  Scientific notation: Converts strings like '2.85231E+07' into a standard
        float.

    If a value is just text (e.g., a component name), it is left unchanged.

    Args:
        value_str (str): The text value read from the XML file.

    Returns:
        float or str: The cleaned number for calculations, or the original text.
    """
    if isinstance(value_str, str) and value_str.strip():
        # Replace comma decimal separator with a dot for standard float conversion
        cleaned_str = value_str.replace(',', '.')

        try:
            # Attempt to convert the string to a floating-point number
            return float(cleaned_str)
        except ValueError:
            # If conversion fails, it's not a number; return the original text
            return value_str.strip()

    # Return original value if it's not a non-empty string
    return value_str


def parse_xml_data(file_path):
    """
    Reads the raw DWSIM XML data file and extracts all process stream information.

    Uses a state machine to correctly filter out phase-specific component properties
    (e.g., 'Ethanol' under 'Molar Fraction (Solid)') while keeping general stream
    properties and overall 'Mixture' composition data.

    Args:
        file_path (str): The location of the 'Data.xml' file.

    Returns:
        tuple: A pair of (extracted_data, property_order).
            - extracted_data (list): A list where each item is a dictionary

              representing a single stream and its properties.
            - property_order (list): An ordered list of all unique property names
              found, used to define the rows in the final CSV file.
    """
    print(f"--- Reading and parsing XML file at: {os.path.abspath(file_path)} ---")
    extracted_data = []
    root = None

    # This list will store the unique property names in the order they appear
    property_order = []
    seen_properties = set()

    try:
        # WORKAROUND: Handle potential XML encoding errors
        print("Applying workaround for potential XML encoding errors...")
        with open(file_path, 'r', encoding='latin-1') as f:
            xml_content = f.read()

        # Use a regular expression to find and remove the XML declaration
        xml_declaration_match = re.search(r'<\?xml[^>]*\?>', xml_content, re.IGNORECASE)
        if xml_declaration_match:
            # Remove the XML declaration from the content to avoid parsing errors
            cleaned_content = xml_content.replace(xml_declaration_match.group(0), '', 1).strip()
        else:
            cleaned_content = xml_content.strip()

        # Parse the cleaned XML content from the string
        root = ET.fromstring(cleaned_content)

    except FileNotFoundError:
        print(f"\nERROR: File not found at '{file_path}'.")
        return [], []
    except ET.ParseError as e:
        print(f"\nXML Parsing Error: {e}. The file may be malformed.")
        return [], []
    except Exception as e:
        print(f"\nAn unexpected error occurred: {e}")
        return [], []

    # --- Data Extraction Logic ---
    if root is not None:
        try:
            print(f"XML loaded successfully. Root element: <{root.tag}>")
            print("\n--- Extracting Stream Data ---")

            # In the DWSIM XML, each process stream is an <Object> element.
            for obj in root.findall('./Object'):
                stream_name = obj.get('name')
                stream_type = obj.get('type')

                stream_properties = {'Name': stream_name, 'Type': stream_type}
                state = STATE_GENERAL  # Reset state for each new stream

                # Each engineering property is a <Property> tag within the <Object>
                for prop in obj.findall('Property'):
                    prop_name = prop.get('name')
                    prop_value_str = prop.get('value')
                    prop_units = prop.get('units')

                    # --- START STATE-MACHINE FILTERING LOGIC ---

                    # 0. Check for explicit exclusion (the requested properties or Envelopes)
                    if prop_name in EXPLICIT_EXCLUSION_PROPERTIES or 'Envelope' in prop_name:
                        if 'Envelope' in prop_name:
                            print(f"DEBUG: Skipping envelope property: {prop_name}")
                        continue  # Skip this property

                    is_mixture_property = 'Mixture' in prop_name
                    is_phase_property = any(p in prop_name for p in PHASE_KEYWORDS)

                    # 1. Check for specific, qualified properties (contains /)
                    if '/' in prop_name:
                        if is_mixture_property:
                            state = STATE_GENERAL  # Always keep Mixture qualified properties
                        else:
                            continue  # Omit non-Mixture qualified phase properties

                    # 2. Check for phase-related properties (contains (PhaseName))
                    elif is_phase_property:
                        if is_mixture_property:
                            state = STATE_GENERAL  # Keep properties like 'Density (Mixture)'
                        else:
                            # Unwanted Phase Header (e.g., Molar Fraction (Vapor)).
                            # Enter unwanted block state and skip the header itself.
                            state = STATE_IN_UNWANTED_BLOCK
                            continue

                    # 3. Check for simple component names (Ethanol, Water, etc.)
                    elif prop_name in SIMPLE_COMPONENT_PROPERTIES:
                        if state == STATE_IN_UNWANTED_BLOCK:
                            # Omit component property that falls immediately after an unwanted phase header
                            continue
                            # Else: Keep component property (it's the overall stream value).
                        state = STATE_GENERAL

                    # 4. All other properties (General stream properties, like Temperature, Mass Flow)
                    else:
                        state = STATE_GENERAL  # Keep, and exit any prior unwanted block

                    # --- END STATE-MACHINE FILTERING LOGIC ---

                    # If we reached here, the property is kept.

                    # Create a standardized name for the property including its unit
                    combined_name = f"{prop_name} [{prop_units}]" if prop_units else f"{prop_name} [-]"

                    # Keep track of the order of properties as they first appear
                    if combined_name not in seen_properties:
                        property_order.append(combined_name)
                        seen_properties.add(combined_name)

                    # Standardize the property's value (e.g., convert '1,23' to 1.23)
                    stream_properties[combined_name] = clean_float_value(prop_value_str)

                extracted_data.append(stream_properties)

            print(f"Successfully extracted data for {len(extracted_data)} streams.")
            return extracted_data, property_order

        except Exception as e:
            print(f"\nAn error occurred during data extraction: {e}")

    return [], []


def save_data_to_csv(data, output_path, row_order):
    """
    Pivots the extracted data into a "wide-format" table and saves it to a CSV file.

    The resulting CSV has Property names as the first column, and stream names as the
    headers for subsequent columns, making it easy to paste into spreadsheets.
    """
    if not data:
        print("\nNo data to save to CSV.")
        return

    print(f"\n--- Formatting data into a wide table and saving to CSV at: {output_path} ---")

    def sort_key_natural(stream_dict):
        """
        Custom sorting function to ensure streams are sorted naturally (e.g., S-1, S-2, S-10)
        rather than alphabetically (S-1, S-10, S-2).
        """
        name = stream_dict['Name']
        # Extract numbers from the stream name for sorting
        numbers = re.findall(r'(\d+)', name)
        # Sort by the last number found in the name, fallback to a high number if no number is found
        return (int(numbers[-1]), name) if numbers else (999999, name)

    # 1. Sort the streams to ensure columns are in a logical order
    sorted_streams = sorted(data, key=sort_key_natural)

    # 2. Define the header row for the CSV file. It will be "Property" followed
    #    by the name of each sorted stream.
    stream_names = [s['Name'] for s in sorted_streams]
    fieldnames = ['Property'] + stream_names

    # 3. Build the pivoted data row by row
    pivoted_data = []
    for prop_name in row_order:
        # Start a new row with the property name as the first cell
        row = {'Property': prop_name}
        # Iterate through each stream to get the value for the current property
        for stream in sorted_streams:
            # Get the value for this stream and property, or an empty string if it's missing
            row[stream['Name']] = stream.get(prop_name, "")
        pivoted_data.append(row)

    # 4. Write the structured data to the CSV file
    try:
        # 'utf-8-sig' encoding helps Excel open the file correctly, especially
        # with special characters.
        with open(output_path, 'w', newline='', encoding='utf-8-sig') as csvfile:
            writer = csv.DictWriter(csvfile, fieldnames=fieldnames, delimiter=',')
            writer.writeheader()
            writer.writerows(pivoted_data)
        print(f"Successfully saved data to '{output_path}'.")
    except Exception as e:
        print(f"ERROR: Could not write CSV file: {e}")


if __name__ == "__main__":

    def reorder_properties(order_list, moved_fragment, anchor_fragment):
        """
        Reorders a list of property names by moving the 'moved_fragment'
        (matching the full name including units) to immediately follow the
        'anchor_fragment' (matching the full name including units).
        """
        # Find the full names including units
        moved_name = next((name for name in order_list if moved_fragment in name), None)
        anchor_name = next((name for name in order_list if anchor_fragment in name), None)

        if not moved_name:
            print(f"Warning: Property '{moved_fragment}' not found in data. Skipping reorder.")
            return order_list
        if not anchor_name:
            print(f"Warning: Anchor property '{anchor_fragment}' not found in data. Skipping reorder.")
            return order_list

        # Ensure we don't try to remove something that is not there
        try:
            # Remove the moved property from its current location
            order_list.remove(moved_name)
        except ValueError:
            # Should not happen if moved_name was found by next() above, but safe guard.
            pass

            # Find the index of the anchor property in the new, shorter list
        anchor_index = order_list.index(anchor_name)

        # Insert the moved property immediately after the anchor property
        order_list.insert(anchor_index + 1, moved_name)

        print(f"Property '{moved_fragment}' reordered after '{anchor_fragment}'.")
        return order_list


    # --- Main Workflow ---
    # This block runs when the script is executed directly.

    # 1. Read and parse the raw data from the XML file.
    # This returns a list of all streams and a list of all unique properties found.
    data, property_order = parse_xml_data(XML_FILE_PATH)

    # Proceed only if data was successfully extracted
    if data:
        # --- Apply User-Requested Reordering ---
        ANCHOR_FRAGMENT = "Volumetric Flow"
        MOVED_FRAGMENT = "Molar Enthalpy (Mixture)"

        # Reorder the list in place
        property_order = reorder_properties(property_order, MOVED_FRAGMENT, ANCHOR_FRAGMENT)

        # 2. Pivot the data and save it to a wide-format CSV file.
        # This creates the final table for analysis in other software.
        save_data_to_csv(data, CSV_OUTPUT_PATH, property_order)

        # --- Optional: Example Analysis ---
        # The following section demonstrates how you can directly access and use the
        # extracted data within Python for quick checks or calculations.
        print("\n--- Example: Quick Data Check (First 3 Streams) ---")

        # Define property name fragments for the example, which will be matched
        # against the full property names (with units) in the stream data.
        properties_to_show_fragments = [
            'Mass Flow',
            'Energy Flow',
            'Temperature',
            'Pressure',
            'Molar Flow',
            'Volumetric Flow',
            'Molar Enthalpy (Mixture)',  # Now included in the example printout
            'Ethanol',  # Overall component
            'Water',  # Overall component
            'Density (Mixture)',
        ]

        for stream in data[:3]:
            print("-" * 40)
            print(f"Stream Name: {stream['Name']} ({stream['Type']})")

            for prop_fragment in properties_to_show_fragments:
                # Find the actual key in the stream dictionary that contains the fragment
                # This handles the fact that we don't know the exact unit beforehand.
                full_prop_key = next((key for key in stream.keys() if prop_fragment in key), None)

                if full_prop_key:
                    # Display the full key (name + unit) and its value
                    print(f"  {full_prop_key}: {stream[full_prop_key]}")

        # Example calculation: Sum the energy flow of all 'Energy Stream' types
        total_energy_streams = [
            s.get('Energy Flow [kJ/h]', 0)
            for s in data
            if s['Type'] == 'Energy Stream' and isinstance(s.get('Energy Flow [kJ/h]'), (int, float))
        ]

        if total_energy_streams:
            print("\n" + "=" * 50)
            print("Example Calculation:")
            print(f"Total Energy Flow from all Energy Streams: {sum(total_energy_streams):,.2f} kJ/h")
            print("=" * 50)