import os
import sys
from transformers import PretrainedConfig, PreTrainedModel, AutoModel, AutoConfig
sys.path.append(os.path.dirname(os.path.dirname(__file__)))
from ultra.models import Ultra
import torch


class UltraConfig(PretrainedConfig):

    model_type = "ultra"

    def __init__(
            self,
            relation_model_layers: int = 6,
            relation_model_dim: int = 64,
            entity_model_layers: int = 6, 
            entity_model_dim: int = 64, 
            **kwargs):
        
        self.relation_model_cfg = dict(
            input_dim=relation_model_dim,
            hidden_dims=[relation_model_dim]*relation_model_layers,
            message_func="distmult",
            aggregate_func="sum",
            short_cut=True,
            layer_norm=True
        )

        self.entity_model_cfg = dict(
            input_dim=entity_model_dim,
            hidden_dims=[entity_model_dim]*entity_model_layers,
            message_func="distmult",
            aggregate_func="sum",
            short_cut=True,
            layer_norm=True
        )

        super().__init__(**kwargs)

class UltraLinkPrediction(PreTrainedModel):

    config_class = UltraConfig

    def __init__(self, config):
        super().__init__(config)

        self.model = Ultra(
            rel_model_cfg=config.relation_model_cfg,
            entity_model_cfg=config.entity_model_cfg,
        )

    def forward(self, data, batch):
        # data: PyG data object 
        # batch shape: (bs, 1+num_negs, 3)
        return self.model.forward(data, batch)
    

if __name__ == "__main__":

    ultra_config = UltraConfig()
    ultra_model = UltraLinkPrediction(ultra_config)
    ultra_model = UltraLinkPrediction.from_pretrained("mgalkin/ultra_3g")
    print(ultra_model)
    print("Done")