# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/19_pctexamples.ipynb.

# %% auto #0
__all__ = ['PCTExamples']

# %% ../nbs/19_pctexamples.ipynb #90320abe
import json, os
import matplotlib.pyplot as plt
from   matplotlib import colors

# %% ../nbs/19_pctexamples.ipynb #b785a203
from .hierarchy import PCTHierarchy
from .environment_processing import EnvironmentProcessingFactory

# %% ../nbs/19_pctexamples.ipynb #a4875394
class PCTExamples:
    """
    PCTExamples class provides methods to load a PCT hierarchy from a configuration file, 
    summarize the hierarchy, get the configuration, draw the hierarchy, run the hierarchy 
    for a specified number of steps, and close the environment.
    Attributes:
        hierarchy (PCTHierarchy): The loaded PCT hierarchy.
        env (Environment): The environment associated with the PCT hierarchy.
    """

    def __init__(self, config_file, min=True, early_termination=False, history=False, additional_props=None, 
                 render=False, video_params=None, suffixes=False):
        """
        Initializes the PCTExamples instance by loading the hierarchy from the given configuration file.
        """
        self.config_file = config_file
        self.hierarchy, self.env, self.environment_properties = PCTHierarchy.load_from_file(
            config_file, min=min, early_termination=early_termination, history=history, 
            additional_props=additional_props, render=render, video=video_params, suffixes=suffixes
        )
        self.history_data = None
        self.history = history

    def summary(self):
        """
        Prints a summary of the hierarchy.
        """
        self.hierarchy.summary()

    def get_config(self):
        """
        Returns the configuration of the hierarchy.
        """
        return self.hierarchy.get_config()
        
    def get_model_details(self):
        """
        Returns details about the model including total nodes and total parameters in JSON format.
        
        Returns:
            dict: A dictionary containing model statistics
                - total_nodes: Total number of nodes in the hierarchy
                - total_parameters: Total number of links/parameters in the hierarchy
        """
        
        # Get total nodes directly using the hierarchy's get_num_nodes method
        total_nodes = self.hierarchy.get_num_nodes()
        
        # Get total parameters (links) using the hierarchy's get_num_links method
        total_parameters = self.hierarchy.get_num_links() - 2 * total_nodes
        
        # Create the result dictionary
        model_details = {
            "total_nodes": total_nodes,
            "total_parameters": total_parameters
        }
        
        return model_details

    def draw(self, with_labels=True, with_edge_labels=False, font_size=12, font_weight='bold', font_color='black', 
                       color_mapping={'PL':'aqua','OL':'limegreen','CL':'goldenrod', 'RL':'red', 'I':'silver', 'A':'yellow'},
                       node_size=500, arrowsize=25, align='horizontal', file=None, figsize=(8,8), move={}, draw_fig=True,
                       node_color=None, layout={'r':2,'c':1,'p':2, 'o':0}, funcdata=False, interactive_mode=False, experiment=None):
        """
        Draws the hierarchy with the specified parameters and returns the figure.
        """
        fig = self.hierarchy.draw(with_labels, with_edge_labels, font_size, font_weight, font_color, color_mapping, node_size, 
                            arrowsize, align, file, figsize, move, draw_fig, node_color, layout, funcdata, interactive_mode, experiment)
        
        return fig

    def run(self, steps=1, verbose=False):#, render=False):
        """
        Runs the hierarchy for the specified number of steps and returns the result.
        """
        # self.hierarchy.get_preprocessor()[0].set_render(render)
        return self.hierarchy.run(steps, verbose)

    def results(self, verbose=False, args=None, environment_properties=None):
        print(self.environment_properties)
        env_name = self.env.get_name()
        env_proc = EnvironmentProcessingFactory.createEnvironmentProcessing(f'{env_name}EnvironmentProcessing')
        allargs={'drive':'', 'verbosed':{'hpct_verbose':verbose}, 'max': False} | args
        env_proc.set_properties(allargs)
        env_proc.results(filepath=self.config_file, environment_properties=self.environment_properties|environment_properties, hierarchy=self.hierarchy)

    def close(self):
        """
        Closes the environment.
        """
        self.env.close()

    def get_history_keys(self):
        """
        Returns the keys of the history data of the hierarchy.
        """
        if self.history :
            if self.history_data is None:
                self.history_data = self.hierarchy.get_history_data()
            return list(self.history_data.keys())
        else:
            print("History is not enabled")

        return None

    def set_history_data(self):
        """
        Sets the history data of the hierarchy.
        """
        if self.history :

            if self.history_data is None:
                self.history_data = self.hierarchy.get_history_data()

            return self.history_data
        else:
            print("History is not enabled")

        return None

    def get_config_file_contents(self):
        """
        Reads and returns the contents of the config file.
        """
        with open(self.config_file, 'r') as file:
            return file.read()

    def plot_history(self, plots=None, title_prefix='', plots_dir=None, plots_figsize=(12, 6), history_data=None):
        """
        Plots the history of the hierarchy.
        """
        print(plots)
        plots = self.hierarchy.get_plots_config(plots, title_prefix)
        print(plots)

        for plot in plots:
            print(plot)
            plotfile=None
            if plots_dir:
                plotfile = plots_dir + os.sep + plot['title'] + '-' + str(self.hierarchy.get_namespace()) + '.png'
            fig = self.hierarchy.hierarchy_plots(title=plot['title'], plot_items=plot['plot_items'], figsize=plots_figsize, file=plotfile, history=history_data)
        return fig

    def plot_single(self, plot=None, title_prefix='', plots_dir=None, plots_figsize=(12, 6), history_data=None):
        """
        Plot one item of the history of the hierarchy.
        """

        plotfile=None
        if plots_dir:
            if not os.path.exists(plots_dir):
                os.makedirs(plots_dir)
            plotfile = plots_dir + os.sep + title_prefix + plot['title'] + '-' + str(self.hierarchy.get_namespace()) + '.png'
        fig = self.hierarchy.hierarchy_plots(title=title_prefix + plot['title'], plot_items=plot['plot_items'], figsize=plots_figsize, file=plotfile, history=history_data)
        return fig

    @classmethod
    def run_example(cls, config_file=None, 
                    run_hierarchy=True,
                    render=False, 
                    early_termination=True,
                    steps=None,
                    print_summary=False,
                    return_config=False,
                    display_usage=False,
                    verbose=False,
                    image_params=None,
                    video_params=None,
                    plot_params=None,
                    history=None,  # Allow explicit history control
                    **kwargs):
        """
        Class method to run a PCT hierarchy example with various options.
        
        Args:
            config_file (str): Path to the configuration file
            run_hierarchy (bool): Whether to run the hierarchy (default: True)
            render (bool): Whether to render the environment (default: False)
            early_termination (bool): Whether to enable early termination (default: False)
            steps (int): Number of steps to run (overrides config if specified)
            print_summary (bool): Whether to print hierarchy summary (default: False)
            return_config (bool): Whether to return the configuration (default: False)
            display_usage (bool): Whether to display usage information (default: False)
            verbose (bool): Whether to enable verbose output (default: False)
            image_params (dict or None): If not None, parameters for image creation (triggers image creation)
            video_params (dict or None): If not None, parameters for video creation (triggers video creation)
            plot_params (dict or None): If not None, parameters for plot creation (triggers plot creation)
            history (bool): Explicit control over history tracking (default: None - auto-determined)
            **kwargs: Additional parameters passed to PCTExamples constructor
        
        Returns:
            dict: Results containing requested outputs and run results
        """
        
        if display_usage:
            usage_text = """
            PCTExamples.run_example() Usage:
            
            Basic usage:
                PCTExamples.run_example('testfiles/MountainCar/MountainCar-cdf7cc1497ad143c0b04a3d9e72ab783.properties')
            
            With options:
                PCTExamples.run_example(
                    config_file='testfiles/MountainCar/MountainCar-cdf7cc1497ad143c0b04a3d9e72ab783.properties',
                    run_hierarchy=True,      # Run the hierarchy
                    render=True,             # Render environment
                    image_params={...},      # Create hierarchy diagram if not None
                    video_params={...},      # Create video if not None
                    plot_params={...},       # Create plots if not None
                    early_termination=True,  # Enable early termination
                    steps=1000,              # Override step count
                    print_summary=True,      # Print hierarchy summary
                    return_config=True,      # Return configuration
                    verbose=True             # Verbose output
                )
            
            Returns dictionary with results and any requested outputs.
            """
            print(usage_text)
            if config_file is None or (isinstance(config_file, (str, bytes)) and not config_file) or (isinstance(config_file, (tuple, list)) and not config_file):
                return {"usage_displayed": True}
            if not any([run_hierarchy, image_params is not None, video_params is not None, plot_params is not None, print_summary, return_config]):
                return {"usage_displayed": True}
        
         
        results = {}
        
        try:
            # History is only needed for plots, not videos
            # Video is controlled separately by hierarchy video property
            enable_history = False
            if history is not None:
                enable_history = history
            elif plot_params is not None:
                enable_history = True

            # Create video if video_params is not None
            if video_params is not None:
                if verbose:
                    print("=== Creating Video ===")
                
                # Set default filename if not provided
                if video_params.get('filename') is None:
                    base_name = os.path.splitext(os.path.basename(config_file))[0]
                    video_params['filename'] = f"/tmp/{base_name}_video.mp4"                

            # Create PCTExamples instance, pass video_params to constructor
            example = cls(
                config_file=config_file,
                early_termination=early_termination,
                render=render,
                history=enable_history,
                video_params=video_params,
                **kwargs
            )
            
            # Print summary if requested
            if print_summary:
                print("=== Hierarchy Summary ===")
                example.summary()
                results['summary_printed'] = True
            
            # Get configuration if requested
            if return_config:
                config = example.get_config()
                results['config'] = config
                if verbose:
                    print("=== Configuration Retrieved ===")
            
            # Create hierarchy image if image_params is not None
            if image_params is not None:
                if verbose:
                    print("=== Creating Hierarchy Image ===")
                
                # Set default filename if not provided
                if image_params.get('file') is None:
                    base_name = os.path.splitext(os.path.basename(config_file))[0]
                    image_params['file'] = f"/tmp/{base_name}_hierarchy.png"
                
                fig = example.draw(**image_params)
                results['image_created'] = True
                results['image_file'] = image_params['file']
                results['figure'] = fig
                
                if verbose:
                    print(f"Hierarchy image saved as: {image_params['file']}")
            
            # Run hierarchy if requested
            if run_hierarchy:
                if verbose:
                    print("=== Running Hierarchy ===")
                
                # Use provided steps or default from config
                run_steps = steps if steps is not None else 500
                
                run_result = example.run(steps=run_steps, verbose=verbose)
                results['run_result'] = run_result
                results['steps_completed'] = run_steps
                
                if verbose:
                    print(f"Hierarchy completed {run_steps} steps")
            
            
            # Create plots if plot_params is not None
            if plot_params is not None:
                if verbose:
                    print("=== Creating Plots ===")
                
                # Get history data for plotting
                history_data = example.set_history_data()
                if history_data:
                    if plot_params.get('single_plot') and plot_params.get('plots'):
                        # Create single plot
                        plot_fig = example.plot_single(
                            plot=plot_params['plots'],
                            title_prefix=plot_params['title_prefix'],
                            plots_dir=plot_params['plots_dir'],
                            plots_figsize=plot_params['plots_figsize'],
                            history_data=history_data
                        )
                        results['plot_figure'] = plot_fig
                        results['plot_type'] = 'single'
                    else:
                        # Create multiple plots
                        plot_fig = example.plot_history(
                            plots=plot_params['plots'],
                            title_prefix=plot_params['title_prefix'],
                            plots_dir=plot_params['plots_dir'],
                            plots_figsize=plot_params['plots_figsize'],
                            history_data=history_data
                        )
                        results['plot_figure'] = plot_fig
                        results['plot_type'] = 'multiple'
                    
                    results['plots_created'] = True
                    results['history_keys'] = example.get_history_keys()
                    
                    if verbose:
                        print("Plots created successfully")
                        print(f"Available history keys: {results['history_keys']}")
                else:
                    results['plots_error'] = "No history data available for plotting"
                    if verbose:
                        print("Warning: No history data available for plotting")
            
            # Clean up
            example.close()
            results['success'] = True
            
        except Exception as e:
            results['success'] = False
            results['error'] = str(e)
            if verbose:
                print(f"Error: {e}")
        
        return results

    # Usage examples in docstring format
    def _usage_examples():
        """
        Usage Examples:
        
        # Basic run with MountainCar config
        result = PCTExamples.run_example('testfiles/MountainCar/MountainCar-cdf7cc1497ad143c0b04a3d9e72ab783.properties')
        
        # Full featured run
        result = PCTExamples.run_example(
            config_file='testfiles/MountainCar/MountainCar-cdf7cc1497ad143c0b04a3d9e72ab783.properties',
            run_hierarchy=True,
            render=True,
            image_params={'figsize': (16, 10), 'with_labels': True},
            video_params={'fps': 60, 'filename': '/tmp/mountaincar_demo.mp4'},
            plot_params={'plots_figsize': (14, 8), 'title_prefix': 'MountainCar_'},
            early_termination=True,
            steps=5000,
            print_summary=True,
            return_config=True,
            verbose=True
        )
        
        # Just create image and get config
        result = PCTExamples.run_example(
            config_file='testfiles/MountainCar/MountainCar-cdf7cc1497ad143c0b04a3d9e72ab783.properties',
            run_hierarchy=False,
            image_params={'figsize': (16, 10), 'with_labels': True},
            plot_params={'single_plot': True, 'plots': {'title': 'position_plot', 'plot_items': ['position']}},
            return_config=True,
            print_summary=True
        )
        
        # Display usage help
        PCTExamples.run_example('', display_usage=True)
        """
        pass
