"""Predictions - Collection of methods to store and manipulate SCO model runs
and their outputs (predictions).
"""

import datetime
import os
import uuid

import attribute
import datastore


# ------------------------------------------------------------------------------
#
# Constants
#
# ------------------------------------------------------------------------------

# Timestamp of run creation
RUN_CREATED = 'createdAt'
# Timestamp of run start
RUN_STARTED = 'startedAt'
# Timestamp of run end
RUN_FINISHED = 'finishedAt'

# Run states
STATE_FAILED = 'FAILED'
STATE_IDLE = 'IDLE'
STATE_RUNNING = 'RUNNING'
STATE_SUCCESS = 'SUCCESS'


# ------------------------------------------------------------------------------
#
# Model Run State Objects
#
# ------------------------------------------------------------------------------

class ModelRunState(object):
    """Object containing information about the state and potential results or
    error messages generated by a predictive model run. This is considered an
    abstract class that is extended by three sub-classes for states: RUNNING,
    FAILED, and SUCCESS.
    """
    @staticmethod
    def from_json(json_obj):
        # Decide on type of returned object based on fact that failed states
        # have errors while success states have model output
        if json_obj['type'] == STATE_FAILED and 'errors' in json_obj:
            return ModelRunFailed(json_obj['errors'])
        elif json_obj['type'] == STATE_SUCCESS and 'modelOutput' in json_obj:
            return ModelRunSuccess(json_obj['modelOutput'])
        elif json_obj['type'] == STATE_IDLE:
            return ModelRunIdle()
        elif json_obj['type'] == STATE_RUNNING:
            return ModelRunActive()

    @property
    def is_failed(self):
        """Flag indicating if the model run has exited in a failed state.

        Returns
        -------
        Boolean
            True, if model run is in falied state.
        """
        return False

    @property
    def is_idle(self):
        """Flag indicating if the model run is waiting to start execution.

        Returns
        -------
        Boolean
            True, if model run is in idle state.
        """
        return False

    @property
    def is_running(self):
        """Flag indicating if the model run is in a running state.

        Returns
        -------
        Boolean
            True, if model run is in running state.
        """
        return False

    @property
    def is_success(self):
        """Flag indicating if the model run has finished with success.

        Returns
        -------
        Boolean
            True, if model run is in success state.
        """
        return False

    @staticmethod
    def to_json(obj):
        """Generate a JSON serialization for the run state object.

        Returns
        -------
        Json-like object
            Json serialization of model run state object
        """
        # Have text description of state in Json object (for readability)
        json_obj = {'type' : repr(obj)}
        # Add state-specific elements
        if obj.is_failed:
            json_obj['errors'] = obj.errors
        elif obj.is_success:
            json_obj['modelOutput'] = obj.model_output
        return json_obj


class ModelRunActive(ModelRunState):
    """Object indicating an active model run."""
    def __repr__(self):
        """String representation of the run state object."""
        return STATE_RUNNING

    @property
    def is_running(self):
        """Override is_running flag to indicate that this object represents a
        failed model run.
        """
        return True


class ModelRunFailed(ModelRunState):
    """Object indicating a failed model run. Contains a list of error messages
    that may have been generated as result of an exception during model run
    execution.

    Attributes
    ----------
    errors : list(string), optional
        List of error messages
    """
    def __init__(self, errors=[]):
        """Initialize list of errors. Set as an empty list if no error messages
        are given.

        Parameters
        ----------
        errors : list(string), optional
            List of error messages
        """
        self.errors = errors

    def __repr__(self):
        """String representation of the run state object."""
        return STATE_FAILED

    @property
    def is_failed(self):
        """Override is_failed flag to indicate that this object represents a
        failed model run.
        """
        return True


class ModelRunIdle(ModelRunState):
    """Object indicating an idle model run."""
    def __repr__(self):
        """String representation of the run state object."""
        return STATE_IDLE

    @property
    def is_idle(self):
        """Override is_idle flag to indicate that this object represents an
        idle model run.
        """
        return True


class ModelRunSuccess(ModelRunState):
    """Object indicating a succesful completed model run. Contains the result
    of the model run in form of a functional data object (reference).

    Attributes
    ----------
    model_output : string
        Unique identifier of functional data object containing the model run
        output
    """
    def __init__(self, model_output):
        """Initialize reference to model output object.

        Parameters
        ----------
        model_output : string
            Unique identifier of functional data object containing the model run
            output
        """
        self.model_output = model_output

    def __repr__(self):
        """String representation of the run state object."""
        return STATE_SUCCESS

    @property
    def is_success(self):
        """Override is_success flag to indicate that this object represents a
        failed model run.
        """
        return True


# ------------------------------------------------------------------------------
#
# Database Objects
#
# ------------------------------------------------------------------------------


class ModelRunHandle(datastore.ObjectHandle):
    """Handle to access and manipulate an object representing a model run and
    its state information.

    The status of the model run is maintained as a separate object. A run that
    has completed successfully will have a prediction result associated
    with its state. In case of failure, there will be a list of error
    messages associated with its state.

    The state information is replicated into the properties list to allow
    object listing filters based on run state.

    Attributes
    ----------
    arguments: Dictionary(attribute.Attribute)
        Dictionary of typed attributes defining the image group options
    experiment : string
        Unique experiment object identifier
    schedule : Dictionary(string)
        Timestamps for model run state changes
    state: ModelRunState
        Model run state object
    """
    def __init__(
        self,
        identifier,
        properties,
        state,
        experiment,
        arguments,
        schedule=None,
        timestamp=None,
        is_active=True):
        """Initialize the subject handle.

        Parameters
        ----------
        identifier : string
            Unique object identifier
        properties : Dictionary
            Dictionary of experiment specific properties
        state : ModelRunState
            Model run state object
        experiment : string
            Unique experiment object identifier
        arguments: Dictionary(attribute.Attribute)
            Dictionary of typed attributes defining the model run arguments
        schedule : Dictionary(string), optional
            Timestamps for model run state changes. Only optinal if timestamp is
            missing as well.
        timestamp : datetime, optional
            Time stamp of object creation (UTC).
        is_active : Boolean, optional
            Flag indicating whether the object is active or has been deleted.
        """
        # Initialize super class
        super(ModelRunHandle, self).__init__(
            identifier,
            timestamp,
            properties,
            is_active=is_active
        )
        # Initialize class specific Attributes
        self.state = state
        self.experiment = experiment
        self.arguments = arguments
        # Set state change information. Only allowed to be missing at run
        # creation, i.e., if timestamp is none.
        if schedule is None:
            if not timestamp is None:
                raise ValueError('missing schedule information')
            self.schedule = {RUN_CREATED : str(self.timestamp.isoformat())}
        else:
            self.schedule = schedule

    @property
    def is_model_run(self):
        """Override the is_experiment property of the base class."""
        return True


# ------------------------------------------------------------------------------
#
# Object Stores
#
# ------------------------------------------------------------------------------

class DefaultModelRunManager(datastore.MongoDBStore):
    """Manager for model runs and their outputs.

    This is a default implentation that uses MongoDB as storage backend.

    Attributes
    ----------
    parameters : list(attribute.AttributeDefinition)
        List of parameters that can be passed to the model when running it.
    """
    def __init__(self, mongo_collection):
        """Initialize the MongoDB collection and base directory where to store
        model runs and ouput files. Set immutable and mandatory properties.

        Parameters
        ----------
        mongo_collection : Collection
            Collection in MongoDB storing model run information
        """
        # Initialize the super class
        super(DefaultModelRunManager, self).__init__(
            mongo_collection,
            [datastore.PROPERTY_STATE])
        # Initialize the definition of image group options attributes
        self.parameters = {
            'gabor_orientations' : attribute.AttributeDefinition(
                'gabor_orientations',
                attribute.IntegerType(),
                default_value=8
            ),
            'max_eccentricity' : attribute.AttributeDefinition(
                'max_eccentricity',
                attribute.FloatType(),
                default_value=12
            ),
            'normalized_pixels_per_degree' : attribute.AttributeDefinition(
                'normalized_pixels_per_degree',
                attribute.FloatType()
            )
        }

    def create_object(self, name, experiment, arguments=None, properties=None):
        """Create a model run object with the given list of arguments. The
        initial state of the object is RUNNING.

        Parameters
        ----------
        name : string
            User-provided name for the model run
        experiment : string
            Unique identifier of associated experiment object
        arguments : list(attribute.Attribute), optional
            List of attribute instances
        properties : Dictionary, optional
            Set of model run properties.
        Returns
        -------
        PredictionHandle
            Object handle for created model run
        """
        # Create a new object identifier.
        identifier = str(uuid.uuid4())
        # By default all model runs are in IDLE state at creation
        state = ModelRunIdle()
        # Create the initial set of properties.
        run_properties = {
            datastore.PROPERTY_NAME: name,
            datastore.PROPERTY_STATE: str(state)
        }
        if not properties is None:
            for prop in properties:
                if not prop in run_properties:
                    run_properties[prop] = properties[prop]
        # If argument list is not given then the initial set of arguments is
        # given by those parameter for which a default value is defined. Use
        # default values for those arguments in the definition that are not
        # present in a given arguments list.
        run_arguments = attribute.get_default_attributes(self.parameters)
        if not arguments is None:
            args = attribute.get_and_validate_attributes(arguments, self.parameters)
            for arg_key in args:
                run_arguments[arg_key] = args[arg_key]
        # Create the image group object and store it in the database before
        # returning it.
        obj = ModelRunHandle(
            identifier,
            run_properties,
            state,
            experiment,
            run_arguments
        )
        self.insert_object(obj)
        return obj

    def from_json(self, document):
        """Create model run object from JSON document retrieved from database.

        Parameters
        ----------
        document : JSON
            Json document in database

        Returns
        -------
        PredictionHandle
            Handle for model run object
        """
        # Get object identifier from Json document
        identifier = str(document['_id'])
        # Create model run handle.
        return ModelRunHandle(
            identifier,
            document['properties'],
            ModelRunState.from_json(document['state']),
            document['experiment'],
            attribute.attributes_from_json(document['arguments']),
            schedule=document['schedule'],
            timestamp=datetime.datetime.strptime(
                document['timestamp'], '%Y-%m-%dT%H:%M:%S.%f'
            ),
            is_active=document['active']
        )

    def to_json(self, model_run):
        """Create a Json-like dictionary for a model run object. Extends the
        basic object with run state, arguments, and optional prediction results
        or error descriptions.

        Parameters
        ----------
        model_run : PredictionHandle

        Returns
        -------
        (JSON)
            Json-like object, i.e., dictionary.
        """
        # Get the basic Json object from the super class
        json_obj = super(DefaultModelRunManager, self).to_json(model_run)
        # Add run state
        json_obj['state'] = ModelRunState.to_json(model_run.state)
        # Add run scheduling Timestamps
        json_obj['schedule'] = model_run.schedule
        # Add experiment information
        json_obj['experiment'] = model_run.experiment
        # Transform dictionary of attributes into list of key-value pairs.
        json_obj['arguments'] = attribute.attributes_to_json(model_run.arguments)
        return json_obj

    def update_state(self, identifier, state):
        """Update state of identified model run.

        Raises exception if state change results in invalid run life cycle.

        Parameters
        ----------
        identifier : string
            Unique model run identifier
        state : ModelRunState
            Object representing new run state

        Returns
        -------
        ModelRunHandle
            Modified model run handle or None if no run with given identifier
            exists
        """
        # Get model run to ensure that it exists
        model_run = self.get_object(identifier)
        if model_run is None:
            return None
        # Set timestamp of state change. Raise exception if state change results
        # in invalid life cycle
        timestamp = str(datetime.datetime.utcnow().isoformat())
        if state.is_idle:
            raise ValueError('invalid state change: run cannot become idle')
        elif state.is_running:
            # Current state is required to be IDLE
            if not model_run.state.is_idle:
                raise ValueError('invalid state change: finished run cannot start again')
            model_run.schedule[RUN_STARTED] = timestamp
        elif state.is_failed or state.is_success:
            # Current state is required to be RUNNING
            if not model_run.state.is_running:
                raise ValueError('invalid state change: cannot finish inactive run')
            model_run.schedule[RUN_FINISHED] = timestamp
        # Update model run state and replace object in database
        model_run.state = state
        model_run.properties[datastore.PROPERTY_STATE] = str(state)
        self.replace_object(model_run)
        # Return modified model run
        return model_run
