# This parses a Modelica model to extract its components (variables and parameters) and their parameters. 
# It uses the ANTLR-generated lexer and parser to build a parse tree and then walks the tree to extract relevant information.

# Configuration
file_name = 'SimpleHeatingSystem.mo'

import sys
import os
from antlr4 import *

# Import lexer, parser, and listener generated by ANTLR for Modelica grammar.
from modelicaLexer import modelicaLexer
from modelicaParser import modelicaParser
from modelicaListener import modelicaListener

# Listener class to extract components and their parameters from the parse tree
class ModelicaFeatureExtractor(modelicaListener):
    def __init__(self, include_equations):
        self.components = []  # List to store extracted components
        self.equations = []  # List to store extracted equations
        self.include_equations = include_equations

    # Method called when entering a component clause in the parse tree
    def enterComponent_clause(self, ctx: modelicaParser.Component_clauseContext):
        # Extract the component type and name
        type_prefix = ctx.type_prefix().getText() if ctx.type_prefix() else ""
        component_type = type_prefix + " " + ctx.type_specifier().getText()
        # print(f"Component type: {component_type}")

        # Iterate over each component declaration in the component list
        for component in ctx.component_list().component_declaration():
            component_name = component.declaration().IDENT().getText()
            # print(f"Component name: {component_name}")  ## uncomment for debugging

            parameters = {}  # Dictionary to store parameters and their values

            # Check if there are any modifications (initial values or other modifications)
            if component.declaration().modification() is not None:
                modification = component.declaration().modification()
                param_value = self.get_modification_value(modification)
                if param_value is not None:
                    parameters["value"] = param_value  # Store the parameter value
                    # print(f"Param: {component_name} = {param_value}")  uncomment for debugging

            # Add the component to the list
            self.components.append((component_type.strip(), component_name, parameters))

    # Helper method to extract the value from a modification context
    def get_modification_value(self, modification_ctx):
        if modification_ctx is None:
            return None
        if modification_ctx.expression():  # Direct value assignment
            return modification_ctx.expression().getText()
        elif modification_ctx.class_modification():
            class_mod = modification_ctx.class_modification()  # Class modification
            if class_mod.argument_list() is not None:
                values = []
                for argument in class_mod.argument_list().argument():
                    if hasattr(argument, 'element_modification'):
                        element_mod = argument.element_modification()
                        param_name = element_mod.name().getText()
                        param_value = self.get_modification_value(element_mod.modification())
                        values.append(f"{param_name}={param_value}")
                return "({})".format(", ".join(values))
        return None

    # Method to extract equations
    def enterEquation_section(self, ctx: modelicaParser.Equation_sectionContext):
        if self.include_equations:
            for equation in ctx.equation():
                equations_text = equation.getText()
                self.equations.append(equations_text)

# Function to parse a Modelica file and extract components
def parse_model(file_name, include_equations):

    input_stream = FileStream(file_name, encoding='utf-8')  # Read the file
    lexer = modelicaLexer(input_stream)  # Tokenize the input
    stream = CommonTokenStream(lexer)  # Create a token stream
    parser = modelicaParser(stream)  # Create a parser
    tree = parser.stored_definition()  # Parse the input into a parse tree

    extractor = ModelicaFeatureExtractor(include_equations)  # Create a listener
    walker = ParseTreeWalker()  # Create a parse tree walker
    walker.walk(extractor, tree)  # Walk the parse tree with the listener

    return extractor.components, extractor.equations  # Return the extracted components

if __name__ == "__main__":
    components, equations = parse_model(file_name, include_equations)  # Parse the file and extract components
    for component in components:
        print(component)  # Print the extracted components
    for equation in equations:
        print(equation)