import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModel
import numpy as np
from sklearn.preprocessing import StandardScaler
from joblib import load
import random

def set_seed(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
# set_seed(42)

class BertEmbedder:
    def __init__(self, model_name = "./embedding-model/matscibert", device="cuda"):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModel.from_pretrained(model_name).to(device)
        self.device = device
        self.model.eval()

    @torch.no_grad()
    def encode(self, text):
            inputs = self.tokenizer(
                text,
                padding=True,
                truncation=True,
                max_length=512,
                return_tensors="pt"
            ).to(self.device)

            outputs = self.model(**inputs)
            last_hidden = outputs.last_hidden_state
            mask = inputs.attention_mask.unsqueeze(-1)
            embedding = (last_hidden * mask).sum(1) / mask.sum(1)

            return embedding.cpu().numpy().squeeze().tolist()


class FeatureExtractor(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(779, 1024), 
            nn.ReLU(), 
            nn.Linear(1024, 512), 
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.Linear(256, 32),
        )

    def forward(self, x):
        return self.net(x).squeeze(-1)


class MCDropoutModel(nn.Module):
    def __init__(self, input_dim=32, dropout_rate=0.1):
        super().__init__()
        # self.fc1 = nn.Linear(input_dim, 64)
        self.fc2 = nn.Linear(32, 128)
        self.fc3 = nn.Linear(128, 32)
        self.fc4 = nn.Linear(32, 1)
        self.dropout = nn.Dropout(dropout_rate)
        
    def forward(self, x):
        # x = F.leaky_relu(self.fc1(x), negative_slope=0.05)
        x = F.leaky_relu(self.dropout(self.fc2(x)), negative_slope=0.05)
        x = F.leaky_relu(self.dropout(self.fc3(x)), negative_slope=0.05)
        x = self.fc4(x)
        return x.squeeze(-1)
    
def mc_dropout_predict(model, X, n_samples = 50):
    set_seed(42)
    device = next(model.parameters()).device
    X_tensor = torch.FloatTensor(X).to(device)
    
    model.eval()
    for m in model.modules():
        if isinstance(m, nn.Dropout):
            m.train()
    
    predictions = []
    with torch.no_grad():
        for _ in range(n_samples):
            pred = model(X_tensor)
            predictions.append(pred.cpu().numpy())
    
    predictions = np.array(predictions)
    return predictions.mean(axis=0), predictions.std(axis=0)

def predict(Al_content=None,
            Nb_content=None,
            Ta_content=None,
            Ti_content=None,
            Zr_content=None,
            Mo_content=None,
            V_content=None,
            Cr_content=None,
            W_content=None,
            Hf_content=None,
            Ni_content=None,
            process_text = None, 
            scaler_model = "./ANN-modeling/mc_dropout_scaler.pkl",
            model_path = './ANN-modeling/mc_dropout_best_vam_only.pth', 
            embedding_model = "./embedding-model/matscibert", 
            extract_model = "./DANN-modeling/feature_extractor.pth",
            device = 'cuda'):
    
    comp_list = [
        Al_content or 0.0,
        Nb_content or 0.0,
        Ta_content or 0.0,
        Ti_content or 0.0,
        Zr_content or 0.0,
        Mo_content or 0.0,
        V_content or 0.0,
        Cr_content or 0.0,
        W_content or 0.0,
        Hf_content or 0.0,
        Ni_content or 0.0,
    ]
    
    comp_list = [float(x) for x in comp_list]

    if process_text is not None:
        # print("Embedding process text...")
        embedder = BertEmbedder(embedding_model)
        emb = embedder.encode(process_text)
        proc_emb = np.array(emb)
    
    # print("Extracting features...")
    extractor = FeatureExtractor().to(device)
    extractor.load_state_dict(torch.load(extract_model, map_location=device))
    
    # combinate process and composition
    # X = np.hstack([comp_list, proc_emb])
    comp_array = np.atleast_2d(comp_list)  # Shape: (n_samples, 11)
    proc_emb_array = np.tile(proc_emb, (comp_array.shape[0], 1))  # Repeat proc_emb for each sample
    X = np.hstack([comp_array, proc_emb_array])
    lantent_features = extractor(torch.FloatTensor(X).to(device)).cpu().detach().numpy()

    # print("Predicting...")
    model = MCDropoutModel().to(device)
    model.load_state_dict(torch.load(model_path, map_location=device))

    # scaler = StandardScaler()
    scaler = load(scaler_model)
    lantent_features_scaled = scaler.transform(lantent_features)
    pred, std = mc_dropout_predict(model, lantent_features_scaled)

    return np.round(pred, 2), np.round(std, 2)

if __name__ == "__main__":
    # set_seed(42)
    process_text = "prepared by laser powder bed fusion technique using a feedstock of pre-alloyed powders with optimized process parameters: laser power 230 W, scanning speed 900 mm/s, hatch spacing 62 μm, and layer thickness 30 μm."
    ys, std = predict(Al_content=25,
            Nb_content=25,
            Ta_content=None,
            Ti_content=25,
            Zr_content=None,
            Mo_content=None,
            V_content=25,
            Cr_content=None,
            W_content=None,
            Hf_content=None,
            Ni_content=None,
            process_text = process_text, 
            scaler_model = "./ANN-modeling/mc_dropout_scaler.pkl",
            model_path = './ANN-modeling/mc_dropout_best_vam_only.pth', 
            embedding_model = "./embedding-model/matscibert", 
            extract_model = "./DANN-modeling/feature_extractor.pth",
            device = 'cuda')
    print(ys, std)