from fastapi import FastAPI
from clearml import Task, Dataset
import ast
import uvicorn
import os
from pycaret.classification import load_model, predict_model
from tvDatafeed import Interval as TVInterval

interval_dict = {i.value: i for i in TVInterval}

# load model
model_task = Task.get_task(os.getenv("MODEL_TASK_ID"))
model_filename = model_task.artifacts['model'].get_local_copy()
model = load_model(model_filename[:-4])

# get data info
dataset_id = model_task.data.hyperparams['Args']['dataset_id'].value
data_params = Dataset.get(dataset_id).get_metadata()

# create get_data definition from task
parsed_script = ast.parse(Task.get_task(task_name='get_data').export_task()['script']['diff'])
get_data_script = [node for node in ast.walk(parsed_script) if isinstance(node, ast.FunctionDef) and node.name == 'get_data'][0]
get_data_module = ast.Module([get_data_script])
get_data_module.type_ignores = []
get_data = compile(get_data_module, filename='<string>', mode='exec')
exec(get_data)

# create process_data definition from task
parsed_script = ast.parse(Task.get_task(task_name='process_data').export_task()['script']['diff'])
process_data_script = [node for node in ast.walk(parsed_script) if isinstance(node, ast.FunctionDef) and node.name == 'process_data'][0]
process_data_module = ast.Module([process_data_script])
process_data_module.type_ignores = []
process_data = compile(process_data_module, filename='<string>', mode='exec')
exec(process_data)


app = FastAPI()

@app.get(f'/{data_params["exchange"]}/{data_params["symbol"]}')
def predict():
    data = get_data(data_params['symbol'], data_params['exchange'], data_params['interval'], 100)
    processed_data = process_data(data, int(data_params['window_length']), int(data_params['target_length']), training=False)
    predictions = predict_model(model, processed_data[-1:])
    return predictions[['prediction_label', 'prediction_score']]


if __name__ == '__main__':
    uvicorn.run(app)