"""Reward analysis functions."""
from __future__ import absolute_import

import numpy as np
import pandas as pd
import itertools as it
from collections import defaultdict
import warnings

from scipy.ndimage.morphology import binary_dilation

from . import behavior_analysis as ba
from . import new_intervals as ints
from ..misc.analysis_helpers import rewards_by_condition
from ..classes.classes import ExperimentGroup
from ..classes import exceptions as exc


def lick_to_reward_distance(expt_grp, rewardPositions=None):
    """Calculate the average lick to reward distance.

    Parameters
    ----------
    rewardPositions : {str, None, np.ndarray}
        If a string, assumed to be a condition label, and will use the
        reward positions used for each mouse during the condition.
        If 'None', uses the actual reward positions during the experiment.
        Otherwise pass in normalized reward positions.

    Returns
    -------
    pd.DataFrame

    """
    result = []

    if rewardPositions is None:
        rewards_by_expt = {
            expt: expt.rewardPositions(units='normalized')
            for expt in expt_grp}
    elif isinstance(rewardPositions, basestring):
        rewards_by_expt = rewards_by_condition(
            expt_grp, rewardPositions, condition_column='condition')
    else:
        rewards_by_expt = defaultdict(lambda: np.array(rewardPositions))

    for expt in expt_grp:

        rewards = rewards_by_expt[expt]

        for trial in expt.findall('trial'):
            bd = trial.behaviorData(imageSync=False)
            position = ba.absolutePosition(
                trial, imageSync=False, sampling_interval='actual')

            if np.any(rewards >= 1.0):
                trial_rewards = rewards / bd['trackLength']
            else:
                trial_rewards = rewards

            licking = bd['licking'][:, 0]
            licking = licking[np.isfinite(licking)]
            licking = licking / bd['samplingInterval']
            licking = licking.astype('int')

            licking_positions = position[licking] % 1

            # meshgrid sets up the subtraction below
            # basically tile expands the arrays
            rewards_mesh, licking_mesh = np.meshgrid(
                trial_rewards, licking_positions)

            reward_distance = licking_mesh - rewards_mesh
            # All distances should be on [-0.5, 0.5)
            reward_distance[reward_distance >= 0.5] -= 1.0
            reward_distance[reward_distance < -0.5] += 1.0

            reward_distance = np.amin(np.abs(reward_distance), axis=1)

            assert len(licking_positions) == len(reward_distance)
            for lick, position in it.izip(
                    reward_distance, licking_positions):
                result.append(
                    {'trial': trial, 'position': position, 'value': lick})
    return pd.DataFrame(result, columns=['trial', 'position', 'value'])


def fraction_rewarded_lick_intervals(expt_grp, **lick_interval_kwargs):
    """Fraction of lick intervals that were rewarded.

    Parameters
    ----------
    **lick_interval_kwargs : dict
        All additional keyword parameters are passed to
        ba.calculateRewardedLickIntervals.

    Returns
    -------
    pd.DataFrame

    """
    result = []
    for expt in expt_grp:
        rewarded_intervals, unrewarded_intervals \
            = ba.calculateRewardedLickIntervals(expt, **lick_interval_kwargs)
        rewarded_count = sum(
            [intervals.shape[0] for intervals in rewarded_intervals])
        unrewarded_count = sum(
            [intervals.shape[0] for intervals in unrewarded_intervals])

        try:
            fraction = rewarded_count \
                / float(rewarded_count + unrewarded_count)
        except ZeroDivisionError:
            fraction = np.nan
        result.append({
            'expt': expt, 'rewarded_count': rewarded_count,
            'unrewarded_counts': unrewarded_count, 'value': fraction})
    return pd.DataFrame(result, columns=[
        'expt', 'rewarded_count', 'unrewarded_count', 'value'])


def fraction_licks_in_rewarded_intervals(expt_grp, **lick_interval_kwargs):
    """Fraction of licks that are within a rewarded lick interval.

    Parameters
    ----------
    **lick_interval_kwargs : dict
        All additional keyword parameters are passed to
        ba.calculateRewardedLickIntervals.

    Returns
    -------
    pd.DataFrame

    """
    result = []
    for expt in expt_grp:
        rewarded_intervals, _ = ba.calculateRewardedLickIntervals(
            expt, **lick_interval_kwargs)

        licks = []
        for trial in expt.findall('trial'):
            sampling_interval = trial.behavior_sampling_interval()
            licks.append(
                trial.behaviorData(imageSync=False)['licking'][:, 0] /
                float(sampling_interval))

        rewarded_licks = np.zeros(len(rewarded_intervals))
        total_licks = 0

        for trial_idx, intervals_trial, licks_trial in it.izip(
                it.count(), rewarded_intervals, licks):
            total_licks += licks_trial.shape[0]
            for interval in intervals_trial:
                rewarded_licks[trial_idx] += np.sum(
                    (interval[0] <= licks_trial) &
                    (interval[1] >= licks_trial))

        total_rewarded_licks = np.sum(rewarded_licks)

        try:
            fraction = total_rewarded_licks / float(total_licks)
        except ZeroDivisionError:
            fraction = np.nan
        result.append({
            'expt': expt, 'rewarded_licks': total_rewarded_licks,
            'total_licks': total_licks, 'value': fraction})

    return pd.DataFrame(result, columns=[
        'expt', 'rewarded_licks', 'total_licks', 'value'])


def fraction_licks_rewarded(expt_grp):
    """Fraction of possible licks rewarded.

    Takes in to account the operant reward rate.

    Returns
    -------
    pd.DataFrame

    """
    result = []
    for expt in expt_grp:

        totalLicks = sum([
            trial.behaviorData()['licking'].shape[0]
            for trial in expt.findall('trial')])

        totalWater = sum([
            trial.behaviorData()['water'].shape[0]
            for trial in expt.findall('trial')])

        rewardRate = expt.reward_parameters().get('operant_rate', 1)

        try:
            fraction = float(totalWater) / (totalLicks / float(rewardRate))
        except ZeroDivisionError:
            fraction = np.nan

        result.append({
            'expt': expt, 'lick': totalLicks, 'water': totalWater,
            'value': fraction})
    return pd.DataFrame(result, columns=['expt', 'lick', 'water', 'value'])


def fractionLicksNearRewardsPerLap(
        expGrp, anticipStartCM=-5, anticipEndCM=-0.1, compareStartCM=-15,
        compareEndCM=-0.1, fractionColName="value", rewardPositions=None,
        exclude_reward=False, exclude_reward_duration=10.0):
    """Fraction of licks in the anticipatory zone vs a compare zone, per lap.

    Parameters
    ----------
    anticipStartCM, anticipEndCM : float
        licks in this spatial window is counted anticipatory.
        Units are in cm. The reward zone start is considered as 0.
        Prereward space is negative, and post reward space is positive.
    copareStartCM, compareEndCM : float
        licks in this window is counted toward total licks.
        Usually, this window should contain the anticipatory window.
    fractionColName : str
        the name to give to the lick fraction column of the returned
        dataframe,  default is "value"
    rewardPositions : {str, None, np.ndarray}
        If a string, assumed to be a condition label, and will use the
        reward positions used for each mouse during the condition.
        If 'None', uses the actual reward positions during the experiment.
        Otherwise pass in normalized reward positions.
    exclude_reward : bool
        exclude licks that occurs after water rewards. This is trying to
        ignore licks that are for drinking water. Default is false.
    exclude_duration : float
        number of seconds after the onset of water reward for which the
        licking should be ignored. Default is 10 seconds.

    Returns
    -------
    pd.DataFrame
        Each row is the licking calculation of a lap, with columns:
            trial - trial instance
            rewardPosition - reward location
            lapNum - the lap number
            anticipLicks - number of licks in the anticipatory zone
            compareLicks - number of licks in the compare zone
            value - anticipLicks/compareLicks.
                the name of this column is set by the kwarg fractionColName,
                when the compare zone contains the anticipatory zone, it is
                similar to fraction of licks near rewards
    """
    result = []

    if(rewardPositions is None):
        rewards_by_exp = {exp: exp.rewardPositions(units='normalized')
                          for exp in expGrp}
    elif(isinstance(rewardPositions, basestring)):
        rewards_by_exp = rewards_by_condition(
            expGrp, rewardPositions, condition_column="condition")
    else:
        rewards_by_exp = defaultdict(lambda: np.array(rewardPositions))

    for exp in expGrp:
        try:
            belt_length = exp.belt().length()
        except exc.NoBeltInfo:
            warnings.warn('Missing belt length. All results may be wrong.')
            belt_length = 220
        anticipStart = anticipStartCM / float(belt_length)
        anticipEnd = anticipEndCM / float(belt_length)
        compareStart = compareStartCM / float(belt_length)
        compareEnd = compareEndCM / float(belt_length)

        rewards = rewards_by_exp[exp]

        for trial in exp.findall("trial"):
            position = ba.absolutePosition(
                trial, imageSync=False, sampling_interval="actual")
            bd = trial.behaviorData(
                imageSync=False, sampling_interval="actual")

            lapNum = position.astype("int32")

            for reward in rewards:
                for i in np.r_[0:np.max(lapNum)]:
                    absRewardPos = i + reward
                    anticipS = absRewardPos + anticipStart
                    anticipE = absRewardPos + anticipEnd
                    compareS = absRewardPos + compareStart
                    compareE = absRewardPos + compareEnd

                    anticipBA = (position >= anticipS) & (position < anticipE)
                    compareBA = (position >= compareS) & (position < compareE)

                    if exclude_reward:
                        numExcPoints = np.int(np.float(
                            exclude_reward_duration)/bd["samplingInterval"])
                        try:
                            firstWater = np.where(compareBA & bd["water"])[0]
                            if(firstWater.size > 0):
                                firstWater = firstWater[0]
                                compareBA[firstWater:firstWater +
                                          numExcPoints] = False
                                anticipBA[firstWater:firstWater +
                                          numExcPoints] = False
                        except KeyError:
                            print("No water signal exist, " +
                                  "exclude_reward not in effect")

                    numAnticipLicks = np.sum(bd["licking"][anticipBA])
                    numCompareLicks = np.sum(bd["licking"][compareBA])

                    fraction = numAnticipLicks / float(numCompareLicks)

                    result.append({"trial": trial,
                                   "rewardPos": reward,
                                   "lapNum": i,
                                   "anticipLicks": numAnticipLicks,
                                   "compareLicks": numCompareLicks,
                                   fractionColName: fraction})

    return pd.DataFrame(result, columns=[
        'trial', 'rewardPos', 'lapNum', 'anticipLicks', 'compareLicks'])


def fractionLicksNearRewards(expt_grp, **kwargs):
    """Aggregate fractionLicksNearRewardsPerLap by Trial."""
    fraction_per_lap = fractionLicksNearRewardsPerLap(expt_grp, **kwargs)
    result = fraction_per_lap.groupby('trial', as_index=False).apply(