# Copyright 2024 Leadscope, Inc.; Instem LSS Ltd.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from . import LSEClient, StructureSet
from ._util import print_progress, DataFrameBuilder, check_task_error

import logging

class NitrosamineAssessment:
    """
    :class:`NitrosamineAssessment`
    ==============================

    The ``NitrosamineAssessment`` class is used to run N-nitrosamine potency assessments via the Leadscope Python Interface.

    The ``apply`` method requires a server session to be active, a session can be started using ``start_session``, \
    this must be subsequently ended with ``end_session``. In most cases it is preferable to use ``with_session`` \
    which starts and ends the session when done, including in the case of a thrown exception.

    Parameters
    ----------
    client : :class:`lserest.LSEClient`
        An :class:`LSEClient` object containing server connection details
    structures : :class:`lserest.StructureSet`
        A :class:`StructureSet` containing the structures to apply the models to
    exclude_images : `bool`
        Don't return images from the server (default: `True`)
    image_size : `int`
        The size (square) in pixels for images if returned (default: 50)
    progress_callback : `callable`
        A progress reporting function (optional: prints to console if not specified)
    
    Attributes
    ----------
    json_results : JSON
        Results for an assessment once it has been run, populated by ``apply``
    """

    def __init__(self, client:LSEClient, structures:StructureSet, exclude_images:bool=True, image_size:int=50):

        self._client = client
        self.structures = structures

        self.exclude_images = exclude_images
        self.image_size = image_size
        
        self.json_results = None
    
    # ---------------------------------------------------------------------------------------------------------------------------------
    # Public API methods:
    # ---------------------------------------------------------------------------------------------------------------------------------
    
    def apply(self, progress_callback=print_progress):
        """
        Run N-nitrosamine potency assessment. \
        This method requires an active session.

        This is a long task, the parameter `progress_callback` can be set to use a custom progress reporting function.

        This saves the returned JSON to `self.json_results`, results can also be obtained as a :class:`pandas.DataFrame` via `results()`.

        Warning: if `apply()` is re-run, the results are overwritten, not appended to.

        Parameters
        ----------
        progress_callback : callable
            A progress reporting function (optional: prints to console if not specified)
        """
        if self._client.session_id:
            self._client.submit_post('/nitrosamines/start', {
                'molfiles': self.structures.mol_strings,
                'excludeImages': self.exclude_images,
                'imageSize': self.image_size
            })

            state = self._client.wait_for_task(progress_callback)
            if state['complete']: self.json_results = self._client.submit_get('/nitrosamines/results')
            else: check_task_error(state)
        else:
            logging.error("The method NitrosamineAssessment.apply requires a session to be opened (see LSEClient.with_session).")

    def results(self):
        """
        Gets results for the N-nitrosamine potency assessment.

        Returns
        -------
        A :class:`pandas.DataFrame` of consensus results
        """
        data = DataFrameBuilder(len(self.json_results))
        for i, app_result in enumerate(self.json_results):
            data.add_name(i, app_result['name'])
            data.add_result(i, app_result['validNitrosamine'], "Valid Nitrosamine?")
            try: data.add_result(i, int(app_result['aiValue']), "AI Value")
            except KeyError: pass
            
        return data.df()
    
    # ---------------------------------------------------------------------------------------------------------------------------------
    # Internal methods
    # ---------------------------------------------------------------------------------------------------------------------------------
    
    def _apply(self, progress_callback, client:LSEClient):
        client.submit_post('/nitrosamines/start', {
            'molfiles': self.structures.mol_strings,
            'excludeImages': self.exclude_images,
            'imageSize': self.image_size
        })

        state = client.wait_for_task(progress_callback)
        if state['complete']: self.json_results = client.submit_get('/nitrosamines/results')
        else: check_task_error(state)