import json
import os
from concurrent import futures
import ast
import numpy as np
import traceback
from model_pb2_grpc import *
import psycopg
import sql
import model_pb2_grpc
import pandas as pd
from psycopg import sql as pgsql
from sklearn.ensemble import GradientBoostingRegressor
from model_pb2 import TrainingResponse,GetAmountResponse

"""
accompanied modules are generated by calling the following command

```
python3 -m grpc_tools.protoc -I. --python_out=. --grpc_python_out=. model.proto
```

"""
class TrainingModelService(ModelServicer):
    def __init__(self):
        super(TrainingModelService, self).__init__()
        self.model = GradientBoostingRegressor()
        self.conn = psycopg.connect("dbname=dev user=root host=frontend-node-0 port=4566")
        print("training service up!")
    def Training(self, request, context):
        print(f"get data for training!")
        try:
            with self.conn.cursor() as cur:
                cur.execute(sql.GET_BATCH_FOR_TRAINING)
                results = cur.fetchall()
                df = pd.DataFrame(list(results))
                train_y = df.loc[:, 0]
                train_x = df.drop(columns=[0,1,2,14,15])
                print(f"training!")
                self.model.fit(train_x, train_y)
                return TrainingResponse()
        except Exception as e:
            print(traceback.format_exc())
            context.set_code(grpc.StatusCode.INTERNAL)
            context.set_details(e)
            return TrainingResponse()

    def GetAmount(self, request, context):
        do_location_id = request.do_location_id
        pu_location_id = request.pu_location_id
        try:
            with self.conn.cursor() as cur:
                cur.execute(sql.GET_FEATURE_DO_LOCATION % do_location_id)
                results = cur.fetchall()
                df = pd.DataFrame(list(results))
                train_x1 = df.drop(columns=[0,1])
                cur.execute(sql.GET_FEATURE_PU_LOCATION % pu_location_id)
                results = cur.fetchall()
                df = pd.DataFrame(list(results))
                train_x2 = df.drop(columns=[0,1])
                train_x = pd.concat([train_x1, train_x2], axis=1)
                result = self.model.predict(train_x)[0]
                return GetAmountResponse(amount = result)
        except Exception as e:
            print(traceback.format_exc())
            context.set_code(grpc.StatusCode.INTERNAL)
            context.set_details(e)
            return GetAmountResponse(amount = 0.0)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.conn.close()

    def serve(self):
        server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
        model_pb2_grpc.add_ModelServicer_to_server(
            self, server)
        server.add_insecure_port('[::]:8080')
        server.start()
        server.wait_for_termination()

if __name__ == '__main__':
    for i in range(1):
        print("making the", i, "th connection")

        conn = psycopg.connect("dbname=dev user=root host=frontend-node-0 port=4566")
        with conn.cursor() as cur:
            # the following code will panic
            # because placeholder is completely not supported
            # even varchar parsed to TypeOid 0, which isn't even defined
            cur.execute("select (2, %s);", ("2333333",))
            results = cur.fetchall()
        print(results)

        conn.close()