if __name__ != "__main__":
    raise ImportError("This file is not intended to be imported!")

import inspect
import subprocess

from torch import nn

DOCS = """
This file is automatically generated by torchmix and should not be changed.

This module contains wrapped versions of the objects in the PyTorch `nn`
module, intended to provide a more user-friendly interface and added
configuration management functionality.

These classes can be used as a drop-in replacement for the original PyTorch
`nn` module, allowing you to access the full configuration of the module
through the config attribute.
"""


def lazy_conv_template(name):
    f"""with suppress(AttributeError):
    class {name}(MixModule, nn.{name}):
        __init__ = functools.partial(nn.{name}.__init__)  # type: ignore[misc]"""


def sequential_template(name):
    return f"""with suppress(AttributeError):
    class {name}(MixModule, nn.{name}):
        build_mode = BuildMode.WITH_ARGS
        __init__ = functools.partial(nn.{name}.__init__)  # type: ignore"""


def module_template(name):
    return f"""with suppress(AttributeError):
    class {name}(MixModule, nn.{name}):
        __init__ = functools.partial(nn.{name}.__init__)"""


with open("torchmix/nn/__init__.py", "w") as f:
    f.write('"""')
    f.write(DOCS)
    f.write('"""\n\n')

    def write(line=None):
        return f.write(line + "\n" if line else "\n")

    write("from contextlib import suppress")
    write("import functools")
    write()
    write("from torch import nn")
    write()
    write("from torchmix.core._builds import BuildMode")
    write("from torchmix.core._module import MixModule")
    write()
    write()
    write("Module = nn.Module")
    write("Parameter = nn.Parameter")
    write("UninitializedBuffer = nn.UninitializedBuffer")
    write("UninitializedParameter = nn.UninitializedParameter")
    write()

    for name, obj in inspect.getmembers(nn):
        if name in ("Module", "Parameter") or name.startswith(("_", "__")):
            continue
        try:
            if issubclass(obj, nn.Module):
                pass
        except:  # noqa: E722
            if name not in (
                "common_types",
                "grad",
                "intrinsic",
                "qat",
                "quantizable",
                "quantized",
            ):
                write("with suppress(AttributeError):")
                write(f"    {name} = nn.{name}\n")

    processed = set()
    for name, obj in inspect.getmembers(nn):
        if name in ("Module", "Parameter") or name.startswith(("_", "__")):
            continue

        try:
            template: str
            if name == "Sequential":
                template = sequential_template(name)
            elif name.startswith("LazyConv"):
                template = lazy_conv_template(name)
            elif issubclass(obj, nn.Module):
                template = module_template(name)
            else:
                raise ValueError(f"{name} not handled")

            write()
            write()
            write(template)

            processed.add(name)

        except:  # noqa: E722
            pass

try:
    subprocess.run(["poetry", "run", "black", "-l", "78", "torchmix/nn/__init__.py"])
except:  # noqa: E722
    pass
