from fastapi import FastAPI
from pydantic import BaseModel
from typing import List
from src.processor import *
from sentence_transformers import SentenceTransformer

app = FastAPI()

class Input(BaseModel):
    text1 : List
    text2 : List
    model : str = "sentence-transformers/all-MiniLM-L6-v2"

class Output(BaseModel):
    matrix : List

@app.post("/saoke-to-heatmap", response_model=Output)
def saoke_to_heatmap(payload: Input):
    saoke_spec = text_to_saoke(payload.text1)
    saoke_patent = text_to_saoke(payload.text2)

    try:
        model = SentenceTransformer(payload.model)
    except:
        model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')

    embeddings1, embeddings2 = text_to_embeddings(saoke_spec, saoke_patent, model)
    
    matrix = embeddings_to_matrix(embeddings1, embeddings2)
    print({"matrix": matrix})
    return {"matrix": matrix}

@app.post("/text-to-heatmap", response_model=Output)
def text_to_heatmap(payload: Input):
    try:
        model = SentenceTransformer(payload.model)
    except:
        model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')

    embeddings1, embeddings2 = text_to_embeddings(payload.text1, payload.text2, model)
    matrix = embeddings_to_matrix(embeddings1, embeddings2)
    print({"matrix": matrix})
    return {"matrix": matrix}