# coding: utf-8
import torch
import logging
from .architectures import (
    EmotionTransformer,
    PersonalityTransformer,
    FusionTransformer,
)


def load_pretrained_emotion_encoder(checkpoint_path, device):
    emotion_model = EmotionTransformer(
        input_dim_emotion=2560,
        input_dim_personality=2560,
        hidden_dim=512,
        out_features=1024,
        tr_layer_number=2,
        num_transformer_heads=16,
        positional_encoding=False,
        dropout=0.1,
        per_activation = "relu"
    ).to(device)

    checkpoint = torch.load(checkpoint_path, map_location=device)
    state_dict = checkpoint["model_state_dict"] if "model_state_dict" in checkpoint else checkpoint
    emotion_model.load_state_dict(state_dict)

    def extract_features(inputs, lengths):
        features = emotion_model.emo_proj(inputs)
        for block in emotion_model.emotion_encoder:
            features = block(features)
        return features

    emotion_model.extract_features = extract_features
    emotion_model.eval()
    return emotion_model

def load_pretrained_personality_encoder(checkpoint_path, device):
    personality_model = PersonalityTransformer(
        input_dim_emotion=2560,
        input_dim_personality=2560,
        hidden_dim=128,
        out_features=1024,
        tr_layer_number=2,
        num_transformer_heads=4,
        positional_encoding=False,
        dropout=0.1,
        per_activation = "relu"
    ).to(device)

    checkpoint = torch.load(checkpoint_path, map_location=device)
    personality_model.load_state_dict(checkpoint)

    def extract_features(inputs, lengths):
        features = personality_model.per_proj(inputs)
        for block in personality_model.personality_encoder:
            features = block(features, features, features)
        return features

    personality_model.extract_features = extract_features
    personality_model.eval()
    return personality_model

def load_fusion_model(
    fusion_checkpoint_path: str,
    emotion_encoder_checkpoint: str,
    personality_encoder_checkpoint: str,
    device: str = "cpu",
):
    device = torch.device(device)

    emotion_encoder = load_pretrained_emotion_encoder(emotion_encoder_checkpoint, device)
    personality_encoder = load_pretrained_personality_encoder(personality_encoder_checkpoint, device)

    checkpoint = torch.load(fusion_checkpoint_path, map_location=device)

    fusion_model = FusionTransformer(
        emo_model=emotion_encoder,
        per_model=personality_encoder,
        hidden_dim=1024,
        out_features=256,
        tr_layer_number=2,
        num_transformer_heads=4,
        positional_encoding=False,
        per_activation="relu",
        dropout=0.1
    ).to(device)
    fusion_model.load_state_dict(checkpoint)
    fusion_model.eval()
    return fusion_model, device