# modeling_ohca.py
import os
import numpy as np
from transformers import PreTrainedModel, AutoConfig, PretrainedConfig

class OhcaConfig(PretrainedConfig):
    model_type = "cvd_ohca_predictor"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

try:
    import tensorflow
except ImportError:
    raise ImportError(
        "\n\n"
        "[MISSING DEPENDENCY ERROR]\n"
        "TensorFlow v2.16+ is required to run models with the OHCAPredictorPipeline.\n\n"
        "How to fix this:\n"
        "Install via the instructions at https://www.tensorflow.org/install\n"
        "Verify: \"import tensorflow as tf; print(tf.__version__)\" \n\n"
        "Copyright (c) 2026 Rihaan Meher. Licensed Under the Terms of the Apache License, v2.0\n"
        "---------------------------------------------------------------------------"
    )
try:
    from ohca_predictor.pipeline import OHCAPredictorPipeline
    from ohca_predictor.utils.io_utils import WindowSample, create_padded_dataset
except ImportError:
    raise ImportError(
        "\n\n"
        "[MISSING DEPENDENCY ERROR]\n"
        "The 'cvdpredict' package version 1.0.0 is required to run models with the OHCAPredictorPipeline.\n\n"
        "How to fix this:\n"
        "Build from source:        pip install git+https://github.com/sharktide/CVD-Predict.git@v1.0.0\n"
        "Verify: python3 -c \"import ohca_predictor; print(ohca_predictor.__version__)\"\n"
        "Copyright (c) 2026 Rihaan Meher. Licensed Under the Terms of the Apache License, v2.0\n"
        "---------------------------------------------------------------------------"
    )

class TFCVDPredictorForOHCADetection(PreTrainedModel):
    config_class = OhcaConfig

    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        self.pipeline_instance = None

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        config = kwargs.pop("config", None) or OhcaConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
        model = cls(config)
        
        weights_path = "hf://sharktide/ohca-predictor-v1"
        model.pipeline_instance = OHCAPredictorPipeline.load(weights_path)
        return model

    def forward(self, inputs, **kwargs):
        """
        Accepts a preprocessed batch dictionary OR a single/list of WindowSample objects.
        """
        # If the input is a raw WindowSample instance, run it through the package preprocessor
        if isinstance(inputs, WindowSample) or (isinstance(inputs, list) and len(inputs) > 0 and isinstance(inputs[0], WindowSample)):
            samples_list = [inputs] if isinstance(inputs, WindowSample) else inputs
            dataset = create_padded_dataset(samples_list, batch_size=len(samples_list), shuffle=False)
            # Pull the preprocessed tensor dictionary from the single-batch dataset
            for batch in dataset:
                inputs = batch
                break

        # Forward directly into the verified Keras transformer core
        outputs = self.pipeline_instance.model(inputs, training=False)
        
        return {
            "ohca_risk": outputs["ohca_risk"],
            "survival_outputs": outputs.get("survival_outputs", None)
        }
        
    def __call__(self, *args, **kwargs):
        return self.forward(*args, **kwargs)