from __future__ import annotations
import bentoml
import json
import torch
import pandas as pd
import faiss
import typing as t

from joblib import load
from typing import List
from torch_geometric.nn import global_mean_pool
from collections import OrderedDict
from preprocessing import DataToGraph, Query
import logging

#Set up logger with custom format
logger = logging.getLogger('bentoml')
new_formatter = logging.Formatter("%(asctime)s :: %(levelname)s :: %(name)s :: %(funcName)s :: %(message)s")

# Get the existing handlers of the logger
for handler in logger.handlers:
    # Update the format of each handler
    handler.setFormatter(new_formatter)

file_handler = logging.FileHandler('service.log')
file_handler.setLevel(logging.DEBUG)
file_handler.setFormatter(new_formatter)
logger.addHandler(file_handler)

#Create Custom Runners for Dynamic Model Loading
class LearningRunner(bentoml.Runnable):
    SUPPORTED_RESOURCES = ("cpu",)
    SUPPORTS_CPU_MULTI_THREADING = True

    def __init__(self, root):
        #Fetch all needed supporting data
        self.model = torch.jit.load(root + '/model.pt')
        index = faiss.read_index(root + '/labeled_embeddings.index')
        with open(root + '/index_tags.json', 'r') as f:
            tags = json.load(f)
        self.query = Query(index, tags)
        self.stats = torch.load(root + '/stats.pt')
        self.localizer_model = load(root + '/localizer_clf.joblib')
    
    @bentoml.Runnable.method(batchable=False)
    def inference(self, data: bentoml.io.JSON) -> t.Any:
        with torch.no_grad():
            results = self.model(data.x.float(), data.edge_index, data.edge_weight)
        return results
    
    @bentoml.Runnable.method(batchable=False)
    def preprocess(self, data: bentoml.io.JSON) -> t.Any:
        graph = DataToGraph(self.stats, data).convert_data_to_graph[0]
        return graph
    
    @bentoml.Runnable.method(batchable=False)
    def fetch_n_closest(self, arr, n: int=5) -> t.Any:
        neighbors = self.query.find_closest_n(arr, n)
        return neighbors
    

runner = bentoml.Runner(LearningRunner, runnable_init_params={
         "root": "/data",
         })

svc = bentoml.Service(
    name="service", 
    runners=runner,
)

@svc.api(input=bentoml.io.JSON(), output=bentoml.io.JSON())
async def generate_embedding(data) -> bentoml.io.JSON:
        data = await runner.preprocess.async_run(data)
        inference = await runner.inference.async_run(data)
        exam_array = global_mean_pool(inference, None)
        return {
            'embedding': exam_array.numpy().tolist(),
            }