import torch
import torch.onnx
import onnx
import tensorflow as tf
from onnx_tf.backend import prepare
import numpy as np
import os

# Script producción-ready para conversión modelo PyTorch → TFLite con quantization
# Autor: BCloud Consulting | Fecha: 2025


class ModelConverter:
    """
    Convierte modelo PyTorch a TensorFlow Lite optimizado para edge deployment.
    Soporta quantization INT8 para reducir tamaño 75% y mejorar speed 85%.
    """

    def __init__(self, pytorch_model, input_shape, output_path="model_optimized.tflite"):
        """
        Args:
            pytorch_model: Modelo PyTorch en eval mode
            input_shape: Tuple con shape input (ej: (1, 3, 224, 224) para imágenes)
            output_path: Path donde guardar .tflite file
        """
        self.model = pytorch_model
        self.input_shape = input_shape
        self.output_path = output_path

    def export_to_onnx(self, onnx_path="model_temp.onnx"):
        """
        Paso 1: Exportar PyTorch → ONNX (formato intermediate hardware-agnostic)
        """
        print(f"[1/4] Exportando PyTorch → ONNX...")

        # Crear dummy input para tracing
        dummy_input = torch.randn(self.input_shape)

        # Export con dynamic axes para permitir batch size variable
        torch.onnx.export(
            self.model,
            dummy_input,
            onnx_path,
            export_params=True,
            opset_version=12,  # Opset 12 compatible con la mayoría de runtimes
            do_constant_folding=True,  # Optimización: fold constants en compile time
            input_names=['input'],
            output_names=['output'],
            dynamic_axes={
                'input': {0: 'batch_size'},  # Primer dimensión es batch
                'output': {0: 'batch_size'}
            }
        )

        # Verificar modelo ONNX es válido
        onnx_model = onnx.load(onnx_path)
        onnx.checker.check_model(onnx_model)

        print(f"✓ ONNX model guardado en {onnx_path}")
        return onnx_path

    def onnx_to_tensorflow(self, onnx_path, tf_path="model_tf"):
        """
        Paso 2: Convertir ONNX → TensorFlow SavedModel
        """
        print(f"[2/4] Convirtiendo ONNX → TensorFlow...")

        # Cargar modelo ONNX
        onnx_model = onnx.load(onnx_path)

        # Preparar TensorFlow representation
        tf_rep = prepare(onnx_model)

        # Exportar como TensorFlow SavedModel
        tf_rep.export_graph(tf_path)

        print(f"✓ TensorFlow model guardado en {tf_path}")
        return tf_path

    def quantize_and_convert_tflite(self, tf_path, calibration_dataset=None):
        """
        Paso 3: Convertir TF SavedModel → TFLite con quantization INT8

        Args:
            calibration_dataset: Generator yielding calibration inputs para quantization
                                Debería ser representativo de datos producción
        """
        print(f"[3/4] Convirtiendo TensorFlow → TFLite con quantization...")

        # Cargar SavedModel
        converter = tf.lite.TFLiteConverter.from_saved_model(tf_path)

        # Configurar optimizations
        converter.optimizations = [tf.lite.Optimize.DEFAULT]

        if calibration_dataset is not None:
            # Post-training quantization completa (INT8) usando calibration data
            converter.representative_dataset = calibration_dataset
            converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
            converter.inference_input_type = tf.uint8  # Input será uint8
            converter.inference_output_type = tf.uint8  # Output será uint8
            print(" → Quantization INT8 activada con calibration dataset")
        else:
            # Dynamic range quantization (fallback sin calibration data)
            print(" → Dynamic range quantization (sin calibration)")

        # Convertir
        tflite_model = converter.convert()

        # Guardar modelo cuantizado
        with open(self.output_path, 'wb') as f:
            f.write(tflite_model)

        # Reportar tamaño
        original_size = os.path.getsize(f"{tf_path}/saved_model.pb") / (1024**2)  # MB
        optimized_size = os.path.getsize(self.output_path) / (1024**2)  # MB
        reduction = ((original_size - optimized_size) / original_size) * 100

        print(f"✓ TFLite model guardado en {self.output_path}")
        print(f"  Tamaño original: {original_size:.2f} MB")
        print(f"  Tamaño optimizado: {optimized_size:.2f} MB")
        print(f"  Reducción: {reduction:.1f}%")

    def validate_conversion(self, test_input):
        """
        Paso 4: Validar numerical equivalence PyTorch vs TFLite
        """
        print(f"[4/4] Validando conversión...")

        # Inferencia PyTorch original
        with torch.no_grad():
            pytorch_output = self.model(test_input).numpy()

        # Inferencia TFLite
        interpreter = tf.lite.Interpreter(model_path=self.output_path)
        interpreter.allocate_tensors()

        input_details = interpreter.get_input_details()
        output_details = interpreter.get_output_details()

        # Preparar input (convertir a uint8 si quantized)
        test_input_np = test_input.numpy()

        if input_details[0]['dtype'] == np.uint8:
            # Quantize input
            scale, zero_point = input_details[0]['quantization']
            test_input_np = (test_input_np / scale + zero_point).astype(np.uint8)

        interpreter.set_tensor(input_details[0]['index'], test_input_np)
        interpreter.invoke()
        tflite_output = interpreter.get_tensor(output_details[0]['index'])

        # Dequantize output si necesario
        if output_details[0]['dtype'] == np.uint8:
            scale, zero_point = output_details[0]['quantization']
            tflite_output = (tflite_output.astype(np.float32) - zero_point) * scale

        # Calcular similarity (cosine similarity)
        from scipy.spatial.distance import cosine
        similarity = 1 - cosine(pytorch_output.flatten(), tflite_output.flatten())

        print(f"  Cosine similarity PyTorch ↔ TFLite: {similarity:.4f}")

        if similarity > 0.99:
            print("  ✓ Conversión exitosa (similarity >0.99)")
        elif similarity > 0.95:
            print("  ⚠ Similarity acceptable pero verificar accuracy (0.95-0.99)")
        else:
            print("  ✗ WARNING: Low similarity (

from flask import Flask, request, jsonify
import onnxruntime as ort
import numpy as np
import time
from prometheus_client import Counter, Histogram, generate_latest, REGISTRY
import logging

# Servidor inferencia edge production-ready con ONNX Runtime
# Características: Batching, health checks, Prometheus metrics, error handling
# Autor: BCloud Consulting | 2025

app = Flask(__name__)
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Prometheus metrics
REQUEST_COUNT = Counter('inference_requests_total', 'Total inference requests')
REQUEST_LATENCY = Histogram('inference_latency_seconds', 'Inference latency in seconds')
ERROR_COUNT = Counter('inference_errors_total', 'Total inference errors')


class ONNXInferenceEngine:
    """
    Motor inferencia ONNX optimizado para edge deployment.
    Soporta CPU y GPU (CUDA) execution providers.
    """

    def __init__(self, model_path, use_gpu=False, max_batch_size=8):
        """
        Args:
            model_path: Path al archivo .onnx
            use_gpu: Si True, usar CUDA execution provider (requiere GPU NVIDIA)
            max_batch_size: Máximo batch size para inferencia
        """
        self.model_path = model_path
        self.max_batch_size = max_batch_size

        # Configurar execution providers (GPU o CPU)
        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if use_gpu else ['CPUExecutionProvider']

        # Cargar modelo ONNX
        try:
            self.session = ort.InferenceSession(model_path, providers=providers)
            logger.info(f"✓ Modelo ONNX cargado desde {model_path}")
            logger.info(f"  Execution providers: {self.session.get_providers()}")
        except Exception as e:
            logger.error(f"✗ Error cargando modelo: {e}")
            raise

        # Obtener input/output metadata
        self.input_name = self.session.get_inputs()[0].name
        self.output_name = self.session.get_outputs()[0].name
        self.input_shape = self.session.get_inputs()[0].shape

        logger.info(f"  Input: {self.input_name}, shape: {self.input_shape}")
        logger.info(f"  Output: {self.output_name}")

    def preprocess(self, image_array):
        """
        Preprocessing input image (normalización, resize, etc.)
        Ajustar según tu modelo específico.
        """
        # Ejemplo: Normalize to [0, 1] y convert a formato esperado
        if image_array.max() > 1.0:
            image_array = image_array / 255.0

        # Asegurar shape correcto (batch, channels, height, width)
        if len(image_array.shape) == 3:
            image_array = np.expand_dims(image_array, axis=0)  # Add batch dimension

        return image_array.astype(np.float32)

    def infer(self, input_data):
        """
        Ejecutar inferencia con timing.

        Args:
            input_data: Numpy array con shape compatible con modelo

        Returns:
            (output, latency_ms): Tuple con predicción y latencia en milisegundos
        """
        start_time = time.time()

        try:
            # Ejecutar inferencia
            output = self.session.run(
                [self.output_name],
                {self.input_name: input_data}
            )[0]

            latency_ms = (time.time() - start_time) * 1000

            return output, latency_ms

        except Exception as e:
            logger.error(f"Inference error: {e}")
            raise

    def batch_infer(self, inputs_list):
        """
        Inferencia batch para múltiples inputs.
        Agrupa requests para mejor throughput.

        Args:
            inputs_list: Lista de numpy arrays

        Returns:
            List de predictions
        """
        # Agregar inputs en batches
        batches = [
            inputs_list[i:i+self.max_batch_size]
            for i in range(0, len(inputs_list), self.max_batch_size)
        ]

        results = []
        total_latency = 0

        for batch in batches:
            # Stack batch
            batch_array = np.vstack(batch)

            # Infer
            output, latency = self.infer(batch_array)
            results.extend(output)
            total_latency += latency

        avg_latency = total_latency / len(batches)
        logger.info(f"Batch inference: {len(inputs_list)} samples, avg latency: {avg_latency:.2f}ms")

        return results


# Inicializar motor inferencia (global)
MODEL_PATH = "/models/optimized-model.onnx"
USE_GPU = True  # Cambiar a False si no hay GPU disponible
inference_engine = ONNXInferenceEngine(MODEL_PATH, use_gpu=USE_GPU)


@app.route('/healthz', methods=['GET'])
def health_check():
    """
    Health check endpoint para Kubernetes liveness probe.
    """
    return jsonify({"status": "healthy", "model": MODEL_PATH}), 200


@app.route('/ready', methods=['GET'])
def readiness_check():
    """
    Readiness check para Kubernetes readiness probe.
    Verifica que modelo está cargado y listo.
    """
    try:
        # Test inference con dummy input
        dummy_input = np.random.rand(1, 3, 224, 224).astype(np.float32)
        _, latency = inference_engine.infer(dummy_input)

        return jsonify({
            "status": "ready",
            "test_latency_ms": latency
        }), 200

    except Exception as e:
        return jsonify({"status": "not ready", "error": str(e)}), 503


@app.route('/predict', methods=['POST'])
def predict():
    """
    Endpoint inferencia principal.
    Acepta JSON con array de imágenes (base64 o numpy array serializado).
    """
    REQUEST_COUNT.inc()

    try:
        # Parse request
        data = request.get_json()

        if 'image' not in data:
            ERROR_COUNT.inc()
            return jsonify({"error": "Missing 'image' field"}), 400

        # Preprocess
        image_array = np.array(data['image'])
        input_data = inference_engine.preprocess(image_array)

        # Infer con timing
        with REQUEST_LATENCY.time():
            output, latency_ms = inference_engine.infer(input_data)

        # Post-process (ej: aplicar softmax, obtener top-k predictions)
        predictions = output.tolist()

        logger.info(f"Inference successful: latency={latency_ms:.2f}ms")

        return jsonify({
            "predictions": predictions,
            "latency_ms": latency_ms,
            "model_version": "v1.2.0"
        }), 200

    except Exception as e:
        ERROR_COUNT.inc()
        logger.error(f"Prediction error: {e}")
        return jsonify({"error": str(e)}), 500


@app.route('/metrics', methods=['GET'])
def metrics():
    """
    Prometheus metrics endpoint.
    Scraped por Prometheus para monitoring.
    """
    return generate_latest(REGISTRY), 200, {'Content-Type': 'text/plain; charset=utf-8'}


if __name__ == '__main__':
    # Producción: usar Gunicorn/uWSGI en lugar de Flask dev server
    # Comando: gunicorn -w 4 -b 0.0.0.0:8080 onnx_inference_server:app
    app.run(host='0.0.0.0', port=8080, debug=False)

# Prometheus configuration para monitoring edge AI fleet
# Scrape edge devices inference metrics + system metrics

global:
scrape_interval: 15s  # Scrape cada 15 segundos
evaluation_interval: 15s  # Evaluar rules cada 15 segundos
external_labels:
    cluster: 'edge-ai-production'
    environment: 'production'

# Alertmanager configuration (para notificaciones)
alerting:
alertmanagers:
    - static_configs:
        - targets: ['alertmanager:9093']

# Scrape configs: Qué targets monitorear
scrape_configs:
# Job 1: Edge inference servers (nuestro servidor ONNX)
- job_name: 'edge-inference-servers'
    static_configs:
    - targets:
        - 'edge-device-01:8080'  # Factory site 1
        - 'edge-device-02:8080'  # Factory site 2
        - 'edge-device-03:8080'  # Factory site 3
        # ... más devices
    metrics_path: '/metrics'  # Endpoint métricas
    # Labels adicionales por device
    relabel_configs:
    - source_labels: [__address__]
        regex: 'edge-device-(\d+):.*'
        target_label: 'device_id'
        replacement: '$1'

# Job 2: Node exporters (system metrics: CPU, RAM, disk)
- job_name: 'node-exporter'
    static_configs:
    - targets:
        - 'edge-device-01:9100'
        - 'edge-device-02:9100'
        - 'edge-device-03:9100'

# Job 3: NVIDIA GPU metrics (si devices tienen GPUs)
- job_name: 'nvidia-gpu'
    static_configs:
    - targets:
        - 'edge-device-01:9445'  # DCGM exporter port
        - 'edge-device-02:9445'

# Job 4: Kubernetes metrics (si usas K3s/KubeEdge)
- job_name: 'kubernetes-pods'
    kubernetes_sd_configs:
    - role: pod
        namespaces:
        names: ['production']
    relabel_configs:
    # Solo scrape pods con annotation prometheus.io/scrape: "true"
    - source_labels: [__meta_kubernetes_pod_annotation_prometheus_io_scrape]
        action: keep
        regex: true
    - source_labels: [__meta_kubernetes_pod_annotation_prometheus_io_path]
        action: replace
        target_label: __metrics_path__
        regex: (.+)
    - source_labels: [__address__, __meta_kubernetes_pod_annotation_prometheus_io_port]
        action: replace
        regex: ([^:]+)(?::\d+)?;(\d+)
        replacement: $1:$2
        target_label: __address__

# Alerting rules
rule_files:
- '/etc/prometheus/alerts/*.yml'

# Alerting rules para edge AI monitoring

groups:
- name: edge_ai_alerts
    interval: 30s
    rules:
    # Alert: Latencia alta
    - alert: HighInferenceLatency
        expr: histogram_quantile(0.95, rate(inference_latency_seconds_bucket[5m])) > 0.100
        for: 5m
        labels:
        severity: warning
        annotations:
        summary: "High inference latency on {{ $labels.device_id }}"
        description: "P95 latency is {{ $value }}s (threshold 100ms) for last 5min"

    # Alert: Error rate alta
    - alert: HighErrorRate
        expr: rate(inference_errors_total[5m]) > 0.05
        for: 2m
        labels:
        severity: critical
        annotations:
        summary: "High error rate on {{ $labels.device_id }}"
        description: "Error rate {{ $value }} errors/sec (threshold 0.05/sec)"

    # Alert: Device offline
    - alert: EdgeDeviceDown
        expr: up{job="edge-inference-servers"} == 0
        for: 1m
        labels:
        severity: critical
        annotations:
        summary: "Edge device {{ $labels.device_id }} is down"
        description: "Device unreachable for 1+ minutes"

    # Alert: GPU utilization muy baja (underutilized)
    - alert: GPUUnderutilized
        expr: avg_over_time(DCGM_FI_DEV_GPU_UTIL[1h]) < 30
        for: 30m
        labels:
        severity: info
        annotations:
        summary: "GPU underutilized on {{ $labels.device_id }}"
        description: "Average GPU util {{ $value }}% last hour (consider downsizing)"

    # Alert: Disk space bajo
    - alert: LowDiskSpace
        expr: (node_filesystem_avail_bytes / node_filesystem_size_bytes) * 100 < 10
        for: 5m
        labels:
        severity: warning
        annotations:
        summary: "Low disk space on {{ $labels.device_id }}"
        description: "Only {{ $value }}% disk space remaining" 

apiVersion: apps/v1
kind: Deployment
metadata:
name: edge-ai-inference
namespace: production
labels:
    app: edge-ai
    version: v1.2.0
spec:
replicas: 3  # High availability edge cluster
selector:
    matchLabels:
    app: edge-ai
template:
    metadata:
    labels:
        app: edge-ai
        version: v1.2.0
    spec:
    # Tolerations para edge nodes (pueden tener taints específicos)
    tolerations:
        - key: "edge-node"
        operator: "Exists"
        effect: "NoSchedule"

    # Node selector para asegurar deployment solo en edge nodes
    nodeSelector:
        node-type: edge
        gpu: "true"

    containers:
        - name: onnx-inference-server
        image: bcloud/edge-ai-inference:v1.2.0-onnx
        imagePullPolicy: IfNotPresent

        # Resources requests/limits específicos para edge devices
        resources:
            requests:
            memory: "2Gi"
            cpu: "1000m"
            nvidia.com/gpu: "1"  # Requiere NVIDIA GPU
            limits:
            memory: "4Gi"
            cpu: "2000m"
            nvidia.com/gpu: "1"

        # Health checks críticos para edge reliability
        livenessProbe:
            httpGet:
            path: /healthz
            port: 8080
            initialDelaySeconds: 30
            periodSeconds: 10
            timeoutSeconds: 5
            failureThreshold: 3

        readinessProbe:
            httpGet:
            path: /ready
            port: 8080
            initialDelaySeconds: 10
            periodSeconds: 5

        # Environment variables para configuración
        env:
            - name: MODEL_PATH
            value: "/models/optimized-model.onnx"
            - name: BATCH_SIZE
            value: "8"
            - name: MAX_LATENCY_MS
            value: "50"  # SLA latency edge AI
            - name: LOG_LEVEL
            value: "INFO"

        # Volume mounts para modelo + logs
        volumeMounts:
            - name: model-storage
            mountPath: /models
            readOnly: true
            - name: logs
            mountPath: /var/log/inference

        ports:
            - containerPort: 8080
            name: http
            protocol: TCP
            - containerPort: 9090
            name: metrics
            protocol: TCP

    volumes:
        # Persistent volume para modelo (actualizado vía GitOps)
        - name: model-storage
        persistentVolumeClaim:
            claimName: edge-ai-model-pvc
        # EmptyDir para logs (sincronizados a cloud async)
        - name: logs
        emptyDir: {}
---
# Service para exponer inference endpoint
apiVersion: v1
kind: Service
metadata:
name: edge-ai-service
namespace: production
labels:
    app: edge-ai
spec:
type: ClusterIP
selector:
    app: edge-ai
ports:
    - name: http
    port: 80
    targetPort: 8080
    protocol: TCP
    - name: metrics
    port: 9090
    targetPort: 9090
    protocol: TCP
---
# HorizontalPodAutoscaler para scaling automático basado en CPU/memory
apiVersion: autoscaling/v2
kind: HorizontalPodAutoscaler
metadata:
name: edge-ai-hpa
namespace: production
spec:
scaleTargetRef:
    apiVersion: apps/v1
    kind: Deployment
    name: edge-ai-inference
minReplicas: 2
maxReplicas: 10
metrics:
    - type: Resource
    resource:
        name: cpu
        target:
        type: Utilization
        averageUtilization: 70
    - type: Resource
    resource:
        name: memory
        target:
        type: Utilization
        averageUtilization: 80
behavior:
    scaleDown:
    stabilizationWindowSeconds: 300  # Evitar flapping
    policies:
        - type: Percent
        value: 50
        periodSeconds: 60
    scaleUp:
    stabilizationWindowSeconds: 0
    policies:
        - type: Percent
        value: 100
        periodSeconds: 30