##############################################################################bl
# MIT License
#
# Copyright (c) 2021 - 2024 Advanced Micro Devices, Inc. All Rights Reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
##############################################################################el

from omniperf_analyze.analysis_base import OmniAnalyze_Base
from utils.utils import demarcate
from utils import file_io, parser
from utils.gui import build_bar_chart, build_table_chart

import os
import logging
import random
import copy
import dash
import dash_bootstrap_components as dbc
from dash import html
from dash import dcc
from dash.dependencies import Input, Output, State


class webui_analysis(OmniAnalyze_Base):
    def __init__(self, args, supported_archs):
        super().__init__(args, supported_archs)
        self.app = dash.Dash(__name__, external_stylesheets=[dbc.themes.CYBORG])
        self.dest_dir = os.path.abspath(args.path[0][0])
        self.arch = None
        
        self.__hidden_sections = ["Memory Chart", "Kernels"]
        self.__hidden_columns = ["Tips", "coll_level"]
        # define different types of bar charts
        self.__barchart_elements = {
            "instr_mix": [1001, 1002],
            "multi_bar": [1604, 1704],
            "sol": [1101, 1201, 1301, 1401, 1601, 1701],
            #"l2_cache_per_chan": [1802, 1803]
        }
        # define any elements which will have full width
        self.__full_width_elements = {1801}
        

    @demarcate
    def build_layout(
        self,
        input_filters,
        arch_configs
    ):
        """
        Build gui layout
        """
        from utils.gui_components.header import get_header
        from utils.gui_components.memchart import get_memchart

        comparable_columns = parser.build_comparable_columns(self.get_args().time_unit)
        base_run, base_data = next(iter(self._runs.items()))
        self.app.layout = html.Div(style={"backgroundColor": "rgb(50, 50, 50)"})

        filt_kernel_names = []
        kernel_top_df = base_data.dfs[1]
        for kernel_id in base_data.filter_kernel_ids:
            filt_kernel_names.append(kernel_top_df.loc[kernel_id, "Kernel_Name"])
        
        self.app.layout.children = html.Div(
            children=[
                dbc.Spinner(
                    children=[
                        get_header(base_data.raw_pmc, input_filters, filt_kernel_names),
                        html.Div(id="container", children=[]),
                    ],
                    fullscreen=True,
                    color="primary",
                    spinner_style={"width": "6rem", "height": "6rem"},
                )
            ]
        )

        @self.app.callback(
            Output("container", "children"),
            [Input("disp-filt", "value")],
            [Input("kernel-filt", "value")],
            [Input("gcd-filt", "value")],
            [Input("norm-filt", "value")],
            [Input("top-n-filt", "value")],
            [State("container", "children")],
        )
        def generate_from_filter(
            disp_filt, kernel_filter, gcd_filter, norm_filt, top_n_filt, div_children
        ):
            logging.debug("[analysis] gui normalization is %s" % norm_filt)

            base_data = self.initalize_runs() # Re-initalizes everything
            panel_configs = copy.deepcopy(arch_configs.panel_configs)
            # Generate original raw df
            base_data[base_run].raw_pmc = file_io.create_df_pmc(self.dest_dir, self.get_args().verbose)
            logging.debug("[analysis] gui dispatch filter is %s" % disp_filt)
            logging.debug("[analysis] gui kernel filter is %s" % kernel_filter)
            logging.debug("[analysis] gui gpu filter is %s" % gcd_filter)
            logging.debug("[analysis] gui top-n filter is %s" % top_n_filt)
            base_data[base_run].filter_kernel_ids = kernel_filter
            base_data[base_run].filter_gpu_ids = gcd_filter
            base_data[base_run].filter_dispatch_ids = disp_filt
            base_data[base_run].filter_top_n = top_n_filt

            # Reload the pmc_kernel_top.csv for Top Stats panel
            file_io.create_df_kernel_top_stats(
                raw_data_dir=str(self.dest_dir),
                filter_gpu_ids=base_data[base_run].filter_gpu_ids,
                filter_dispatch_ids=base_data[base_run].filter_dispatch_ids,
                time_unit=self.get_args().time_unit,
                max_kernel_num=base_data[base_run].filter_top_n,
            )
            # Only display basic metrics if no filters are applied
            if not (disp_filt or kernel_filter or gcd_filter):
                temp = {}
                keep = [1, 201, 101, 301]
                for key in base_data[base_run].dfs:
                    if keep.count(key) != 0:
                        temp[key] = base_data[base_run].dfs[key]

                base_data[base_run].dfs = temp
                temp = {}
                keep = [0, 100, 200, 300]
                for key in panel_configs:
                    if keep.count(key) != 0:
                        temp[key] = panel_configs[key]
                panel_configs = temp
            # All filtering will occur here
            parser.load_table_data(
                workload=base_data[base_run], 
                dir=self.dest_dir, 
                is_gui=True,
                debug=self.get_args().debug,
                verbose=self.get_args().verbose,
            )

            # ~~~~~~~~~~~~~~~~~~~~~~~
            # Generate GUI content
            # ~~~~~~~~~~~~~~~~~~~~~~~
            div_children = []

            # Append memory chart and roofline
            div_children.append(
                get_memchart(panel_configs[300]["data source"], base_data[base_run])
            )
            has_roofline = os.path.isfile(os.path.join(self.dest_dir, "roofline.csv"))
            if has_roofline and hasattr(self.get_socs()[self.arch], "roofline_obj"):
                # update roofline for visualization in GUI
                self.get_socs()[self.arch].analysis_setup(
                    roofline_parameters={
                        'path_to_dir': self.dest_dir,
                        'device_id': 0,
                        'sort_type': 'kernels',
                        'mem_level': 'ALL',
                        'include_kernel_names': False,
                        'is_standalone': False
                    }
                )
                roof_obj = self.get_socs()[self.arch].roofline_obj
                div_children.append(
                    roof_obj.empirical_roofline(
                        ret_df=parser.apply_filters(workload=base_data[base_run], dir=self.dest_dir, is_gui=True, debug=self.get_args().debug)
                    )
                )
            
            # Iterate over each section as defined in panel configs
            for panel_id, panel in panel_configs.items():
                title = str(panel_id // 100) + ". " + panel["title"]
                section_title = (
                    panel["title"]
                    .replace("(", "")
                    .replace(")", "")
                    .replace("/", "")
                    .replace(" ", "_")
                    .lower()
                )
                html_section = []

                if panel["title"] not in self.__hidden_sections:
                    # Iterate over each table per section
                    for data_source in panel["data source"]:
                        for t_type, table_config in data_source.items():
                            original_df = base_data[base_run].dfs[table_config["id"]]
                            # The sys info table need to add index back
                            if t_type == "raw_csv_table" and "Info" in original_df.keys():
                                original_df.reset_index(inplace=True)
                            
                            content = determine_chart_type(
                                original_df=original_df,
                                table_config=table_config,
                                hidden_columns=self.__hidden_columns,
                                barchart_elements=self.__barchart_elements,
                                norm_filt=norm_filt,
                                comparable_columns=comparable_columns,
                                decimal=self.get_args().decimal
                            )

                            # Update content for this section
                            if table_config["id"] in self.__full_width_elements:
                                # Optionally override default (50%) width
                                html_section.append(
                                    html.Div(
                                        className="float-child",
                                        children=content,
                                        style={"width": "100%"},
                                    )
                                )
                            else:
                                html_section.append(
                                    html.Div(className="float-child", children=content)
                                )

                    # Append the new section with all of it's contents
                    div_children.append(
                        html.Section(
                            id=section_title,
                            children=[
                                html.H3(
                                    children=title,
                                    style={"color": "white"},
                                ),
                                html.Div(className="float-container", children=html_section),
                            ],
                        )
                    )

            # Display pop-up message if no filters are applied
            if not (disp_filt or kernel_filter or gcd_filter):
                div_children.append(
                    html.Section(
                        id="popup",
                        children=[
                            html.Div(
                                children="To dive deeper, use the top drop down menus to isolate particular kernel(s) or dispatch(s). You will then see the web page update with additional low-level metrics specific to the filter you've applied.",
                            ),
                        ],
                    )
                )
            
            return div_children


    #-----------------------
    # Required child methods
    #-----------------------
    @demarcate
    def pre_processing(self):
        """Perform any pre-processing steps prior to analysis.
        """
        super().pre_processing()
        if len(self._runs) == 1:
            args = self.get_args()
            file_io.create_df_kernel_top_stats(
                raw_data_dir=self.dest_dir,
                filter_gpu_ids=self._runs[self.dest_dir].filter_gpu_ids,
                filter_dispatch_ids=self._runs[self.dest_dir].filter_dispatch_ids,
                time_unit=args.time_unit,
                max_kernel_num=args.max_kernel_num,
            )
            # create 'mega dataframe'
            self._runs[self.dest_dir].raw_pmc = file_io.create_df_pmc(
                self.dest_dir, args.verbose
            )
            # create the loaded kernel stats
            parser.load_kernel_top(self._runs[self.dest_dir], self.dest_dir)
            # set architecture
            self.arch = self._runs[self.dest_dir].sys_info.iloc[0]['gpu_soc']
            
        else:
            self.error("Multiple runs not yet supported in GUI. Retry without --gui flag.")

    
    @demarcate
    def run_analysis(self):
        """Run CLI analysis.
        """
        super().run_analysis()
        args = self.get_args()
        input_filters = {
            "kernel": self._runs[self.dest_dir].filter_kernel_ids,
            "gpu": self._runs[self.dest_dir].filter_gpu_ids,
            "dispatch": self._runs[self.dest_dir].filter_dispatch_ids,
            "normalization": args.normal_unit,
            "top_n": args.max_kernel_num,
        }
        
        self.build_layout(
            input_filters,
            self._arch_configs[self.arch],
        )
        if args.random_port:
            self.app.run_server(debug=False, host="0.0.0.0", port=random.randint(1024, 49151))
        else:
            self.app.run_server(debug=False, host="0.0.0.0", port=args.gui)


@demarcate
def determine_chart_type(
        original_df,
        table_config,
        hidden_columns,
        barchart_elements,
        norm_filt,
        comparable_columns,
        decimal
    ):
    content = []

    display_columns = original_df.columns.values.tolist().copy()
    # Remove hidden columns. Better way to do it?
    for col in hidden_columns:
        if col in display_columns:
            display_columns.remove(col)
    display_df = original_df[display_columns]

    # Determine chart type:
    # a) Barchart
    if table_config["id"] in [
        x for i in barchart_elements.values() for x in i
    ]:
        d_figs = build_bar_chart(display_df, table_config, barchart_elements, norm_filt)
        # Smaller formatting if barchart yeilds several graphs
        if (
            len(d_figs) > 2
            and not table_config["id"]
            # in barchart_elements["l2_cache_per_chan"]
        ):
            temp_obj = []
            for fig in d_figs:
                temp_obj.append(
                    html.Div(
                        className="float-child",
                        children=[
                            dcc.Graph(
                                figure=fig, style={"margin": "2%"}
                            )
                        ],
                    )
                )
            content.append(
                html.Div(
                    className="float-container", children=temp_obj
                )
            )
        # Normal formatting if < 2 graphs
        else:
            for fig in d_figs:
                content.append(
                    dcc.Graph(figure=fig, style={"margin": "2%"})
                )
    # B) Tablechart
    else:
        d_figs = build_table_chart(
            display_df,
            table_config,
            original_df,
            display_columns,
            comparable_columns,
            decimal,
        )
        for fig in d_figs:
            content.append(html.Div([fig], style={"margin": "2%"}))

    # subtitle for each table in a panel if existing
    if "title" in table_config and table_config["title"]:
        subtitle = (
            str(table_config["id"] // 100)
            + "."
            + str(table_config["id"] % 100)
            + " "
            + table_config["title"]
            + "\n"
        )

        content.insert(
            0,
            html.H4(
                children=subtitle,
                style={"color": "white"},
            ),
        )
    return content