import torch


def extract_text_feature(prompt, model, processor, device='cpu'):
    """Extract text features

    Args:
        prompt: a single text query
        model: OwlViT model
        processor: OwlViT processor
        device (str, optional): device to run. Defaults to 'cpu'.
    """
    device = 'cpu'
    if torch.cuda.is_available():
        device = 'cuda'
    with torch.no_grad():
        input_ids = torch.as_tensor(processor(text=prompt)[
                                    'input_ids']).to(device)
        print(input_ids.device)
        text_outputs = model.owlvit.text_model(
            input_ids=input_ids,
            attention_mask=None,
            output_attentions=None,
            output_hidden_states=None,
            return_dict=None,
        )
        text_embeds = text_outputs[1]
        text_embeds = model.owlvit.text_projection(text_embeds)
        text_embeds /= text_embeds.norm(p=2, dim=-1, keepdim=True) + 1e-6
        query_embeds = text_embeds
    return input_ids, query_embeds


def prompt2vec(prompt: str, model, processor):
    """ Convert prompt into a computational vector

    Args:
        prompt (str): Text to be tokenized

    Returns:
        xq: vector from the tokenizer, representing the original prompt
    """
    # inputs = tokenizer(prompt, return_tensors='pt')
    # out = clip.get_text_features(**inputs)
    input_ids, xq = extract_text_feature(prompt, model, processor)
    input_ids = input_ids.detach().cpu().numpy()
    xq = xq.detach().cpu().numpy()
    return input_ids, xq


def tune(clf, X, y, iters=2):
    """ Train the Zero-shot Classifier

    Args:
        X (numpy.ndarray): Input vectors (retreived vectors)
        y (list of floats or numpy.ndarray): Scores given by user
        iters (int, optional): iterations of updates to be run
    """
    assert len(X) == len(y)
    # train the classifier
    clf.fit(X, y, iters=iters)
    # extract new vector
    return clf.get_weights()


class Classifier:
    """Multi-Class Zero-shot Classifier
    This Classifier provides proxy regarding to the user's reaction to the probed images.
    The proxy will replace the original query vector generated by prompted vector and finally
    give the user a satisfying retrieval result.

    This can be commonly seen in a recommendation system. The classifier will recommend more 
    precise result as it accumulating user's activity.
    
    This is a multiclass classifier. For N queries it will set the all queries to the first-N classes
    and the last one takes the negative one.
    """

    def __init__(self, xq: list):
        init_weight = torch.Tensor(xq)
        self.num_class = xq.shape[0]
        DIMS = xq.shape[1]
        # note that the bias is ignored, as we only focus on the inner product result
        self.model = torch.nn.Linear(DIMS, self.num_class, bias=False)
        # convert initial query `xq` to tensor parameter to init weights
        self.model.weight = torch.nn.Parameter(init_weight)
        # init loss and optimizer
        self.loss = torch.nn.BCEWithLogitsLoss()
        self.optimizer = torch.optim.SGD(self.model.parameters(), lr=0.1)

    def fit(self, X: list, y: list, iters: int = 5):
        # convert X and y to tensor
        X = torch.Tensor(X)
        X /= torch.norm(X, p=2, dim=-1, keepdim=True)
        y = torch.Tensor(y).long()
        # Generate labels for binary classification and ignore outbound labels
        non_ind = y > self.num_class
        y = torch.nn.functional.one_hot(y % self.num_class, num_classes=self.num_class).float()
        y[non_ind] = 0
        for i in range(iters):
            # zero gradients
            self.optimizer.zero_grad()
            # Normalize the weight before inference
            # This will constrain the gradient or you will have an explosion on query vector
            self.model.weight.data /= torch.norm(self.model.weight.data, p=2, dim=-1, keepdim=True)
            # forward pass
            out = self.model(X)
            # compute loss
            loss = self.loss(out, y)
            # backward pass
            loss.backward()
            # update weights
            self.optimizer.step()

    def get_weights(self):
        xq = self.model.weight.detach().numpy()
        return xq
    
class SplitLayer(torch.nn.Module):
    def forward(self, x):
        return torch.split(x, 1, dim=-1)