from typing import Optional, Union, List, Dict, Tuple, Iterator

from itertools import chain
from typing import AnyStr, Callable
import statistics
import pandas as pd
import numpy as np
from uuid import UUID, uuid4
from pathlib import Path

from datalogue.clients.http import _HttpClient, HttpMethod
from datalogue.clients.ontology import _OntologyClient
from datalogue.models.ontology import Ontology, OntologyNode, DataRef, TrainingDataColumn
from datalogue.models.model import Model, ModelType, TrainingStatusType, OrderList, TrainingState, \
    training_status_type_from_str, Deployment, Section, SynthesizedSection, OriginalSection, SynthesisStats, \
    TrainingData
from datalogue.errors import DtlError
from datalogue.dtl_utils import _parse_list, _parse_uuid
from uuid import UUID
from datalogue.models.metrics import ArgoMetrics


class _SynthesisClient:
    """
    Client to interact with SynthesizedData
    """
    def __init__(self, http_client: _HttpClient):
        self.http_client = http_client

    def _get_sections(self, endpoint: str, datastore_ids: List[Union[str, UUID]], page: int, size: int) -> \
            Union[DtlError, dict]:
        params = {
            "datastoreId": datastore_ids,
            "page": page,
            "size": size
        }
        res = self.http_client.make_authed_request(endpoint, HttpMethod.GET, params=params)
        if isinstance(res, DtlError):
            return res
        return res

    def get(self, datastore_ids: List[Union[str, UUID]], ontology_id: Union[str, UUID], page: int, page_size: int) -> \
            Union[DtlError, List[SynthesizedSection]]:
        """
        Retrieves all data that was generated by the run of the synthesize function. The parameters filter the list of section returned

        :param datastore_ids: list of datastores to be used to filter the section to be used as a baseline. If None is used, then it uses all the synthesize all the data annotated for a given ontology
        :param ontology_id: id of the ontology for which we are doing the synthesizing
        :param page: which page of data to be returned
        :param page_size: number of elements in the page
        :return: List of Sections

        """
        res = self._get_sections(f"/themis/sections/ontology/{ontology_id}/synthesized", datastore_ids, page, page_size)
        if isinstance(res, DtlError):
            return res
        return _parse_list(SynthesizedSection.from_payload)(res)

    def _original(self, datastore_ids: List[Union[str, UUID]], ontology_id: Union[str, UUID], page: int, size: int):
        res = self._get_sections(f"/themis/sections/ontology/{ontology_id}/original", datastore_ids, page, size)
        if isinstance(res, DtlError):
            return res
        return _parse_list(OriginalSection.from_payload)(res)

    def _count_original(self, datastore_ids: List[Union[str, UUID]], ontology_id: Union[str, UUID]):
        params = {
            "datastoreId": datastore_ids,
        }
        res = self.http_client.make_authed_request(f"/themis/sections/ontology/{ontology_id}/original/count", HttpMethod.GET, params=params)
        if isinstance(res, DtlError):
            return res
        if isinstance(res, int):
            return res
        return DtlError(f"{res} is expected to be int")

    def _create(self, ontology_id: Union[str, UUID], section: SynthesizedSection) -> Union[DtlError, None]:
        body = section._as_payload()
        return self.http_client.make_authed_request(f"/themis/sections/ontology/{ontology_id}/synthesized",
                                                       HttpMethod.POST, body=body)

    def _clean(self, datastore_ids: List[Union[str, UUID]],
               ontology_id: Union[str, UUID]) -> Union[DtlError, None]:
        params = {
            "datastoreId": datastore_ids,
        }
        return self.http_client.make_authed_request(f"/themis/sections/ontology/{ontology_id}/synthesized", HttpMethod.DELETE, params=params)

    def synthesize(self,
                   datastore_ids: List[Union[str, UUID]],
                   ontology_id: Union[str, UUID],
                   synthesis_algorithm: Callable[[OriginalSection], List[SynthesizedSection]],
                   overwrite_synthesis=False
                   ) -> Union[DtlError, SynthesisStats]:
        """
        Creates a synthesized training data set. That Data is generated from some original annotated training data.

        :param datastore_ids: list of datastores to be used to filter the section to be used as a baseline. If None is used, then it uses all the data annotated for a given ontology
        :param ontology_id: id of the ontology for which we are doing the synthesization
        :param synthesis_algorithm: function that is to be called on sections and synthesizes training data
        :param overwrite_synthesis: when set to true, the previously synthesized data is deleted and replaced by the one generated by this call. This flag takes into account the filters that are activated (ontology or datastore)
        :return: Statistics about the synthesis
        """
        if overwrite_synthesis:
            self._clean(datastore_ids, ontology_id)

        current_page = 1
        batch_size = 10
        to_process = self._count_original(datastore_ids, ontology_id)
        if isinstance(to_process, DtlError):
            return DtlError("Nothing to process")
        page = self._original(datastore_ids, ontology_id, current_page, batch_size)
        processed = 0
        created = []

        while not isinstance(page, DtlError) and len(page) > 0:
            prepared_to_send = [synthesis_algorithm(item) for item in page]
            to_create = chain.from_iterable(prepared_to_send)
            created_in_batch = 0
            for section in to_create:
                res = self._create(ontology_id, section)
                if isinstance(res, DtlError):
                    return res
                created_in_batch = created_in_batch + 1
            created.append(created_in_batch)
            processed = processed + len(page)
            current_page = current_page + 1
            print(f"{processed/to_process*100}%")
            if len(page) < batch_size:
                break
            page = self._original(datastore_ids, ontology_id, current_page, batch_size)
        print("Success!")
        if len(created) > 1:
            return SynthesisStats(sum(created), processed, statistics.mean(created), statistics.stdev(created))
        else:
            return SynthesisStats(sum(created), processed, statistics.mean(created), 0)


class TrainingCoach:

    @staticmethod
    def check(ontology: Ontology) -> pd.DataFrame:
        leaves = ontology.leaves()
        rows = []
        for leaf in leaves:
            is_without_data = leaf.total_training_count == 0
            unique_names = set(map(lambda x: x.datastore_name, leaf.training_data_columns))
            has_got_only_one_source = len(unique_names) < 2
            total_smaller_than_x = leaf.total_training_count < 2000
            rows.append([leaf.name, leaf.id, is_without_data, has_got_only_one_source, total_smaller_than_x])

        return pd.DataFrame(rows,
                            columns=['Class', 'Id', 'WARNING: No data', 'WARNING: <2 sources', 'WARNING: <2000 data points'])


class _MLBase:
    def __init__(self, http_client: _HttpClient):
        self.http_client = http_client
    
    def get(self, model_id: UUID) -> Union[DtlError, TrainingState]:
        """""
        Get :class:`Trainings` based on the given model_id

        :param model_id:
        :return: Requested training if successful, else returns :class:`DtlError`
        """

        res = self.http_client.make_authed_request(f'/argus/trainings/{str(model_id)}',
                                                   HttpMethod.GET)
        if isinstance(res, DtlError):
            return DtlError(f"Something happened while trying to get training for {model_id}", res)

        return TrainingState._from_payload(res)
    
    def cancel(self, model_id: UUID) -> Union[DtlError, bool]:
        """
        Cancel training given a model_id

        :param model_id:
        :return: status of the training if successful, else returns :class:`DtlError`
        """

        res = self.http_client.make_authed_request(
            f'/argus/trainings/{str(model_id)}/cancel',
            HttpMethod.POST)

        if isinstance(res, DtlError):
            return res
        else:
            return True

    def get_model_url(self, model_id: UUID) -> Union[DtlError, Optional[str]]:
        """
        Return url of a deployment for a given model_id

        :param model_id:
        :return: url if there is an active deployment for the training, otherwise returns None
        """

        res = self.http_client.make_authed_request(f'/argus/deployments/model-url/{str(model_id)}', HttpMethod.GET)

        if isinstance(res, DtlError):
            return res

        return res["url"]

    def get_model_history(self, ontology_id: UUID, limit: Optional[int] = 20) -> Union[DtlError, List[Model]]:
        """
        Get :class:`Training` history based on the given ontology_id
        
        :param ontology_id:
        :return: List of trainings if successful, else returns :class:`DtlError`
        """

        params = {}
        if limit is not None:
            if not isinstance(limit, int) or limit < 0:
                return DtlError("limit should be a positive integer")
            params["limit"] = limit

        res = self.http_client.make_authed_request(
            f'/yggy/ontology/{str(ontology_id)}/trainings/history',
            HttpMethod.GET, params=params)

        if isinstance(res, DtlError):
            return res
        else:
           return _parse_list(Model._from_payload)(res)


class _MLClient:
    """
    Client to interact with the ML
    """
    def __init__(self, http_client: _HttpClient):
        self.http_client = http_client
        self.cbc = _CBCClient(http_client)
        self.ner = _NERClient(http_client)
        self.deployment = _DeploymentClient(http_client)


class _CBCDataClient:
    """
    Client to interact with TrainingData
    """

    def __init__(self, http_client: _HttpClient):
        self.http_client = http_client

    def add(self, store_id: UUID, store_name: str, refs: List[DataRef], count: Optional[int] = None) -> Union[DtlError, List[UUID]]:
        """
        Attaches paths in the given `store_id` to the the nodes of the ontology.

        :param store_id: Id of the datastore that it going to be read
        :param store_name: Name of the datastore that it going to be read
        :param refs: List of data references to show which paths are going to be attached which ontology nodes
        :param count: Number of n first rows which we want to attach, if not passed we attach all rows
        :return: List of stream ids which are the jobs that is transferring data from datastore to Themis
        """
        stream_ids = []

        for dataRef in refs:
            for path in dataRef.path_list:

                # Creates a dataset per attachment as UI does.
                dataset_id = self.__create_dataset(store_name)
                if isinstance(dataset_id, DtlError):
                    return dataset_id

                stream_id = self.__transfer_data_from_datastore(store_id, dataset_id, dataRef.node_id, path, count)
                if isinstance(stream_id, DtlError):
                    return stream_id

                self.__update_node(dataRef.node_id, path, dataset_id, stream_id)
                stream_ids.append(stream_id)
        return stream_ids

    def __create_dataset(self, store_name: str) -> Union[DtlError, UUID]:
        payload = {
            "title": store_name,
            "tags": [],
            "label_map": {}
        }

        res = self.http_client.make_authed_request('/themis/dataset', HttpMethod.POST, body=payload)
        dataset_id = UUID(res.get("id"))

        if dataset_id is None:
            return DtlError("There is no dataset id in response!")
        return dataset_id

    def __transfer_data_from_datastore(self, store_id: UUID, dataset_id: UUID, node_id: UUID, path: List[str], count: Optional[int]) -> UUID:
        params = {
            'sourceId': str(store_id),
            'trainingDataId': str(dataset_id),
            'class': str(node_id),
            'path': path
        }
        if count is not None:
            params['count'] = count
        res = self.http_client.make_authed_request('/scout/run/training-data', HttpMethod.POST, params=params)

        if isinstance(res, DtlError):
            return res

        stream_id = UUID(res["streamId"])
        return stream_id

    def __update_node(self, node_id: UUID, path: List[str], dataset_id: UUID, stream_id: UUID) -> List[str]:
        #TODO We should remove this part when we move training data adding functionality out of Yggy
        entity_res = self.http_client.make_authed_request(f"/yggy/entity/{node_id}", HttpMethod.GET)
        training_data_list = entity_res.get("trainingDataInfo")
        if training_data_list is None:
            training_data_list = []

        training_data_list.append({
            "datasetId": str(dataset_id),
            "nodePath": path,
            "streamId": str(stream_id)
        })

        self.http_client.make_authed_request(f"/yggy/entity/{node_id}", HttpMethod.POST,
                                                body={"trainingDataInfo": training_data_list})

        dataset_ids = []
        for trainingData in training_data_list:
            dataset_ids.append(trainingData["datasetId"])

        return dataset_ids


class _CBCClient(_MLBase):
    """
    Client to interact with the CBC
    """  
    def __init__(self, http_client: _HttpClient):
        self.http_client = http_client
        self.data = _CBCDataClient(http_client)

    def get_deployments(self,
                        page: int = 1,
                        item_per_page: int = 25) -> Union[DtlError, List[TrainingState]]:
        """
        Retrieve a list of all models currently deployed.

        :param page:
        :param item_per_page:
        :return:
        """
        from_item = (page-1)*item_per_page
        params = {
            'from': from_item,
            'limit': item_per_page,
            'status': 'Successful',
            'model-type': str(ModelType.Cbc)
        }

        res = self.http_client.make_authed_request(
            f'/argus/deployments', HttpMethod.GET, params=params)
        if isinstance(res, DtlError):
            return res
        else:
            return _parse_list(Deployment._from_payload)(res)

    def deploy(self, model_id: UUID, ontology_id: UUID) -> Union[DtlError, UUID]:
        """
        Deploy :class:`Training` based on the given model_id and ontology_id
        :return: True if successful, else returns :class:`DtlError`
        """
        res = self.http_client.make_authed_request(
            f'/yggy/ontology/{str(ontology_id)}/trainings/{str(model_id)}/deploy',
            HttpMethod.POST, {})

        if isinstance(res, DtlError):
            return res
        else:
            deployment_id = UUID(res.get('deploymentId'))
            return deployment_id

    def train(self, ontology_id: UUID, force_run: bool = True) -> Union[DtlError, UUID, List]:
        """
        Starts training for a given ontology_id

        :param ontology_id: id of the ontology to be used for training
        :param force_run: allows to run training even when there are less data than should be
        :return: Either a :class:`DtlError` if any error occurs or :class:`True` if training is requested successfully.
        """

        params = dict()
        params["model-type"] = ModelType.Cbc
        params["force-train"] = force_run

        res = self.http_client.make_authed_request(f"/yggy/ontology/{ontology_id}/train", HttpMethod.POST, params=params)
        if isinstance(res, DtlError):
            return res
        if isinstance(res, List):
            # TODO we should handle response with given content: [{'TagsLessThanLimit': {'currentTotalTags': X, 'recommendedTags': Y}}]
            return res
        else:
            model_id = UUID(res.get('trainingId'))
            return model_id

    def get_models(self, ontology_id: UUID, status_type: Optional[TrainingStatusType] = None, sort: Optional[OrderList] = OrderList.desc, limit: Optional[int] = 20) -> Union[DtlError, List[TrainingState]]:
        """
        Get :class:`Trainings` based on the given ontology_id
        
        :param ontology_id:
        :return: List of trainings if successful, else returns :class:`DtlError`
        """

        params = {}

        if limit is not None:
            if not isinstance(limit, int) or limit < 0:
                return DtlError("limit should be a positive integer")
            params["limit"] = limit

        params["model-type"] = ModelType.Cbc

        if status_type is not None:
            params["status-type"] = status_type.value

        if sort is not None:
            params["sort"] = sort.value

        res = self.http_client.make_authed_request(
            f'/yggy/ontology/{str(ontology_id)}/trainings', HttpMethod.GET, params=params)

        if isinstance(res, DtlError):
            return res
        else:
            return _parse_list(TrainingState._from_payload)(res)


    def upload(self, ontology_id: UUID, name: str, argo_work_dir: Path, experiment_name: str, class_map_path: str, model_id: Optional[UUID] = None) -> Optional[bool]:
        """	
        Processes and pushes via the Argus api to the defined	
        ontology.	
        :param ontology_id: UUID the ontology you are publishing to	
        :param name: str who is pushing this data	
        :param argo_work_dir: str Path to the template's work directory (i.e. nightshade/work/cbc)	
        :param experiment_name: str Name of the experiment to upload	
        :param class_map_path: str	
        :return: Error if it fails or True otherwise	
        """
        if model_id is None:
            model_id = str(uuid4())

        M = ArgoMetrics(str(argo_work_dir), experiment_name, class_map_path)
        endpoint = f'/yggy/ontology/{ontology_id}/trainings/upload-metrics'
        params = {
            "training-id": f"{model_id}",
            "user-firstname": f"{name}",
            "user-lastname": "SDK",
        }

        res = self.http_client.make_authed_request(endpoint, HttpMethod.PUT, M.payload, params)

        if isinstance(res, DtlError):
            return DtlError("Could not publish request.", res)

        return True

    def coach(self, ontology_id: Union[UUID, str]) -> Union[DtlError, pd.DataFrame]:
        """
        Receive machine learning coaching instructions based on the training data attached to an ontology
        For each class, see if the data attached is insufficient in volume or has been sourced from unvaried data, potentially leading to overfit
        :param ontology_id: id of the ontology on which to receive machine learning coaching
        :return: dataFrame containing coaching if successful, or DtlError of failed
        """

        ontology_client = _OntologyClient(self.http_client)
        uuid = _parse_uuid(ontology_id)
        if isinstance(uuid, DtlError):
            return uuid
        ontology = ontology_client.get(uuid)
        if isinstance(ontology, DtlError):
            return ontology
        return TrainingCoach.check(ontology)


class _NERDataClient:
    """
    Client to interact with TrainingData
    """

    def __init__(self, http_client: _HttpClient):
        self.http_client = http_client
        self.synthesis = _SynthesisClient(http_client)

    def delete_section(self, datastore_id: Union[str, UUID], path: List[str]) -> Union[DtlError, None]:
        """
        Deletes a section and all associated annotations by a given section_id
        :param section_id: id of the section to be deleted
        :return: Either a :class:`DtlError` if any error occurs or None if deletion is successful.
        """
        datastore_id = _parse_uuid(datastore_id, f"The supplied datastore_id: {datastore_id} is not a valid UUID")
        if isinstance(datastore_id, DtlError):
            return datastore_id
        res = self.http_client.make_authed_request(f'/themis/sections', HttpMethod.DELETE, body={
            "datastore_id": str(datastore_id),
            "path": path
        })
        if isinstance(res, DtlError):
            return res
        else:
            return None

    def get_annotations(self, ontology_id: Union[str, UUID],
                        datastore_ids: List[Union[str, UUID]],
                        items_per_page: int = 25,
                        page_after_id: Union[str, UUID, None] = None) -> Union[DtlError, Tuple[List[Section], Optional[UUID]]]:
        """
        Retrieve a list of annotated sections tagged from an ontology, with full text, all associated annotations of that section, and other statistics.
        :param ontology_id: id of the ontology to retrieve annotated sections from
        :param datastore_ids: can be used to retrieve only sections imported from specified datastores, if None, will return all annotated sections
        :param items_per_page: the number of AnnotatedSections to return
        :param page_after_id: the id of an AnnotatedSection in sequence, after which the returned page begins
        :return: a list of AnnotatedSections with value which should be used as page_after_id in next call if successful,
        or DtlError if failed

        """
        # validate ontology_id
        ontology_id = _parse_uuid(ontology_id, f"The supplied ontology_id: {ontology_id} is not a valid UUID")
        if isinstance(ontology_id, DtlError):
            return ontology_id

        # validate page_after_id
        if page_after_id is not None:
            page_after_id = _parse_uuid(
                ontology_id, f"The supplied page_after_id: {page_after_id} is not a valid UUID")
            if isinstance(page_after_id, DtlError):
                return page_after_id
        after_query = f'&after={page_after_id}' if page_after_id is not None else ''
        uri = f"/themis/ontology/{ontology_id}/export-annotations-json?size={items_per_page}{after_query}"

        # validate datastore_ids
        datastore_ids = list(map(lambda datastore_id: _parse_uuid(
            datastore_id, f"The supplied datastoreId: {datastore_id} is not a valid UUID"), datastore_ids))
        for datastore_id in datastore_ids:
            if isinstance(datastore_id, DtlError):
                return datastore_id

        body = list(map(lambda datstore_id: str(datstore_id), datastore_ids))
        res = self.http_client.execute_authed_request(uri, HttpMethod.POST, body)

        if isinstance(res, DtlError):
            return res
        sections_statistics = _parse_list(Section.from_payload)(res.json())
        if isinstance(sections_statistics, DtlError):
            return sections_statistics
        next_token_raw = res.headers.get("X-Next-Token", None)
        next_token = _parse_uuid(next_token_raw) if next_token_raw is not None else None
        if isinstance(next_token, DtlError):
            return next_token
        return sections_statistics, next_token


class _NERClient(_MLBase):
    """
    Client to interact with the NER
    """
    
    def __init__(self, http_client: _HttpClient):
        self.http_client = http_client
        self.data = _NERDataClient(http_client)

    def get_deployments(self,
                        page: int = 1,
                        item_per_page: int = 25) -> Union[DtlError, List[TrainingState]]:
        """
        Retrieve a list of all models currently deployed.
        :param page:
        :param item_per_page:
        :return:
        """
        from_item = (page-1)*item_per_page
        params = {
            'from': from_item,
            'limit': item_per_page,
            'status': 'Successful',
            'model-type': str(ModelType.Ner)
        }
        
        res = self.http_client.make_authed_request(
            f'/argus/deployments', HttpMethod.GET, params=params)
        if isinstance(res, DtlError):
            return res
        else:
            return _parse_list(Deployment._from_payload)(res)

    def deploy(self, model_id: UUID, ontology_id: UUID) -> Union[DtlError, UUID]:
        """
        Deploy :class:`Training` based on the given model_id and ontology_id
        :return: True if successful, else returns :class:`DtlError`
        """
        res = self.http_client.make_authed_request(
            f'/yggy/ontology/{str(ontology_id)}/trainings/{str(model_id)}/deploy',
            HttpMethod.POST, {})

        if isinstance(res, DtlError):
            return res
        else:
            deployment_id = UUID(res.get('deploymentId'))
            return deployment_id

    def train(self, ontology_id: UUID, training_data: TrainingData = TrainingData.Original, force_run: bool = True) -> Union[DtlError, UUID, List]:
        """
        Starts training for a given ontology_id

        :param ontology_id: id of the ontology to be used for training
        :param training_data: type of training data to be used to train the model: TrainingData.Original or TrainingData.Synthesized
        :param force_run: allows to run training even when there are less data than should be
        :return: Either a :class:`DtlError` if any error occurs or :class:`True` if training is requested successfully.
        """

        params = dict()
        params["model-type"] = ModelType.Ner
        params["data-origin"] = training_data.value
        params["force-train"] = force_run

        res = self.http_client.make_authed_request(
            f"/yggy/ontology/{ontology_id}/train", HttpMethod.POST, params=params)
        
        if isinstance(res, DtlError):
            return res
        if isinstance(res, List):
            # TODO we should handle response with given content: [{'TagsLessThanLimit': {'currentTotalTags': X, 'recommendedTags': Y}}]
            return res
        else:
            model_id = UUID(res.get('trainingId'))
            return model_id

    def get_models(self, ontology_id: UUID, status_type: Optional[TrainingStatusType] = None, sort: Optional[OrderList] = OrderList.desc, limit: Optional[int] = 20) -> Union[DtlError, List[TrainingState]]:
        """
        Get :class:`Trainings` based on the given ontology_id
        
        :param ontology_id:
        :return: List of trainings if successful, else returns :class:`DtlError`
        """

        params = {}

        if limit is not None:
            if not isinstance(limit, int) or limit < 0:
                return DtlError("limit should be a positive integer")
            params["limit"] = limit

        params["model-type"] = ModelType.Ner

        if status_type is not None:
            params["status-type"] = status_type.value

        if sort is not None:
            params["sort"] = sort.value

        res = self.http_client.make_authed_request(
            f'/yggy/ontology/{str(ontology_id)}/trainings', HttpMethod.GET, params=params)

        if isinstance(res, DtlError):
            return res
        else:
            return _parse_list(TrainingState._from_payload)(res)

class _DeploymentClient:
    """
    Client to interact with the Deployments
    """
    def __init__(self, http_client: _HttpClient):
        self.http_client = http_client

    def get(self, deployment_id: UUID):
        res = self.http_client.make_authed_request(f'/argus/deployments/{str(deployment_id)}', HttpMethod.GET)

        if isinstance(res, DtlError):
            return res

        return Deployment._from_payload(res)
