# Copyright 2015-2018 Thomas W. D. Möbius
#
# This file is part of fmristats.
#
# fmristats is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 3 of the License, or (at your
# option) any later version.
#
# fmristats is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# It is not allowed to remove this copy right statement.

"""

Irritations, paradigms.

"""

from .name import Identifier

import numpy as np

import pickle

class Irritation:
    """
    Irritation design

    Parameters
    ----------
    name : Identifier
        identifier of the session
    """

    def __init__(self, name):
        assert type(name) is Identifier, 'name must be of type Identifier'
        self.name = name

    def design(self, slice_time):
        return np.zeros_like(slice_time)

    def describe(self):
        """
        Give a description of the instance
        """
        return self.name.describe()

    def save(self, file, **kwargs):
        """
        Save irritation instance to disk

        This will save the irritation instance to disk for later use.

        Parameters
        ----------
        file : str
            File name.
        """
        with open(file, 'wb') as output:
            pickle.dump(self, output, **kwargs)

class Block(Irritation):
    """
    Block irritation of the subject

    Parameters
    ----------
    name : Identifier
        identifier of the session
    names : list
    onsets : dict
    durations : dict
    """

    def __init__(self, name, names, onsets, durations):
        assert type(name) is Identifier, 'name must be of type Identifier'

        number = len(names)

        assert len(names) == number, 'number of names does not match number of blocks'
        assert len(durations) == number, 'number of durations does not match number of blocks'
        assert len(onsets) == number, 'number of onsets does not match number of blocks'

        self.name = name
        self.number = number
        self.names = names
        self.durations = durations
        self.onsets = onsets

#    def design(self, slice_time, s, c, offset=0, preset=0):
#        """
#        Defines the irritation design vector
#
#        Parameters
#        ----------
#        slice_time : ndarray, shape (n,) or (n,m)
#            times at which slices have been measured during the
#            experiment.  First dimension must be the same as the number
#            of full scans
#        s : str
#            stimulus block of interest
#        c : str
#            name of the control block to be used
#        offset : float
#            offset to apply at the beginning of an irritation phase
#        preset : float
#            offset to apply at the end of an irritation phase
#
#        Returns
#        -------
#        ndarray, shape_like(slice_time), dtype: float
#            The irritation vector indicating the kind of irritation under
#            which the subject was exposed while the particular slice has
#            been measured.
#        """
#        assert s != c, 'block names: {} and {c} must be different'.format(s,c)
#        assert s in self.names, 'name of stimulus block: {} does not exist'.format(s)
#        assert c in self.names, 'name of control block: {} does not exist'.format(c)
#
#        design = np.empty_like(slice_time)
#        design[...] = np.nan
#
#        onsets = self.onsets[c]
#        durations = self.durations[c]
#        srt = onsets + offset
#        end = onsets + durations - preset
#        mat = np.vstack((srt, end)).T
#
#        for x in iter(mat):
#            design[ (slice_time > x[0]) & (slice_time < x[1]) ] = 0
#
#        onsets = self.onsets[s]
#        durations = self.durations[s]
#        srt = onsets + offset
#        end = onsets + durations - preset
#        mat = np.vstack((srt, end)).T
#
#        for x in iter(mat):
#            design[ (slice_time > x[0]) & (slice_time < x[1]) ] = 1
#
#        return design

    def design(self, slice_timing, s, c, offset=0, preset=0):
        """
        Defines the irritation design vector

        Parameters
        ----------
        slice_timing : ndarray, shape (n,) or (n,m)
            times at which slices have been measured during the
            experiment.  First dimension must be the same as the number
            of full scans
        s : str
            stimulus block of interest
        c : str
            name of the control block to be used
        offset : float
            offset to apply at the beginning of an irritation phase
        preset : float
            offset to apply at the end of an irritation phase

        Returns
        -------
        ndarray, shape_like(slice_timing), dtype: float
            The irritation vector indicating the kind of irritation under
            which the subject was exposed while the particular slice has
            been measured.
        """
        assert s != c, 'block names: {} and {c} must be different'.format(s,c)
        assert s in self.names, 'name of stimulus block: {} does not exist'.format(s)
        assert c in self.names, 'name of control block: {} does not exist'.format(c)

        design = np.empty(slice_timing.shape + (2,) )
        design[...] = np.nan

        onsets = self.onsets[c]
        durations = self.durations[c]
        srt = onsets + offset
        end = onsets + durations - preset
        mat = np.vstack((srt, end)).T

        group = 0
        for x in iter(mat):
            block = (slice_timing > x[0]) & (slice_timing < x[1])
            design[block,0] = 0
            design[block,1] = group
            group += 1

        onsets = self.onsets[s]
        durations = self.durations[s]
        srt = onsets + offset
        end = onsets + durations - preset
        mat = np.vstack((srt, end)).T

        group = 0
        for x in iter(mat):
            block = (slice_timing > x[0]) & (slice_timing < x[1])
            design[block,0] = 1
            design[block,1] = group
            group += 1

        return design

    def number_of_onsets(self):
        return {name : len(self.onsets[name]) for name in self.names}

    def irritated_scans(self, slice_time, offset=0, preset=0):
        """
        Test whether blocks have been aquired during blocks of irritation

        Parameters
        ----------
        slice_time : ndarray, shape (n,) or (n,m), dtype: float

        Returns
        -------
        ndarray, shape_like(slice_time), dtype: bool
            If True, this scan has been aquired during a block of subject
            irritation.
        """
        onsets = self.onsets.items()

        srt = np.hstack([v + offset for k,v in onsets])
        end = np.hstack([v + irr.durations[k] - preset for k,v in onsets])
        mat = np.vstack((srt, end)).T

        irritated = np.zeros_like(slice_time).astype(bool)
        for x in iter(mat):
            irritated = irritated | (slice_time > x[0]) & (slice_time < x[1])

        return irritated

    def describe(self):
        """
        Give a description of the instance
        """
        description = """
        Type of irritation: block design
        Block number: {}
        Block names:  {}
        Onsets: {}
        Durations: {}
        Number of onsets per block: {}"""
        return description.format(
                self.number,
                self.names,
                self.durations,
                self.onsets,
                self.number_of_onsets(),
                )

