
import torch
from torchvision.models import mobilenet_v2  # @manual
from torchvision.models.mobilenetv2 import MobileNet_V2_Weights

#=========================================== helper function ===========================================================
# Get example input
def get_example_inputs(use_real_input):
    tensor_size = (1, 3, 224, 224)
    input_batch = (torch.randn(tensor_size),)
    if use_real_input == True:
        print("Loaded real input image dog.jpg")
        import urllib

        url, filename = (
            "https://github.com/pytorch/hub/raw/master/images/dog.jpg",
            "dog.jpg",
        )
        try:
            urllib.URLopener().retrieve(url, filename)
        except:
            urllib.request.urlretrieve(url, filename)
        from PIL import Image
        from torchvision import transforms

        input_image = Image.open(filename)
        preprocess = transforms.Compose(
            [
                transforms.Resize(256),
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                transforms.Normalize(
                    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
                ),
            ]
        )
        input_tensor = preprocess(input_image)
        input_batch = input_tensor.unsqueeze(0)
        input_batch = (input_batch,)
    return input_batch

# To use Cortex-M backend
from executorch.backends.cortex_m.passes.convert_to_cortex_m_pass import (
    ConvertToCortexMPass,
)

from executorch.backends.cortex_m.passes.quantized_op_fusion_pass import (
    QuantizedOpFusionPass,
)

from executorch.backends.cortex_m.passes.replace_quant_nodes_pass import (
    ReplaceQuantNodesPass,
)

def transform_for_cortex_m_backend(edge_program_manager):
    # Let's make sure we are using optimized Cortex M backend
    # NB: If we can't find and replace ops those are expected to be replaced,
    # bad things will happen at runtime, like "missing operator" errors!

    # Instantiate the mandatory ReplaceQuantNodesPass
    passes = [ReplaceQuantNodesPass]
    # Quantized qdq fusion Op passes
    # passes += [ConvertToCortexMPass, QuantizedOpFusionPass]

    current_edge = edge_program_manager
    for pass_cls in passes:
        transform_pass = (
            pass_cls(current_edge.exported_program())
            if pass_cls.__name__ == "QuantizedLinearFusionPass"
            else pass_cls()
        )
        current_edge = current_edge.transform([transform_pass])
    return current_edge

#=========================================== AOT flow ===========================================================

# Get mobilenetv2 form torchvision
original_model = mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT)
print(type(original_model))

example_inputs = get_example_inputs(True)

model = original_model.eval()

exported_program = torch.export.export(model, example_inputs, strict=False)
graph_module = exported_program.module(check_guards=False)
#_ = graph_module.print_readable()

# We only support Python3.10 and above, so use a later pickle protocol
torch.export.save(
    exported_program,
    f"mv2_exported_program.pt2",
    pickle_protocol=5,
)


#================================================ quantization steps ==============================================================
from executorch.backends.arm.ethosu import EthosUCompileSpec
from executorch.backends.arm.util._factory import create_partitioner, create_quantizer

from executorch.backends.arm.quantizer import (
    EthosUQuantizer,
    get_symmetric_quantization_config,
    get_symmetric_a16w8_quantization_config,
)
from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e

# Create a compilation spec describing the target for configuring the quantizer
# Some args are used by the Arm Vela graph compiler later in the example. Refer to Arm Vela documentation for an
# explanation of its flags: https://gitlab.arm.com/artificial-intelligence/ethos-u/ethos-u-vela/-/blob/main/OPTIONS.md
compile_spec = EthosUCompileSpec(
            target="ethos-u55-256",
            system_config="Ethos_U55_High_End_Embedded",
            memory_mode="Shared_Sram",
            extra_flags=["--verbose-operators", "--verbose-cycle-estimate"],
            config_ini="Arm/vela.ini"
        )

# Create and configure quantizer to use a symmetric quantization config globally on all nodes
print(compile_spec)
quantizer = create_quantizer(compile_spec)

operator_config = get_symmetric_quantization_config(is_per_channel=True)
quantizer.set_global(operator_config)

# Post training quantization
quantized_graph_module = prepare_pt2e(graph_module, quantizer)
quantized_graph_module(*example_inputs) # Calibrate the graph module with the example input
quantized_graph_module = convert_pt2e(quantized_graph_module)
#print(quantized_graph_module)

# Create a new exported program using the quantized_graph_module
quantized_exported_program = torch.export.export(quantized_graph_module, example_inputs)

torch.export.save(
    quantized_exported_program,
    f"mv2_quantize_exported_program.pt2",
    pickle_protocol=5,
)


#============================================== EthosU  backend steps ==============================================================

from executorch.backends.arm.ethosu import EthosUPartitioner
from executorch.exir import (
    EdgeCompileConfig,
    ExecutorchBackendConfig,
    to_edge_transform_and_lower,
)
from executorch.extension.export_util.utils import save_pte_program

# Create partitioner from compile spec
partitioner = EthosUPartitioner(compile_spec)

# Lower the exported program to the Ethos-U backend
edge_program_manager = to_edge_transform_and_lower(
            quantized_exported_program,
            partitioner=[partitioner],
            compile_config=EdgeCompileConfig(
                _check_ir_validity=False,
            ),
        )


# Transform so we can use ops from the Cortex M backend
edge_program_manager = transform_for_cortex_m_backend(edge_program_manager)


# Convert edge program to executorch
executorch_program_manager = edge_program_manager.to_executorch(
            config=ExecutorchBackendConfig(extract_delegate_segments=False)
        )

_ = executorch_program_manager.exported_program().graph_module.print_readable()

# Save pte file
save_pte_program(executorch_program_manager, "mv2_arm_delegate_ethos-u55-256.pte")