import importlib
import os
import sys

import onnxruntime as ort
import torch

# CONFIG SECTION
config = {
    # path to .pth kaloscope lsnet file
    "checkpoint": "/content/Kaloscope/best_checkpoint.pth",
    # Model factory name inside model/lsnet_artist.py
    # (examples: lsnet_t_artist, lsnet_s_artist, lsnet_b_artist, lsnet_xl_artist)
    # lsnet_xl_artist is the one use for kaloscope initial release apparently
    "model_name": "lsnet_xl_artist",
    # where to save exported ONNX model
    "onnx_output": "/content/lsnet_xl_artist-dynamo-opset20.onnx",
    # Image input size (height, width) - derived from model repo readme
    "img_size": (224, 224),
    # ONNX opset version - ignore this, set it below in torch.onnx.export()
    "opset": 18,
    "device": "cpu",
}
# ------------------------------------------------------------------------------------


def strip_prefix_from_state_dict(state_dict, prefix="module."):
    new_state = {}
    for k, v in state_dict.items():
        if k.startswith(prefix):
            new_state[k[len(prefix) :]] = v
        else:
            new_state[k] = v
    return new_state


def load_checkpoint(checkpoint_path, map_location="cpu"):
    ck = torch.load(checkpoint_path, map_location=map_location, weights_only=False)
    if isinstance(ck, dict):
        for key in ("state_dict", "model_state_dict", "model"):
            if key in ck:
                return ck[key], ck
        return ck, ck
    else:
        return None, ck


def main(cfg):
    # module names might be wrong here but the files/modules are in lsnet folder in same dir as this file
    modules_to_delete = ["lsnet.lsnet_artist", "lsnet.ska"]
    for module_name in modules_to_delete:
        if module_name in sys.modules:
            print(f"Forcefully deleting cached module: {module_name}")
            del sys.modules[module_name]

    repo_root = os.getcwd()
    if repo_root not in sys.path:
        sys.path.insert(0, repo_root)

    model_module_name = "lsnet.lsnet_artist"
    try:
        mod = importlib.import_module(model_module_name)

        # originally intended to make google colab not cache this but doesnt work as expected
        importlib.reload(mod)
        if "lsnet.ska" in sys.modules:
            importlib.reload(sys.modules["lsnet.ska"])

    except Exception as e:
        print(f"WARNING - could not import {model_module_name}: {e}")
        mod = None
    model = None
    if mod is not None and hasattr(mod, cfg["model_name"]):
        factory = getattr(mod, cfg["model_name"])
        if callable(factory):
            try:
                # Attempt to instantiate with the number of classes from your checkpoint
                model = factory(num_classes=31770)
                print(f"Instantiated {cfg['model_name']}(num_classes=31770)")
            except Exception as e:
                print(f"Failed to instantiate model with 31770 classes: {e}")
                # Fallback to the original logic if the above fails
                try:
                    model = factory()
                    print(f"Instantiated model using {model_module_name}.{cfg['model_name']}()")
                except TypeError:
                    for ncls in [1000, 100]:
                        try:
                            model = factory(num_classes=ncls)
                            print(f"Instantiated {cfg['model_name']}(num_classes={ncls})")
                            break
                        except Exception:
                            pass
    if model is None:
        print("Model instance not created from source; will try to load pickled model from checkpoint.")

    state_dict, full_ck = load_checkpoint(cfg["checkpoint"], map_location=cfg["device"])
    if state_dict is None and isinstance(full_ck, torch.nn.Module):
        print("Checkpoint is a pickled model object - exporting directly.")
        model = full_ck
    else:
        if model is None:
            if isinstance(full_ck, dict):
                for k in ("model", "net", "module"):
                    if k in full_ck and isinstance(full_ck[k], torch.nn.Module):
                        model = full_ck[k]
                        print(f'Using model found in checkpoint["{k}"]')
                        break

        if model is not None and state_dict is not None:
            state_dict = strip_prefix_from_state_dict(state_dict)
            try:
                model.load_state_dict(state_dict, strict=False)
                print("Loaded state_dict into model (strict=False).")
            except Exception as e:
                print("Warning: load_state_dict failed:", e)
        elif state_dict is not None and model is None:
            raise RuntimeError("Found a state_dict but no model instance to load it into.")

    device = torch.device(cfg["device"])
    model.to(device)
    model.eval()

    h, w = cfg["img_size"]
    dummy_input = torch.randn(1, 3, h, w, device=device)

    onnx_path = cfg["onnx_output"]
    print("Exporting to ONNX ->", onnx_path)
    try:
        torch.onnx.export(
            model,
            dummy_input,
            onnx_path,
            input_names=["input"],
            output_names=["output"],
            # opset_version=cfg["opset"],
            opset_version=None,  # leaving none for auto recommended version, see print(torch.onnx._constants.ONNX_DEFAULT_OPSET) eg torch 2.8.0 = 18 and 2.9.0 = 20 with up to 23 supported
            # if you want dynamic axes, re-enable this AND set dynamo=False
            # dynamic_axes={
            #    "input": {0: "batch_size", 2: "height", 3: "width"},
            #    "output": {0: "batch_size"},
            # },
            verbose=False,
            do_constant_folding=True,
            optimize=True,
            # dynamo here
            dynamo=True,
        )
        print("ONNX export completed.")
    except Exception as e:
        print("ONNX export failed:", e)
        raise

    print("Verifying ONNX model with onnxruntime...")
    sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
    inp_name = sess.get_inputs()[0].name
    out = sess.run(None, {inp_name: dummy_input.cpu().numpy()})
    print("ONNX runtime produced output shapes:", [o.shape for o in out])
    print("Done. Saved:", onnx_path)


if __name__ == "__main__":
    main(config)