import coremltools as ct
import numpy as np
import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModelForCausalLM
import os

# Adjust this path to your downloaded Gemma-3-1b-it model directory
downloaded_hf_model_dir = "/path/your/directory/gemma-3-1b-it" # REPLACE WITH YOUR ACTUAL PATH

print(f"CoreMLTools Version: {ct.__version__}")
print(f"PyTorch Version: {torch.__version__}")
print(f"Transformers Version: {AutoTokenizer.__version__}")
print(f"Numpy Version: {np.__version__}")

try:
    # 1. Hugging Face model loading
    print(f"Loading Hugging Face model from '{downloaded_hf_model_dir}' into memory...")
    model = AutoModelForCausalLM.from_pretrained(downloaded_hf_model_dir, torch_dtype=torch.float16)
    model.eval() # Set model to evaluation mode

    # Disable KV cache for tracing (workaround for potential DynamicCache errors)
    model.config.use_cache = False

    print("Model loaded and configured.")

    # 2. Create a wrapper model for Core ML conversion
    # This addresses the 'dict at output' error during tracing and ensures tuple output
    class GemmaCoreMLWrapper(nn.Module):
        def __init__(self, model):
            super().__init__()
            self.model = model
            # KV cache already handled by self.model.config.use_cache = False above,
            # but reiterating for clarity within wrapper
            self.model.config.use_cache = False

        def forward(self, input_ids, attention_mask):
            # Call the PyTorch model with explicit arguments and force tuple output
            outputs = self.model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                use_cache=False,            # Ensure KV cache is disabled
                return_dict=False,          # Force tuple output (crucial for tracing)
                output_attentions=False,    # Disable attention output for simplicity
                output_hidden_states=False  # Disable hidden states output for simplicity
            )
            # Return only the logits (first element of the tuple output)
            logits = outputs[0]
            return logits

    wrapped_model = GemmaCoreMLWrapper(model) # Use the loaded 'model' here
    print("Wrapper model prepared.")

    # 3. Prepare dummy inputs for TorchScript tracing
    max_seq_length = 1024 # Sequence length for dynamic range upper bound
    tokenizer = AutoTokenizer.from_pretrained(downloaded_hf_model_dir) # Load tokenizer for vocab_size

    # Dummy inputs matching the expected type and a small sequence length for tracing
    dummy_input_ids = torch.randint(0, tokenizer.vocab_size, (1, 10), dtype=torch.long)
    dummy_attention_mask = torch.ones(1, 10, dtype=torch.long)

    # 4. Trace the wrapper model to TorchScript
    print("Tracing wrapper model to TorchScript...")
    traced_model = torch.jit.trace(wrapped_model, (dummy_input_ids, dummy_attention_mask))
    print("TorchScript tracing complete.")

    # 5. Convert TorchScript model to Core ML
    print("Converting TorchScript model to Core ML...")
    coreml_model = ct.convert(
        traced_model, # Pass the TorchScript object
        inputs=[
            # Use np.int32 based on previous warning about int64 downcasting
            ct.TensorType(name="input_ids", shape=(1, ct.RangeDim(upper_bound=max_seq_length)), dtype=np.int32),
            ct.TensorType(name="attention_mask", shape=(1, ct.RangeDim(upper_bound=max_seq_length)), dtype=np.int32)
        ],
        source="pytorch",
        convert_to="mlprogram",
        # Specify a modern iOS deployment target (Gemma is a recent model)
        minimum_deployment_target=ct.target.iOS16 # Or iOS17/iOS18 as appropriate
    )

    # Save the Core ML model
    coreml_model.save("gemma3-1b-it-coreml.mlpackage")
    print("CoreML model saved successfully to gemma3-1b-it-coreml.mlpackage.")

except Exception as e:
    print(f"Conversion to CoreML failed: {e}")