from pathlib import Path
from typing import BinaryIO, Sequence, TextIO, Union

import torch
from torch.fx.passes.split_utils import getattr_recursive
from torch.package import PackageExporter

"""
A tool to package acc submodule as a torch package. The packaged model can be loaded
with just PyTorch library.
"""


def flatten_model(model: torch.fx.GraphModule) -> torch.fx.GraphModule:
    """
    Remove all original modules with an attr holder module so that all original modules
    and names are not present.
    """
    holder_module = torch.nn.Module()
    model._holder = holder_module
    attr_id = 0

    for node in model.graph.nodes:
        assert node.op != "call_module"
        if node.op == "get_attr":
            attr = getattr_recursive(model, node.target)
            setattr(holder_module, f"_attr_{attr_id}", attr)
            with model.graph.inserting_before(node):
                new_node = model.graph.get_attr(f"_holder._attr_{attr_id}")
                node.replace_all_uses_with(new_node)
            attr_id += 1

    model.graph.eliminate_dead_code()
    model.recompile()
    model.delete_all_unused_submodules()
    return model


def generate_standalone_repro(
    model: torch.fx.GraphModule, output: Union[str, Path, TextIO], prelude: str = ""
) -> None:
    """
    Generate a standalone python file for the model where weights are randomized
    and the model flattened.
    This only works if leaf nodes are only torch.nn modules.
    """
    model = flatten_model(model)

    INDENT = "    "
    lines = [
        "",
        "import torch",
        "from torch import nn",
    ]
    code = str(model.code)

    import_modules = set()
    import_map = {
        "torch_tensorrt_fx_tracer_acc_tracer_acc_ops": "torch_tensorrt.fx.tracer.acc_tracer.acc_ops",
        "torch_tensorrt_fx_passes_lower_basic_pass": "torch_tensorrt.fx.passes.lower_basic_pass",
    }
    for line in code.split("\n"):
        for k, v in import_map.items():
            if k in line:
                sub_string = line.split("(")[0].split()[-1]
                if sub_string.startswith(k):
                    mod = sub_string.replace(k + "_", "")
                    import_modules.add(
                        "from " + v + " import " + mod + " as " + sub_string
                    )
    for mod in sorted(import_modules):
        lines.append(mod)

    lines += [
        "",
        "",
        "class ExportedModule(nn.Module):",
        f"{INDENT}def __init__(self):",
        f"{INDENT * 2}super().__init__()",
    ]

    for k, v in model._holder.named_parameters():
        shape = ", ".join([str(i) for i in v.shape])
        rand_func = "randn" if torch.is_floating_point(v) else "randint"
        int_range = "" if torch.is_floating_point(v) else "0, 5, "
        lines.append(
            f"{INDENT * 2}self.{k} = nn.Parameter(torch.{rand_func}({int_range}{shape}, dtype={v.dtype}))"
        )

    def dump(f):
        f.write(prelude)
        f.write("\n".join(lines))
        f.write(
            "\n".join(
                [
                    INDENT + line.replace("self._holder.", "self.")
                    for line in code.split("\n")
                ]
            )
        )
        f.write("\n")

    if isinstance(output, (Path, str)):
        with open(str(output), "w") as f:
            dump(f)
    else:
        dump(output)


class ModelPackager:
    @classmethod
    def set_extern_modules(cls, pe: PackageExporter) -> None:
        pe.extern(
            [
                "builtins",
                "sys",
                "torch.**",
            ]
        )

    @classmethod
    def set_mocked_modules(cls, pe: PackageExporter):
        pe.mock(
            "**",
            exclude=[
                "torch_tensorrt.fx.tracer.acc_tracer.acc_ops",
                "torch_tensorrt.fx.tracer.acc_tracer.acc_normalizer",
                "torch_tensorrt.fx.tracer.acc_tracer.acc_op_properties",
            ],
        )

    @classmethod
    def package_model(
        cls,
        model: torch.nn.Module,
        model_inputs: Sequence[torch.Tensor],
        output: Union[str, Path, BinaryIO],
        preserve_model_structure: bool = False,
    ) -> None:
        if not preserve_model_structure:
            model = flatten_model(model)
        with PackageExporter(output) as pe:
            cls.set_extern_modules(pe)
            cls.set_mocked_modules(pe)
            pe.intern("**")
            pe.save_pickle("repro", "model", model)
            pe.save_pickle("repro", "inputs", model_inputs)