import asyncio

import asyncio
import json
from http import HTTPStatus
from typing import Optional
import bentoml
import numpy as np
from <http://bentoml.io|bentoml.io> import JSON
from pydantic import BaseModel, validator

mayo_icu_runner_1 = bentoml.keras.get("bentomodel_1_service_1:latest").to_runner()
mayo_icu_runner_2 = bentoml.keras.get("bentomodel_2_service_1:latest").to_runner()
mayo_icu_runner_3 = bentoml.keras.get("bentomodel_3_service_1:latest").to_runner()
gates_icu_runner_1 = bentoml.keras.get("bentomodel_4_service_1:latest").to_runner()
gates_icu_runner_2 = bentoml.keras.get("bentomodel_5_service_1:latest").to_runner()
gates_icu_runner_3 = bentoml.keras.get("bentomodel_6_service_1:latest").to_runner()

svc = bentoml.Service("bento_service",
                      runners=[mayo_icu_runner_1, mayo_icu_runner_2, mayo_icu_runner_3,
                               gates_icu_runner_1, gates_icu_runner_2, gates_icu_runner_3])

model_groups = {"mayo_icu": [mayo_icu_runner_1, mayo_icu_runner_2, mayo_icu_runner_3],
                "gates_icu": [gates_icu_runner_1, gates_icu_runner_2, gates_icu_runner_3]}

prediction_groups = {"mayo_icu": ["mayo_icu_runner_1", "mayo_icu_runner_2", "mayo_icu_runner_3"],
                     "gates_icu": ["gates_icu_runner_1", "gates_icu_runner_2", "gates_icu_runner_3"]}


class JsonInput(BaseModel):
    organization: str
    subgroup: str
    patients: dict
    request_id: Optional[int]

    @validator("patients")
    def patients_must_contain_temporal_and_nontemporal_features(cls, v):
        for i in v:
            if "nontemporal_features" not in v[i]:
                raise ValueError('must contain nontemporal_features for patient:', i)
            if "temporal_features" not in v[i]:
                raise ValueError('must contain temporal_features for patient:', i)
            return v


@svc.api(input=JSON(pydantic_model=JsonInput), output=JSON())
async def classify(json_input: JSON, ctx: bentoml.Context) -> json:
    input_dataset = json_input.dict()
    predictions={}
    model_group = input_dataset["organization"] + "_" + input_dataset["subgroup"]
    for patient in input_dataset["patients"]:
        try:
            runners_list = __build_runner_list(__get_runners_from_model_group(model_group),
                                               np.array(input_dataset["patients"][patient]["nontemporal_features"]),
                                               np.array(input_dataset["patients"][patient]["temporal_features"]))

            all_results = await asyncio.gather(*runners_list)

            last_result_from_all_results = [result[-1][-1][-1].item() for result in all_results]

            all_last_results = dict(zip(prediction_groups[model_group], last_result_from_all_results))

            predictions[patient] = all_last_results

        except NotImplementedError:
            ctx.response.status_code = HTTPStatus.NOT_FOUND
            ctx.response.headers.append("Model-group-notfound", model_group)

    return predictions


def __get_runners_from_model_group(model_group: str):
    try:
        return model_groups[model_group]
    except:
        raise NotImplementedError


def __build_runner_list(list_of_runners, model_input_nontemporal, model_input_temporal):
    runner_list = [i.async_run(model_input_nontemporal, model_input_temporal) for i in list_of_runners]
    return runner_list