import torch
import os
import json
from modeling_omegacode import HolographicMasterCodeTransformer

class InferenceHandler:
    def __init__(self):
        self.model = None
        self.initialized = False

    def initialize(self, context):
        # Modelin yüklendiği depo ID'si (Hugging Face tarafından sağlanır)
        self.model_repo_id = os.getenv("HF_MODEL_ID") # Example: "user/repo_name"
        if not self.model_repo_id:
            raise ValueError("HF_MODEL_ID environment variable not set. Cannot determine model repository.")

        # Modelin yerel olarak nerede bulunduğunu belirtir
        # Genellikle Hugging Face Inference API, dosyaları `/tmp/model` altına indirir.
        model_dir = context.properties.get("model_dir", "/tmp/model")

        # `modeling_omegacode.py` dosyasının Python yoluna eklenmesi
        # Bu adım, `modeling_omegacode` modülünün bulunmasını sağlar.
        import sys
        if model_dir not in sys.path:
            sys.path.insert(0, model_dir)

        # config.json dosyasını oku
        config_path = os.path.join(model_dir, "config.json")
        if not os.path.exists(config_path):
            raise FileNotFoundError(f"config.json not found at {config_path}")
        with open(config_path, 'r') as f:
            model_config = json.load(f)

        # Modeli yapılandırmadan oluştur
        self.model = HolographicMasterCodeTransformer(
            input_dim=model_config.get('input_dim', 8),
            d_model=model_config.get('d_model', 128),
            nhead=model_config.get('nhead', 8),
            num_layers=model_config.get('num_layers', 6),
            output_dim=model_config.get('output_dim', 1),
            n_harmonics=model_config.get('n_harmonics', 4)
        )

        # Eğitilmiş ağırlıkları yükle
        model_weights_path = os.path.join(model_dir, "pytorch_model.bin")
        if not os.path.exists(model_weights_path):
            raise FileNotFoundError(f"pytorch_model.bin not found at {model_weights_path}")
        self.model.load_state_dict(torch.load(model_weights_path, map_location='cpu'))
        self.model.eval() # Modeli değerlendirme moduna al

        self.initialized = True

    def preprocess(self, inputs):
        # Gelen girişi PyTorch tensoruna dönüştür
        # Hugging Face Inference API'si genellikle JSON veya benzeri bir format alır.
        # Örneğin, inputs = {'input_data': [[...]]}
        try:
            input_data = inputs[0].get('input_data')
            if not input_data:
                raise ValueError("Missing 'input_data' in input payload")
            # Listeyi tensor'a çevir
            tensor_input = torch.tensor(input_data, dtype=torch.float32)
            return tensor_input
        except Exception as e:
            raise ValueError(f"Input preprocessing failed: {e}. Expected format: [{{'input_data': [[...]]}}]")

    def inference(self, input_batch):
        # Modeli çalıştırma
        if not self.initialized:
            raise RuntimeError("Model not initialized.")
        
        # alpha değerini config'den al veya varsayılanı kullan
        # model_config'i initialize içinde self'e atayarak erişilebilir hale getiriyoruz.
        # Ama burada sadece nominal değeri kullanmak daha güvenli olabilir.
        alpha_value = 1/137.035 # Nominal alpha değeri

        with torch.no_grad():
            output = self.model(input_batch, alpha=alpha_value)
        return output

    def postprocess(self, outputs):
        # Çıktı tensorunu uygun bir formata dönüştür
        # Örneğin, PyTorch tensorunu bir Python listesine veya sözlüğe çevir.
        return [{'output': outputs.tolist()}]

    def handle(self, inputs, context):
        # Tüm süreci yöneten ana metod
        model_input = self.preprocess(inputs)
        model_output = self.inference(model_input)
        return self.postprocess(model_output)

# Bu kısım genellikle Hugging Face inference container tarafından kullanılır.
# Yerel test için bu kısmı yorumdan kaldırabilirsiniz.
# if __name__ == '__main__':
#     handler = InferenceHandler()
#     # Dummy context and inputs for local testing
#     class DummyContext:
#         def __init__(self):
#             self.properties = {"model_dir": "."}
#     dummy_context = DummyContext()
#     handler.initialize(dummy_context)
#     
#     dummy_inputs = [{'input_data': [[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]]}]
#     output = handler.handle(dummy_inputs, dummy_context)
#     print(output)