#!/usr/bin/env python3
"""Export locally fine-tuned Hugging Face models to ONNX format for production inference"""

import os
from pathlib import Path

import torch
from optimum.onnxruntime import ORTModelForSequenceClassification
from transformers import AutoTokenizer, AutoConfig

# Local fine-tuned model paths
MODELS = {
    "finbert": {
        "local_path": "/mnt/dolphinng5_predict/sentiment_engine/models/finbert-crypto-sentiment",
        "output_dir": "/mnt/dolphinng5_predict/sentiment_engine/models/onnx/finbert",
        "labels": ["Bearish", "Bullish", "Neutral"],
        "id2label": {0: "Bearish", 1: "Bullish", 2: "Neutral"},
    },
    "bert-base-event": {
        "local_path": "/mnt/dolphinng5_predict/sentiment_engine/models/bert-crypto-events",
        "output_dir": "/mnt/dolphinng5_predict/sentiment_engine/models/onnx/bert-base-event",
        "labels": ["listing", "delisting", "hack", "regulatory", "governance", 
                  "upgrade", "partnership", "earnings", "macro", "liquidation", "whale", "manipulation"],
        "id2label": {i: l for i, l in enumerate([
            "listing", "delisting", "hack", "regulatory", "governance", 
            "upgrade", "partnership", "earnings", "macro", "liquidation", "whale", "manipulation"
        ])},
    },
    "distilroberta-emotion": {
        "local_path": "/mnt/dolphinng5_predict/sentiment_engine/models/distilroberta-crypto-emotion",
        "output_dir": "/mnt/dolphinng5_predict/sentiment_engine/models/onnx/distilroberta-emotion",
        "labels": ["joy", "fear", "anger", "greed", "sadness", "neutral"],
        "id2label": {i: l for i, l in enumerate(["joy", "fear", "anger", "greed", "sadness", "neutral"])},
    },
    "minilm-l6-v2": {
        "local_path": "/mnt/dolphinng5_predict/sentiment_engine/models/finbert-crypto-sentiment",  # Use finbert tokenizer
        "output_dir": "/mnt/dolphinng5_predict/sentiment_engine/models/onnx/minilm-l6-v2",
        "labels": None,
        "id2label": None,
    },
}

def export_classification_model(model_key: str) -> None:
    """Export a local classification model to ONNX"""
    config = MODELS[model_key]
    local_path = config["local_path"]
    output_dir = Path(config["output_dir"])
    output_dir.mkdir(parents=True, exist_ok=True)

    print(f"Exporting {model_key} from {local_path} to {output_dir}...")

    # Load model config to check problem type
    model_config = AutoConfig.from_pretrained(local_path)
    is_multilabel = getattr(model_config, "problem_type", None) == "multi_label_classification"
    
    print(f"  Problem type: {getattr(model_config, 'problem_type', 'single_label')}")
    print(f"  Labels: {config['labels']}")

    # Load model and export using optimum
    model = ORTModelForSequenceClassification.from_pretrained(
        local_path,
        export=True,
    )
    model.save_pretrained(output_dir)

    # Save tokenizer
    tokenizer = AutoTokenizer.from_pretrained(local_path)
    tokenizer.save_pretrained(output_dir)

    # Save label mapping
    import json
    if config["labels"]:
        with open(output_dir / "label_map.json", "w") as f:
            json.dump({i: label for i, label in enumerate(config["labels"])}, f)
        with open(output_dir / "id2label.json", "w") as f:
            json.dump(config["id2label"], f)

    print(f"  Done: {model_key}")

def export_feature_extraction_model(model_key: str) -> None:
    """Export a feature extraction model to ONNX"""
    config = MODELS[model_key]
    local_path = config["local_path"]
    output_dir = Path(config["output_dir"])
    output_dir.mkdir(parents=True, exist_ok=True)

    print(f"Exporting {model_key} (feature extraction) from {local_path} to {output_dir}...")

    from transformers import AutoModel
    
    # For sentence transformers / feature extraction
    hf_model = AutoModel.from_pretrained(local_path)
    hf_model.eval()
    
    # Create dummy input
    dummy_input = {
        "input_ids": torch.ones(1, 128, dtype=torch.long),
        "attention_mask": torch.ones(1, 128, dtype=torch.long),
    }
    
    # Export to ONNX
    torch.onnx.export(
        hf_model,
        (dummy_input["input_ids"], dummy_input["attention_mask"]),
        output_dir / "model.onnx",
        input_names=["input_ids", "attention_mask"],
        output_names=["last_hidden_state", "pooler_output"],
        dynamic_axes={
            "input_ids": {0: "batch", 1: "sequence"},
            "attention_mask": {0: "batch", 1: "sequence"},
            "last_hidden_state": {0: "batch", 1: "sequence"},
        },
        opset_version=14,
    )
    print(f"  Exported feature extraction model")
    
    # Save tokenizer
    tokenizer = AutoTokenizer.from_pretrained(local_path)
    tokenizer.save_pretrained(output_dir)

    print(f"  Done: {model_key}")

def main():
    print("="*60)
    print("EXPORTING FINE-TUNED MODELS TO ONNX")
    print("="*60)
    
    # Export classification models
    for model_key in ["finbert", "bert-base-event", "distilroberta-emotion"]:
        try:
            export_classification_model(model_key)
        except Exception as e:
            print(f"  ERROR exporting {model_key}: {e}")
            import traceback
            traceback.print_exc()
    
    # Export feature extraction model (MiniLM)
    try:
        export_feature_extraction_model("minilm-l6-v2")
    except Exception as e:
        print(f"  ERROR exporting minilm-l6-v2: {e}")
        import traceback
        traceback.print_exc()
    
    print("\n" + "="*60)
    print("ALL EXPORTS COMPLETE!")
    print("="*60)

if __name__ == "__main__":
    main()