from fastapi import FastAPI, Request
from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
import torch
from pydantic import BaseModel
from typing import Optional

app = FastAPI()


class InputText(BaseModel):
    text : str
    threshold: Optional[float] = None


model_name = "cardiffnlp/twitter-xlm-roberta-base-sentiment"
sentiment_model = AutoModelForSequenceClassification.from_pretrained(model_name)
sentiment_tokenizer = AutoTokenizer.from_pretrained(model_name)
sentiment_model.config.id2label[3] = "mixed"

model_name = 'qanastek/51-languages-classifier'
language_model = AutoModelForSequenceClassification.from_pretrained(model_name)
language_tokenizer = AutoTokenizer.from_pretrained(model_name)







@app.get("/")
def greet_json():
    return {"Hello": "World!"}



@app.post("/language_detection")
async def language_detection(inp: InputText):
    inputs = language_tokenizer(inp.text, return_tensors='pt')
    with torch.no_grad():
        logits = language_model(**inputs).logits
    
    softmax = torch.nn.functional.sigmoid(logits)
    
    # Apply the threshold by creating a mask
    mask = softmax >= inp.threshold
    
    # Filter the tensor based on the threshold
    filtered_x = softmax[mask]
    
    # Get the sorted indices of the filtered tensor
    sorted_indices = torch.argsort(filtered_x, descending=True)
    
    # Map the sorted indices back to the original tensor indices
    original_indices = torch.nonzero(mask, as_tuple=True)[1][sorted_indices]
    
    return [{"label":language_model.config.id2label[predicted_class_id.tolist()], "score":softmax[0, predicted_class_id].tolist()} for predicted_class_id in original_indices]


@app.post("/sentiment_score")
async def sentiment_score(inp: InputText):
    text = inp.text
    inputs = sentiment_tokenizer(text[:2500], return_tensors='pt')
    
    with torch.no_grad():
        logits = sentiment_model(**inputs).logits #+ 1
    
    
    print(logits)
    
    logits = logits + logits[0,1].abs()
    
    # print(torch.nn.functional.sigmoid(logits))
    
    # logits = logits / 10
    
    # print(logits)
    
    # print(torch.abs(logits[0,0] - logits[0,-1]))
    # print(logits[0,1]//torch.max(torch.abs(logits[0,::2])))
    
    logits = torch.cat(
        (
            logits, (
                # ( logits[0,1] + torch.sign(logits[0,0] - logits[0,-1]) * (logits[0,0] - logits[0,-1])/2 )/2 + 
                # (logits[0,0] + logits[0,-1])/20
                (1 - torch.abs(logits[0,0] - logits[0,-1])*(2+(logits[0,1]//torch.max(torch.abs(logits[0,::2])))))
                ).unsqueeze(0).unsqueeze(0)
        ), dim=-1
    )
    
    softmax = torch.nn.functional.softmax(
        logits, 
        dim=-1
    )
    
    return [{"label":sentiment_model.config.id2label[predicted_class_id.tolist()], "score":softmax[0, predicted_class_id].tolist()} for predicted_class_id in softmax.argsort(dim=-1, descending=True)[0]]