from pathlib import Path
 from typing import Optional
 
 import pandas as pd
 
-pd.set_option('display.max_rows', None)
+pd.set_option("display.max_rows", None)
 
 
 def unmangle_helper(func_name: Optional[str]):
     from cpp_demangle import demangle
+
     if func_name is None:
         return None
     if not func_name.startswith("_Z"):
         return func_name
     return demangle(func_name)
 
 
 def parse_elf(elf_path):
     from elftools.elf.elffile import ELFFile
     from elftools.elf.sections import SymbolTableSection
+
     with open(elf_path, "rb") as f:
         elffile = ELFFile(f)
 
         section = elffile.get_section_by_name(".symtab")
 

         return footprint_df
 
 
 def analyze_linker_map_helper(mapFile):
     from mapfile_parser import mapfile
+
     ret = []
     data = mapFile.toJson(humanReadable=False)
     # print("data", data)
     # input("!!!!")
     segments = data["segments"]

     assert linker_map_file.is_file()
 
     mem_footprint_df = parse_elf(elf_file)
 
     from mapfile_parser import mapfile
+
     mapFile = mapfile.MapFile()
     mapFile.readMapFile(linker_map_file)
 
     symbol_map = analyze_linker_map_helper(mapFile)
     symbol_map_df = pd.DataFrame(

             remote_program = workdir / program.name
             self.copy_to_remote(ssh, program, remote_program)
             args_str = " ".join(args)
             qemu = kwargs.get("qemu")
             pre = qemu if qemu is not None else ""
-            command = f"cd {workdir} && chmod +x {remote_program} && {pre} {remote_program} {args_str}; echo SSH EXIT=$?"
+            command = (
+                f"cd {workdir} && chmod +x {remote_program} && {pre} {remote_program} {args_str}; echo SSH EXIT=$?"
+            )
             stdin, stdout, stderr = ssh.exec_command(command)
             # print("stdin", stdin)
             # print("stdout", stdout)
             # print("stderr", stderr)
             output = stderr.read().strip() + stdout.read().strip()

         """Called at the end of a run."""
         platform = report.pre_df["Platform"]
         if (platform != "mlif").any():
             return []
         ret_artifacts = []
-        elf_artifact = lookup_artifacts(
-            artifacts, name="generic_mlonmcu", fmt=ArtifactFormat.BIN, first_only=True
-        )
+        elf_artifact = lookup_artifacts(artifacts, name="generic_mlonmcu", fmt=ArtifactFormat.BIN, first_only=True)
         assert len(elf_artifact) == 1, "ELF artifact not found!"
         elf_artifact = elf_artifact[0]
-        map_artifact = lookup_artifacts(
-            artifacts, name="generic_mlonmcu.map", fmt=ArtifactFormat.TEXT, first_only=True
-        )
+        map_artifact = lookup_artifacts(artifacts, name="generic_mlonmcu.map", fmt=ArtifactFormat.TEXT, first_only=True)
         assert len(map_artifact) == 1, "Linker map artifact not found!"
         map_artifact = map_artifact[0]
         is_ld = "ld" in map_artifact.flags
         assert is_ld, "Non ld linker currently unsupported"
 
         mem_footprint_df = parse_elf(elf_artifact.path)
 
         from mapfile_parser import mapfile
+
         mapFile = mapfile.MapFile()
         mapFile.readMapFile(map_artifact.path)
 
         symbol_map = analyze_linker_map_helper(mapFile)
         symbol_map_df = pd.DataFrame(

 
         if self.per_func:
             mem_footprint_per_func_data = generate_pie_data(mem_footprint_df, x="func", y="bytes", topk=topk)
             # print("per_func\n", mem_footprint_per_func_data, mem_footprint_per_func_data["bytes"].sum())
             if self.to_file:
-                mem_footprint_per_func_artifact = Artifact("mem_footprint_per_func.csv", content=mem_footprint_per_func_data.to_csv(index=False), fmt=ArtifactFormat.TEXT)
+                mem_footprint_per_func_artifact = Artifact(
+                    "mem_footprint_per_func.csv",
+                    content=mem_footprint_per_func_data.to_csv(index=False),
+                    fmt=ArtifactFormat.TEXT,
+                )
                 ret_artifacts.append(mem_footprint_per_func_artifact)
             if self.to_df:
                 post_df = report.post_df.copy()
                 post_df["ROM code (by func)"] = mem_footprint_per_func_data.to_dict()
                 report.post_df = post_df

         if self.per_library:
             library_footprint_df = agg_library_footprint(mem_footprint_df, symbol_map_df, by="library", col="bytes")
             mem_footprint_per_library_data = generate_pie_data(library_footprint_df, x="library", y="bytes", topk=topk)
             # print("per_library\n", mem_footprint_per_library_data, mem_footprint_per_library_data["bytes"].sum())
             if self.to_file:
-                mem_footprint_per_func_artifact = Artifact("mem_footprint_per_library.csv", content=mem_footprint_per_library_data.to_csv(index=False), fmt=ArtifactFormat.TEXT)
+                mem_footprint_per_func_artifact = Artifact(
+                    "mem_footprint_per_library.csv",
+                    content=mem_footprint_per_library_data.to_csv(index=False),
+                    fmt=ArtifactFormat.TEXT,
+                )
                 ret_artifacts.append(mem_footprint_per_func_artifact)
             if self.to_df:
                 post_df = report.post_df.copy()
                 post_df["ROM code (by library)"] = mem_footprint_per_library_data.to_dict()
                 report.post_df = post_df
             if True:  # TODO: generalize
                 # print("if1")
                 if "libmuriscvnn.a" in mem_footprint_per_library_data["library"].unique():
                     # print("if2")
-                    muriscvnn_bytes = mem_footprint_per_library_data[mem_footprint_per_library_data["library"] == "libmuriscvnn.a"]["bytes"].iloc[0]
+                    muriscvnn_bytes = mem_footprint_per_library_data[
+                        mem_footprint_per_library_data["library"] == "libmuriscvnn.a"
+                    ]["bytes"].iloc[0]
                     # print("muriscvnn_bytes", muriscvnn_bytes)
                     post_df = report.post_df.copy()
                     post_df["ROM code (libmuriscvnn.a)"] = muriscvnn_bytes
                     report.post_df = post_df
 
         if self.per_object:
             object_footprint_df = agg_library_footprint(mem_footprint_df, symbol_map_df, by="object", col="bytes")
             mem_footprint_per_object_data = generate_pie_data(object_footprint_df, x="object", y="bytes", topk=topk)
             # print("per_object\n", mem_footprint_per_object_data, mem_footprint_per_object_data["bytes"].sum())
             if self.to_file:
-                mem_footprint_per_func_artifact = Artifact("mem_footprint_per_object.csv", content=mem_footprint_per_object_data.to_csv(index=False), fmt=ArtifactFormat.TEXT)
+                mem_footprint_per_func_artifact = Artifact(
+                    "mem_footprint_per_object.csv",
+                    content=mem_footprint_per_object_data.to_csv(index=False),
+                    fmt=ArtifactFormat.TEXT,
+                )
                 ret_artifacts.append(mem_footprint_per_func_artifact)
             if self.to_df:
                 post_df = report.post_df.copy()
                 post_df["ROM code (by object)"] = mem_footprint_per_object_data.to_dict()
                 report.post_df = post_df

                         else ("Reference" if "tflm" in row.get("Backend") else "-")
                     )
                 )
                 if not (row.get("feature_muriscvnn") or row.get("feature_muriscvnnbyoc"))
                 else (
-                    "Vector" + (" (Portable)" if row.get("config_muriscvnn.use_portable") == 1 or row.get("config_muriscvnnbyoc.use_portable") == 1 else "")
+                    "Vector"
+                    + (
+                        " (Portable)"
+                        if row.get("config_muriscvnn.use_portable") == 1
+                        or row.get("config_muriscvnnbyoc.use_portable") == 1
+                        else ""
+                    )
                     if row.get("config_muriscvnn.use_vext") == 1 or row.get("config_muriscvnnbyoc.use_vext") == 1
                     else (
                         "Packed"
                         if row.get("config_muriscvnn.use_pext") == 1 or row.get("config_muriscvnnbyoc.use_pext") == 1
                         else "Scalar"

     MURISCVNN_PEXT = [["muriscvnn", "pext"] if enable_pext else []]
     CMSISNN_SCALAR = [["cmsisnn"] if enable_scalar else []]
     CMSISNN_DSP = [["cmsisnn", "arm_dsp"] if enable_dsp else []]
     CMSISNN_MVEI = [["cmsisnn", "arm_mvei", "arm_dsp"] if enable_mvei else []]
     TARGET_FEATURES = {
-        **{t: [
-            *([*DEFAULT_SCALAR, *DEFAULT_VEXT, *DEFAULT_PEXT] if enable_default else []),
-            *([*MURISCVNN_SCALAR, *MURISCVNN_VEXT, *MURISCVNN_PEXT] if enable_muriscvnn else []),
-            *([*CMSISNN_SCALAR] if enable_cmsisnn else []),
-        ] for t in SPIKE_TARGETS},
+        **{
+            t: [
+                *([*DEFAULT_SCALAR, *DEFAULT_VEXT, *DEFAULT_PEXT] if enable_default else []),
+                *([*MURISCVNN_SCALAR, *MURISCVNN_VEXT, *MURISCVNN_PEXT] if enable_muriscvnn else []),
+                *([*CMSISNN_SCALAR] if enable_cmsisnn else []),
+            ]
+            for t in SPIKE_TARGETS
+        },
         "canmv_k230_ssh": [
             *([*DEFAULT_SCALAR, *DEFAULT_VEXT] if enable_default else []),
             *([*MURISCVNN_SCALAR, *MURISCVNN_VEXT] if enable_muriscvnn else []),
             *([*CMSISNN_SCALAR] if enable_cmsisnn else []),
         ],

 BACKEND_DEFAULT_CONFIG = {
     "tflmi": {},
     "tvmaot": {"usmp.algorithm": "hill_climb"},
     "tvmaotplus": {"usmp.algorithm": "hill_climb"},
     "tvmrt": {},
-    "tvmllvm": {}
+    "tvmllvm": {},
 }
 
 VLENS = [64, 128, 256, 512, 1024, 2048, 4096, 8192]
 
 DEFAULT_VLENS = [64, 128, 256, 512, 1024, 2048]

                                                                 run.add_model_by_name(model, context=context)
                                                                 run.add_backend_by_name(backend, context=context)
                                                                 run.add_target_by_name(target, context=context)
                                                                 if args.post:
                                                                     run.add_postprocesses_by_name(POSTPROCESSES_0)
-                                                                    run.add_postprocess(CustomPostprocess(), append=True)
+                                                                    run.add_postprocess(
+                                                                        CustomPostprocess(), append=True
+                                                                    )
                                                                     run.add_postprocesses_by_name(
                                                                         POSTPROCESSES_1, append=True
                                                                     )
                                                                     if args.baseline is not None:
                                                                         run.add_postprocess_by_name(

             if args.noop:
                 stage = RunStage.LOAD
             else:
                 stage = RunStage.COMPILE if MEM_ONLY else RunStage.RUN
             success = session.process_runs(
-                until=stage, num_workers=args.parallel, progress=args.progress, context=context, per_stage=args.runs_per_stage,
+                until=stage,
+                num_workers=args.parallel,
+                progress=args.progress,
+                context=context,
+                per_stage=args.runs_per_stage,
             )
             report = session.get_reports()
         report_file = args.output
         report.export(report_file)
         print()