import numpy as np
from transformers import AutoModel, AutoTokenizer
from typing import Any, List
from langchain_core.embeddings import Embeddings
from pydantic import BaseModel, ConfigDict, Field, computed_field
from langchain_community.embeddings import SentenceTransformerEmbeddings
from langchain_community.llms.vllm import VLLM

import torch
import re

class PoolingStrategy():
    def __init__(self):
        pass

    def apply(self, embeddings: List[List[float]]):
        pass


class MeanPoolingStrategy(PoolingStrategy):
    def __init__(self):
        pass

    def apply(self, embeddings: List[List[float]]):
        return np.mean(embeddings, axis=0)


class SumPoolingStrategy(PoolingStrategy):
    def __init__(self):
        pass

    def apply(self, embeddings: List[List[float]]):
        return np.sum(embeddings, axis=0)


class SentenceTransformerDecorator(Embeddings):
    def __init__(self, model_name="all-MiniLM-L6-v2", **kwargs: Any):
        self.embeddings = SentenceTransformerEmbeddings(model_name=model_name, **kwargs)

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        return self.embeddings.embed_documents(texts)

    def embed_query(self, text: str) -> List[float]:
        return self.embeddings.embed_query(text)


class PretrainedEmbeddings(Embeddings):
    def __init__(self, model_name):
        self.model_name = model_name
        self.model = AutoModel.from_pretrained(model_name)
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.shape = self.model.get_input_embeddings().weight.shape

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        output_embeddings = []
        texts = list(map(lambda x: x.replace("\n", " "), texts))
        for text in texts:
            input_ids = self.tokenizer.encode(
                text, return_tensors='pt', padding=True)
            output_embeddings.append(self.model(
                input_ids).last_hidden_state.mean(dim=1))
        return output_embeddings

    def embed_query(self, text: str) -> List[float]:
        return self.embed_documents([text])[0]


class SBertEmbeddings(BaseModel, Embeddings):
    # class Config:
    #     arbitrary_types_allowed = True

    model_config = ConfigDict(
        extra="allow",
        arbitrary_types_allowed = True,
        protected_namespaces=()
    )

    model_name:str

    @computed_field
    @computed_field
    @property
    def model(self) -> Any:
        try:
            from sentence_transformers import SentenceTransformer
        except ImportError as exc:
            raise ImportError(
                "sentence-transformers is required for SBertEmbeddings. "
                "Install it with: pip install sentence-transformers"
            ) from exc
        return SentenceTransformer(self.model_name)
    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        return self.model.encode(texts).tolist()

    def embed_query(self, text: str) -> List[float]:
        return self.embed_documents([text])[0]



class VLLMEmbeddings(BaseModel, Embeddings):
    """
    Custom embedding class for VLLM.

    Args:
        client: A client instance that provides an `embeddings` method for generating embeddings.
        
    Example:
        llm = VLLM(model = "meta-llama/Meta-Llama-3-8B-Instruct", **params)
        client = llm.client
        embeddings = VLLMEmbeddings(client=client)
    """
    model_kwargs: dict[str, Any] = Field(default_factory=dict)
    model_name: str = "meta-llama/Meta-Llama-3-8B-Instruct"
    client: Any|None = None  # The client instance for interacting with the VLLM server
    
    def __init__(self, **kwargs: Any):
        super().__init__(**kwargs)
        # params= {
        #         "trust_remote_code": True,
        #         "max_new_tokens": 8192,
        #         "top_k": 1,
        #         "top_p": 0.95,
        #         "temperature": 0,
        #         "tensor_parallel_size": 1,
        #         "vllm_kwargs": {
        #           "gpu_memory_utilization": 0.8,
        #           "task": "generate"
        #         }
        # }
        
       
        params = self.model_kwargs
        print(self.model_name)
        llm = VLLM(model = self.model_name, **params)
        self.client = llm.client
        if self.client is None:
            raise ValueError("The client must be provided and cannot be None.")
       
    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """
        Embed a list of documents using the VLLM client.

        Args:
            texts: A list of strings representing the documents to embed.

        Returns:
            A list of embeddings, one for each document.
        """
        try:
            # response = self.client.embed(texts)
            # result = []
            # for(i, r) in enumerate(response):
            #     result.append(r.outputs.embedding)
            # return result
            response = self.client.embed(texts)
            return [r.outputs.embedding for r in response]
        except AttributeError:
            raise ValueError("The client must have an `embeddings` method that returns embeddings.")

    def embed_query(self, text: str) -> List[float]:
        """
        Embed a single query using the VLLM client.

        Args:
            text: A string representing the query to embed.

        Returns:
            The embedding for the query.
        """
        try:
            response = self.client.embed([text])
            # result = []
            # for(i, r) in enumerate(response):
            #     result.append(r.outputs.embedding)
            return response[0].outputs.embedding
        except AttributeError:
            raise ValueError("The client must have an `embeddings` method that returns embeddings.")
        
        


class ClinicalBERTEmbeddings(BaseModel, Embeddings):
    model_config = ConfigDict(
        extra="allow",
        arbitrary_types_allowed = True,
        protected_namespaces=()
    )

    model_name: str = "emilyalsentzer/Bio_ClinicalBERT"
    pooling: str = "pooler"   # "cls" | "pooler" | "mean"
    normalize: bool = True
    device: str = "cuda"

    @computed_field
    @property
    def model(self) -> Any:
        _device = self.device or ("cuda" if torch.cuda.is_available() else "cpu")
        _model = AutoModel.from_pretrained(self.model_name)
        _model.to(_device)
        _model.eval()
        return _model

    @computed_field
    @property
    def tokenizer(self) -> Any:
        return AutoTokenizer.from_pretrained(self.model_name)

    def __init__(self, **kwargs: Any):
        super().__init__(**kwargs)
        # This is to trigger the model loading and move to device
        _ = self.model

    def _pool(self, outputs, attention_mask):
        if self.pooling == "cls":
            return outputs.last_hidden_state[:, 0]

        elif self.pooling == "pooler":
            # safer fallback if pooler doesn't exist
            if hasattr(outputs, "pooler_output") and outputs.pooler_output is not None:
                return outputs.pooler_output
            return outputs.last_hidden_state[:, 0]

        elif self.pooling == "mean":
            mask = attention_mask.unsqueeze(-1).expand(outputs.last_hidden_state.size()).float()
            summed = (outputs.last_hidden_state * mask).sum(1)
            counts = mask.sum(1)
            return summed / counts

        else:
            raise ValueError(f"Unknown pooling: {self.pooling}")

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        texts = [t.replace("\n", " ") for t in texts]

        inputs = self.tokenizer(
            texts,
            return_tensors="pt",
            padding=True,
            truncation=True
        ).to(self.device)

        with torch.no_grad():
            outputs = self.model(**inputs)

        embeddings = self._pool(outputs, inputs["attention_mask"])

        if self.normalize:
            embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)

        return embeddings.cpu().numpy().tolist()

    def embed_query(self, text: str) -> List[float]:
        return self.embed_documents([text])[0]
    
    

class BioSentVecEmbeddings(BaseModel, Embeddings):
    model_config = ConfigDict(
        extra="allow",
        arbitrary_types_allowed = True,
        protected_namespaces=()
    )

    model_path: str
    normalize: bool = True
    model_name: str = "biosentvec"

    @computed_field
    @property
    def model(self) -> Any:
        try:
            import sent2vec
        except ImportError as exc:
            raise ImportError(
                "sent2vec is required for BioSentVecEmbeddings. "
                "Install it with: pip install sent2vec"
            ) from exc
        _model = sent2vec.Sent2vecModel()
        _model.load_model(self.model_path)
        return _model

    @computed_field
    @property
    def shape(self) -> tuple:
        return (700,)

    def __init__(self, **kwargs: Any):
        super().__init__(**kwargs)
        # This is to trigger the model loading
        _ = self.model

    # =========================
    # TEXT PREPROCESSING
    # =========================
    def _preprocess_text(self, text: str) -> str:
        if text is None:
            return ""
        text = str(text).strip().lower()
        text = re.sub(r"\s+", " ", text).strip()
        return text

    # =========================
    # SENTENCE SPLITTING
    # =========================
    def _split_into_sentences(self, text: str) -> List[str]:
        if not text:
            return []
        sentences = re.split(r'(?<=[.!?])\s+', text)
        return [s.strip() for s in sentences if s.strip()]

    # =========================
    # CORE EMBEDDING LOGIC
    # =========================
    def _embed_single(self, text: str) -> np.ndarray:
        text = self._preprocess_text(text)

        if not text:
            return np.zeros(self.shape[0], dtype=float)

        sentences = self._split_into_sentences(text)

        if not sentences:
            vec = np.asarray(self.model.embed_sentence(text)).reshape(-1)
            return vec

        vectors = []
        for sent in sentences:
            vec = np.asarray(self.model.embed_sentence(sent)).reshape(-1)
            vectors.append(vec)

        if not vectors:
            return np.zeros(self.shape[0], dtype=float)

        mat = np.vstack(vectors)
        mean_vec = np.mean(mat, axis=0)
        
        # weights = [len(sent.split()) for sent in sentences]
        # mean_vec = np.average(mat, axis=0, weights=weights)

        return mean_vec

    # =========================
    # PUBLIC API
    # =========================
    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        embeddings = []

        for text in texts:
            vec = self._embed_single(text)

            if self.normalize:
                norm = np.linalg.norm(vec)
                if norm > 0:
                    vec = vec / norm

            embeddings.append(vec.tolist())

        return embeddings

    def embed_query(self, text: str) -> List[float]:
        return self.embed_documents([text])[0]

# class EmbeddingGenerator(ABC):

#     def __init__(self):
#         self.shape = None

#     @abstractmethod
#     def generate(self, text):
#         pass

# # https://github.com/huggingface/transformers/issues/1791
# class PretrainedEmbeddingGenerator(EmbeddingGenerator):
#     def __init__(self, model_name, split_text=True):
#         self.model_name = model_name
#         self.model = AutoModel.from_pretrained(model_name)
#         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
#         self.shape = self.model.get_input_embeddings().weight.shape
#         self.split_text = split_text

#     def generate(self, text):
#         if self.split_text:
#             sentences = sent_tokenize(text)
#             output_embeddings = None
#             for sentence in sentences:
#                 input_ids  = self.tokenizer.encode(sentence, return_tensors='pt', padding=True)
#                 if output_embeddings == None:
#                     output_embeddings = self.model(input_ids).last_hidden_state.mean(dim=1)
#                 else:
#                     output_embeddings += self.model(input_ids).last_hidden_state.mean(dim=1)
#             if output_embeddings == None:
#                 output_embeddings = torch.empty((1,self.model.get_input_embeddings().weight.shape[1]))
#         else:
#             encoded_input = self.tokenizer(text, return_tensors='pt')
#             input_ids = self.tokenizer.encode(text, add_special_tokens=True, max_length=self.tokenizer.model_max_length, return_tensors='pt')
#             # input_ids  = encoded_input.input_ids[:self.tokenizer.model_max_length]
#             output_embeddings = self.model(input_ids)
#             # output_embeddings = self.model(**encoded_input)
#             # sentence embedding
#             output_embeddings = output_embeddings.last_hidden_state.mean(dim=1)
#         return torch.flatten(output_embeddings).detach().numpy()


# class LLMEmbeddingGenerator(EmbeddingGenerator):
#     pass


# class SBertEmbeddingGenerator(PretrainedEmbeddingGenerator):
#     def __init__(self, model_name, split_text=True):
#         self.model = SentenceTransformer(model_name)
#         self.shape = self.model._modules['0'].get_word_embedding_dimension()
#         self.split_text = split_text

#     def generate(self, text):
#         if text == None or type(text) != str:
#             embeddings = None
#         else:
#             if self.split_text:
#                 sentences = sent_tokenize(text)
#                 embeddings = self.model.encode(sentences)
#             else:
#                 embeddings = self.model.encode(text)
#         return embeddings


# # NOT modelden input size'ı anlama,
# def create_embeddings_new(generator: EmbeddingGenerator, data, column, drop_column=True):
#     column_embeddings_dict = defaultdict(lambda: np.zeros(generator.shape))
#     for index, row in tqdm(data.iterrows()):
#         # if index == 10:
#         #   break
#         text = data[column][index]
#         embeddings = generator.generate(text)

#     # TODO benzer olan ilacın embedding değerini vererek dene
#         # embedding check none type
#         if embeddings is None or len(embeddings) == 0:
#             sum_of_embeddings = np.zeros(generator.shape)
#         else:
#             sum_of_embeddings = np.sum(embeddings, axis=0)
#         # column_embeddings_dict[row['id']] = sum_of_embeddings.reshape(1, -1) # 2d
#         column_embeddings_dict[row['id']] = sum_of_embeddings
#         # data.iloc[index][column+'_embedding']=sum_of_embeddings

#     data[column+'_embedding'] = pd.Series(column_embeddings_dict.values())
#     if (drop_column):
#         data.drop([column], axis=1, inplace=True)
#     # data[column+'_embedding'] = [column_embeddings_dict[row['name']] for index, row in data.iterrows()]
#     return column_embeddings_dict
