import gradio as gr
import torch
import os
import numpy

from sentence_transformers import SentenceTransformer
#sentence_model = SentenceTransformer("paraphrase-multilingual-mpnet-base-v2")
#sentence_model = SentenceTransformer("all-mpnet-base-v2") 
sentence_model = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1") 

def getText( file ):
    return open( file ).read()

def axis():
    global v, a, s
    doc = getText( './axis/v' )
    embedding = sentence_model.encode( doc , show_progress_bar=False, normalize_embeddings=True )
    torch.save( embedding , './embeds/v' )
    v = embedding
    doc = getText( './axis/a' )
    embedding = sentence_model.encode( doc , show_progress_bar=False, normalize_embeddings=True )
    torch.save( embedding , './embeds/a' )
    a = embedding
    doc = getText( './axis/s' )
    embedding = sentence_model.encode( doc , show_progress_bar=False, normalize_embeddings=True )
    torch.save( embedding , './embeds/s' )
    s = embedding

def zxy( doc ):
    global v, a, s
    if not os.path.exists( './embeds/s' ):
        axis() # run this to setup axis embeddings
    else:
        if not 'v' in globals() :
            print( "loading axes from embeds" )
            v = torch.load( './embeds/v' )
            a = torch.load( './embeds/a' )
            s = torch.load( './embeds/s' )
    embedding = sentence_model.encode( doc , show_progress_bar=False , normalize_embeddings=True )
    # torch.save( embedding , './embeds/new' )
    z = ( 1.0 - embedding @ s ) / 2.0 # level 0 on top
    x = ( 1.0 - embedding @ a ) / 2.0 # ring 0 in the middle
    y = ( embedding @ v + 1.0 ) / 2.0 # sector 0 at twelve oclock
    # add one to make positive / divide by two to make range 0-1
    scale = 1
    levels = 44
    rings = 22
    sectors = 44 * 4 
    zxy = [ numpy.rint( scale * levels * z ) , numpy.rint( scale * rings * x ) , numpy.rint( scale * sectors * y ) ]
    return f"""L{ str( int( zxy[ 0 ] ) ) }/{ str( int( zxy[ 1 ] ) ) }/{ str( int( zxy[ 2 ] ) ) }
        es={ embedding @ s } ex={ embedding @ a } ey={ embedding @ v }
        z={ z } x={ x } y={ y }
        {scale} {levels} {rings} {sectors}
        """

#iface = gr.Interface(fn=zxy, inputs="text", outputs="text")
#iface.launch()
#iface.launch(share=True)  # make public