import os
from pycallgraph2 import PyCallGraph, Config
from pycallgraph2.output import GraphvizOutput
from pycallgraph2.globbing_filter import GlobbingFilter

# Define the path to the folder inside the apps directory
apps_directory = '/local/ca00xebo/repos/pymatlib/apps'
image_folder = os.path.join(apps_directory, 'callgraph_images')

# Create the folder if it does not exist
if not os.path.exists(image_folder):
    os.makedirs(image_folder)

# Create configuration to focus on your specific functions
config = Config()
config.trace_filter = GlobbingFilter(
    include=[
        'pymatlib.core.interpolators.*',
        'pymatlib.core.data_handler.*',
        'InterpolationArrayContainer.*',

        'assignment_converter',
        'interpolate_property',
        'check_equidistant',
        'check_strictly_increasing',
    ],
    exclude=[
        'pycallgraph2.*',
        '*.append',  # Exclude common methods
        '*.join',
    ]
)

# Configure the output with the full path
output_file = os.path.join(image_folder, 'array_container_callgraph.svg')

# Configure the output
graphviz = GraphvizOutput(
    output_file=output_file,
    font_name='Verdana',
    font_size=7,
    group_stdlib=True,
    output_type='svg',
    dpi=1200
)

# Run the specific part of your code with call graph tracking
with PyCallGraph(config=config, output=graphviz):
    import numpy as np
    from pystencils import fields
    from importlib.resources import files
    from pystencilssfg import SourceFileGenerator
    from pymatlib.core.yaml_parser import create_alloy_from_yaml
    from pymatlib.core.interpolators import InterpolationArrayContainer


    with SourceFileGenerator() as sfg:
        u = fields(f"u: double[2D]", layout='fzyx')

        T_bs = np.array([3243.15, 3248.15, 3258.15, 3278.15], dtype=np.float64)
        E_bs = np.array([1.68e10, 1.69e10, 1.70e10, 1.71e10], dtype=np.float64)

        custom_container = InterpolationArrayContainer("BinarySearchTests", T_bs, E_bs)
        sfg.generate(custom_container)

        T_eq = np.array([3243.15, 3253.15, 3263.15, 3273.15], dtype=np.float64)
        E_neq = np.array([1.68e10, 1.69e10, 1.70e10, 1.71e10], dtype=np.float64)

        custom_container = InterpolationArrayContainer("DoubleLookupTests", np.flip(T_eq), np.flip(E_neq))
        sfg.generate(custom_container)

        yaml_path = files('pymatlib.data.alloys.SS304L').joinpath('SS304L_comprehensive.yaml')
        mat = create_alloy_from_yaml(yaml_path, u.center())
        arr_container = InterpolationArrayContainer.from_material("SS304L", mat)
        sfg.generate(arr_container)