# export_onnx.py
import torch
import torch.nn as nn
import onnx
from onnx import shape_inference
import onnxruntime as ort
import numpy as np
from typing import Tuple


class ExportableModel(nn.Module):
    """Example model for edge deployment."""

    def __init__(self, num_classes: int = 10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, 3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.AdaptiveAvgPool2d(1),
        )
        self.classifier = nn.Linear(128, num_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x


def export_to_onnx(
    model: nn.Module,
    save_path: str = "model.onnx",
    input_shape: Tuple[int, ...] = (1, 3, 224, 224),
    opset_version: int = 17,
    dynamic_axes: bool = True,
) -> str:
    """Export PyTorch model to ONNX format."""
    model.eval()

    dummy_input = torch.randn(*input_shape)

    axes = {}
    if dynamic_axes:
        axes = {0: "batch_size", 2: "height", 3: "width"}

    torch.onnx.export(
        model,
        dummy_input,
        save_path,
        opset_version=opset_version,
        input_names=["input"],
        output_names=["output"],
        dynamic_axes={"input": axes, "output": axes} if dynamic_axes else None,
    )

    onnx_model = onnx.load(save_path)
    onnx_model = shape_inference.infer_shapes(onnx_model)
    onnx.checker.check_model(onnx_model)

    print(f"ONNX model exported to {save_path}")
    print(f"Model size: {onnx_model.ByteSize() / 1e6:.2f} MB")

    return save_path


def validate_onnx_export(
    pytorch_model: nn.Module,
    onnx_path: str,
    input_shape: Tuple[int, ...] = (1, 3, 224, 224),
    atol: float = 1e-5,
) -> bool:
    """Validate ONNX export matches PyTorch output."""
    pytorch_model.eval()
    dummy_input = torch.randn(*input_shape)

    with torch.no_grad():
        pytorch_output = pytorch_model(dummy_input).numpy()

    session = ort.InferenceSession(onnx_path)
    onnx_output = session.run(
        None, {"input": dummy_input.numpy()}
    )[0]

    match = np.allclose(pytorch_output, onnx_output, atol=atol)
    max_diff = np.max(np.abs(pytorch_output - onnx_output))

    print(f"ONNX validation: {'PASS' if match else 'FAIL'}")
    print(f"Max absolute difference: {max_diff:.8f}")

    return match

# quantize_model.py
import onnx
import onnxruntime as ort
from onnxruntime.quantization import (
    quantize_static,
    quantize_dynamic,
    CalibrationDataReader,
    QuantType,
    CalibrationMethod,
)
import numpy as np
from typing import List, Tuple


class CalibrationDataReaderImpl(CalibrationDataReader):
    """Custom calibration data reader for quantization."""

    def __init__(self, calibration_data: List[np.ndarray]):
        self.calibration_data = calibration_data
        self.current_index = 0

    def get_next(self) -> dict:
        if self.current_index >= len(self.calibration_data):
            return None

        data = self.calibration_data[self.current_index]
        self.current_index += 1
        return {"input": data}

    def rewind(self):
        self.current_index = 0


class QuantizationPipeline:
    """Complete quantization pipeline for edge deployment."""

    def __init__(self, onnx_path: str):
        self.onnx_path = onnx_path
        self.quantized_models = {}

    def dynamic_quantization(
        self,
        weight_type: QuantType = QuantType.QUInt8,
        save_path: str = "model_dynamic_int8.onnx",
    ) -> str:
        """Apply dynamic quantization (weights only)."""
        quantize_dynamic(
            model_input=self.onnx_path,
            model_output=save_path,
            weight_type=weight_type,
        )

        self.quantized_models["dynamic_int8"] = save_path
        print(f"Dynamic quantization saved to {save_path}")
        return save_path

    def static_quantization(
        self,
        calibration_data: List[np.ndarray],
        activation_type: QuantType = QuantType.QUInt8,
        weight_type: QuantType = QuantType.QUInt8,
        save_path: str = "model_static_int8.onnx",
    ) -> str:
        """Apply static quantization with calibration."""
        data_reader = CalibrationDataReaderImpl(calibration_data)

        quantize_static(
            model_input=self.onnx_path,
            model_output=save_path,
            calibration_data_reader=data_reader,
            activation_type=activation_type,
            weight_type=weight_type,
            calibration_method=CalibrationMethod.MinMax,
            per_channel=True,
            reduce_range=False,
            weight_quantize_mode=0,
            extra_options={"ActivationSymmetric": True},
        )

        self.quantized_models["static_int8"] = save_path
        print(f"Static quantization saved to {save_path}")
        return save_path

    def mixed_precision_quantization(
        self,
        calibration_data: List[np.ndarray],
        sensitive_layers: List[str] = None,
        save_path: str = "model_mixed_int8.onnx",
    ) -> str:
        """Apply mixed precision quantization."""
        if sensitive_layers is None:
            sensitive_layers = ["classifier", "fc"]

        data_reader = CalibrationDataReaderImpl(calibration_data)

        quantize_static(
            model_input=self.onnx_path,
            model_output=save_path,
            calibration_data_reader=data_reader,
            activation_type=QuantType.QUInt8,
            weight_type=QuantType.QUInt8,
            calibration_method=CalibrationMethod.Entropy,
            per_channel=True,
            nodes_to_exclude=sensitive_layers,
        )

        self.quantized_models["mixed"] = save_path
        print(f"Mixed precision quantization saved to {save_path}")
        return save_path

    def benchmark_quantized_models(
        self, test_data: np.ndarray, input_name: str = "input"
    ) -> dict:
        """Benchmark all quantized models."""
        results = {}

        for name, path in [("original", self.onnx_path)] + list(
            self.quantized_models.items()
        ):
            session = ort.InferenceSession(path)
            input_feed = {input_name: test_data}

            import time
            latencies = []
            for _ in range(100):
                start = time.perf_counter()
                session.run(None, input_feed)
                latencies.append(time.perf_counter() - start)

            avg_latency = np.mean(latencies) * 1000
            model_size = onnx.load(path).ByteSize() / 1e6

            output = session.run(None, input_feed)[0]
            top1_acc = np.mean(np.argmax(output, axis=1) == np.argmax(test_data, axis=1)) if test_data.ndim > 1 else 0

            results[name] = {
                "latency_ms": round(avg_latency, 2),
                "size_mb": round(model_size, 2),
                "throughput_fps": round(1000 / avg_latency, 1),
                "accuracy": round(top1_acc * 100, 2),
            }

        return results

# tensorrt_optimize.py
import tensorrt as trt
import numpy as np
from typing import Optional, Tuple
import time

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)


class TensorRTCompiler:
    """TensorRT engine compilation for NVIDIA GPUs."""

    def __init__(
        self,
        fp16: bool = True,
        int8: bool = True,
        max_batch_size: int = 32,
        workspace_size: int = 1 << 30,  # 1GB
    ):
        self.fp16 = fp16
        self.int8 = int8
        self.max_batch_size = max_batch_size
        self.workspace_size = workspace_size

    def build_engine(
        self,
        onnx_path: str,
        engine_path: str = "model.trt",
        input_shape: Optional[Tuple[int, ...]] = None,
    ) -> str:
        """Build TensorRT engine from ONNX model."""
        builder = trt.Builder(TRT_LOGGER)
        network = builder.create_network(
            1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
        )
        parser = trt.OnnxParser(network, TRT_LOGGER)

        with open(onnx_path, "rb") as f:
            if not parser.parse(f.read()):
                for error in range(parser.num_errors):
                    print(parser.get_error(error))
                raise RuntimeError("Failed to parse ONNX model")

        config = builder.create_builder_config()
        config.max_workspace_size = self.workspace_size

        if self.fp16 and builder.platform_has_fast_fp16:
            config.set_flag(trt.BuilderFlag.FP16)
            print("FP16 mode enabled")

        if self.int8 and builder.platform_has_fast_int8:
            config.set_flag(trt.BuilderFlag.INT8)
            print("INT8 mode enabled")

        if input_shape:
            profile = builder.create_optimization_profile()
            profile.set_shape(
                "input",
                min=(1, *input_shape[1:]),
                opt=(self.max_batch_size // 2, *input_shape[1:]),
                max=(self.max_batch_size, *input_shape[1:]),
            )
            config.add_optimization_profile(profile)

        engine = builder.build_serialized_network(network, config)

        with open(engine_path, "wb") as f:
            f.write(engine)

        print(f"TensorRT engine saved to {engine_path}")
        return engine_path


class TensorRTInference:
    """TensorRT inference engine."""

    def __init__(self, engine_path: str):
        self.engine_path = engine_path
        self.engine = None
        self.context = None
        self._load_engine()

    def _load_engine(self):
        """Load TensorRT engine."""
        runtime = trt.Runtime(TRT_LOGGER)
        with open(self.engine_path, "rb") as f:
            self.engine = runtime.deserialize_cuda_engine(f.read())
        self.context = self.engine.create_execution_context()

    def infer(
        self, input_data: np.ndarray, stream=None
    ) -> np.ndarray:
        """Run inference on GPU."""
        import pycuda.driver as cuda
        import pycuda.autoinit

        if stream is None:
            stream = cuda.Stream()

        d_input = cuda.mem_alloc(input_data.nbytes)
        output_shape = (input_data.shape[0], 10)
        output = np.empty(output_shape, dtype=np.float32)
        d_output = cuda.mem_alloc(output.nbytes)

        self.context.set_binding_shape(0, input_data.shape)

        cuda.memcpy_htod_async(d_input, input_data, stream)
        self.context.execute_async_v2(
            bindings=[int(d_input), int(d_output)],
            stream_handle=stream.handle,
        )
        cuda.memcpy_dtoh_async(output, d_output, stream)
        stream.synchronize()

        return output

    def benchmark(
        self, input_data: np.ndarray, num_iterations: int = 100
    ) -> dict:
        """Benchmark inference performance."""
        import pycuda.driver as cuda
        import pycuda.autoinit

        stream = cuda.Stream()
        d_input = cuda.mem_alloc(input_data.nbytes)
        output_shape = (input_data.shape[0], 10)
        output = np.empty(output_shape, dtype=np.float32)
        d_output = cuda.mem_alloc(output.nbytes)

        self.context.set_binding_shape(0, input_data.shape)

        # Warmup
        for _ in range(10):
            cuda.memcpy_htod_async(d_input, input_data, stream)
            self.context.execute_async_v2(
                bindings=[int(d_input), int(d_output)],
                stream_handle=stream.handle,
            )
            stream.synchronize()

        # Benchmark
        start = cuda.Event()
        end = cuda.Event()
        latencies = []

        for _ in range(num_iterations):
            start.record(stream)
            cuda.memcpy_htod_async(d_input, input_data, stream)
            self.context.execute_async_v2(
                bindings=[int(d_input), int(d_output)],
                stream_handle=stream.handle,
            )
            end.record(stream)
            end.synchronize()
            latencies.append(start.time_till(end))

        avg_latency = np.mean(latencies)
        return {
            "avg_latency_ms": round(avg_latency, 2),
            "throughput_fps": round(1000 / avg_latency, 1),
            "p95_latency_ms": round(np.percentile(latencies, 95), 2),
            "p99_latency_ms": round(np.percentile(latencies, 99), 2),
        }

# coreml_convert.py
import coremltools as ct
import torch
import numpy as np
from typing import Tuple


class CoreMLConverter:
    """Convert PyTorch models to CoreML for Apple devices."""

    def __init__(self, model: torch.nn.Module, input_shape: Tuple[int, ...] = (1, 3, 224, 224)):
        self.model = model
        self.input_shape = input_shape

    def convert(
        self,
        save_path: str = "model.mlpackage",
        compute_precision: ct.precision = ct.precision.FLOAT16,
        minimum_deployment_target: ct.target = ct.target.iOS16,
    ) -> str:
        """Convert PyTorch model to CoreML format."""
        self.model.eval()
        example_input = torch.randn(*self.input_shape)

        traced_model = torch.jit.trace(self.model, example_input)

        mlmodel = ct.convert(
            traced_model,
            inputs=[ct.ImageType(
                name="input",
                shape=self.input_shape,
                scale=1/255.0,
                bias=[-0.485/0.229, -0.456/0.224, -0.406/0.225],
            )],
            outputs=[ct.TensorType(name="output")],
            compute_precision=compute_precision,
            minimum_deployment_target=minimum_deployment_target,
        )

        mlmodel.author = "Sustainable AI"
        mlmodel.short_description = "Optimized model for edge deployment"
        mlmodel.save(save_path)

        print(f"CoreML model saved to {save_path}")
        return save_path

    def convert_with_optimizations(
        self,
        save_path: str = "model_optimized.mlpackage",
        target_os: str = "ios16",
    ) -> str:
        """Convert with hardware-specific optimizations."""
        self.model.eval()
        example_input = torch.randn(*self.input_shape)

        traced_model = torch.jit.trace(self.model, example_input)

        deploy_target = {
            "ios16": ct.target.iOS16,
            "macos13": ct.target.macOS13,
        }.get(target_os, ct.target.iOS16)

        mlmodel = ct.convert(
            traced_model,
            inputs=[ct.TensorType(name="input", shape=self.input_shape)],
            outputs=[ct.TensorType(name="output")],
            compute_precision=ct.precision.FLOAT16,
            minimum_deployment_target=deploy_target,
            compute_units=ct.ComputeUnit.ALL,
        )

        mlmodel.save(save_path)
        print(f"Optimized CoreML model saved to {save_path}")
        return save_path