import datetime

import numpy as np
import pandas as pd
from six import iteritems
from six.moves import zip

from zipline.utils.numpy_utils import categorical_dtype, NaTns


def next_event_frame(events_by_sid,
                     dates,
                     missing_value,
                     field_dtype,
                     event_date_field_name,
                     return_field_name):
    """
    Make a DataFrame representing the simulated next known dates or values
    for an event.

    Parameters
    ----------
    dates : pd.DatetimeIndex.
        The index of the returned DataFrame.
    events_by_sid : dict[int -> pd.Series]
        Dict mapping sids to a series of dates. Each k:v pair of the series
        represents the date we learned of the event mapping to the date the
        event will occur.
    event_date_field_name : str
        The name of the date field that marks when the event occurred.

    Returns
    -------
    next_events: pd.DataFrame
        A DataFrame where each column is a security from `events_by_sid` where
        the values are the dates of the next known event with the knowledge we
        had on the date of the index. Entries falling after the last date will
        have `NaT` as the result in the output.


    See Also
    --------
    previous_date_frame
    """
    date_cols = {
        equity: np.full_like(dates, NaTns) for equity in events_by_sid
    }
    value_cols = {
        equity: np.full(len(dates), missing_value, dtype=field_dtype)
        for equity in events_by_sid
    }

    raw_dates = dates.values
    for equity, df in iteritems(events_by_sid):
        event_dates = df[event_date_field_name]
        values = df[return_field_name]
        data = date_cols[equity]
        if not event_dates.index.is_monotonic_increasing:
            event_dates = event_dates.sort_index()

        # Iterate over the raw Series values, since we're comparing against
        # numpy arrays anyway.
        iter_date_vals = zip(event_dates.index.values, event_dates.values,
                             values)
        for knowledge_date, event_date, value in iter_date_vals:
            date_mask = (
                (knowledge_date <= raw_dates) &
                (raw_dates <= event_date)
            )
            value_mask = (event_date <= data) | (data == NaTns)
            data_indices = np.where(date_mask & value_mask)
            data[data_indices] = event_date
            value_cols[equity][data_indices] = value
    return pd.DataFrame(index=dates, data=value_cols)


def previous_event_frame(events_by_sid,
                         date_index,
                         missing_value,
                         field_dtype,
                         event_date_field,
                         previous_return_field):
    """
    Make a DataFrame representing simulated previous dates or values for an
    event.

    Parameters
    ----------
    events_by_sid : dict[int -> DatetimeIndex]
        Dict mapping sids to a series of dates. Each k:v pair of the series
        represents the date we learned of the event mapping to the date the
        event will occur.
    date_index : DatetimeIndex.
        The index of the returned DataFrame.
    missing_value : any
        Data which missing values should be filled with.
    field_dtype: any
        The dtype of the field for which the previous values are being
        retrieved.
    event_date_field: str
        The name of the date field that marks when the event occurred.
    return_field: str
        The name of the field for which the previous values are being
        retrieved.

    Returns
    -------
    previous_events: pd.DataFrame
        A DataFrame where each column is a security from `events_by_sid` and
        the values are the values for the previous event that occurred on the
        date of the index. Entries falling before the first date will have
        `missing_value` filled in as the result in the output.

    See Also
    --------
    next_date_frame
    """
    sids = list(events_by_sid)
    populate_value = None if field_dtype == categorical_dtype else \
        missing_value
    out = np.full(
        (len(date_index), len(sids)),
        populate_value,
        dtype=field_dtype
    )
    d_n = date_index[-1].asm8
    for col_idx, sid in enumerate(sids):
        # events_by_sid[sid] is a DataFrame mapping knowledge_date to event
        # date and values.
        df = events_by_sid[sid]
        df = df[df[event_date_field] <= d_n]
        event_date_vals = df[event_date_field].values
        # Get knowledge dates corresponding to the values in which we are
        # interested
        kd_vals = df[df[event_date_field] <= d_n].index.values
        # The date at which a previous event is first known is the max of the
        #  kd and the event date.
        index_dates = np.maximum(kd_vals, event_date_vals)
        out[
            date_index.searchsorted(index_dates), col_idx
        ] = df[previous_return_field]

    frame = pd.DataFrame(out, index=date_index, columns=sids)
    frame.ffill(inplace=True)
    if field_dtype == categorical_dtype:
        frame[frame.isnull()] = missing_value
    return frame


def normalize_data_query_time(dt, time, tz):
    """Apply the correct time and timezone to a date.

    Parameters
    ----------
    dt : pd.Timestamp
        The original datetime that represents the date.
    time : datetime.time
        The time of day to use as the cutoff point for new data. Data points
        that you learn about after this time will become available to your
        algorithm on the next trading day.
    tz : tzinfo
        The timezone to normalize your dates to before comparing against
        `time`.

    Returns
    -------
    query_dt : pd.Timestamp
        The timestamp with the correct time and date in utc.
    """
    # merge the correct date with the time in the given timezone then convert
    # back to utc
    return pd.Timestamp(
        datetime.datetime.combine(dt.date(), time),
        tz=tz,
    ).tz_convert('utc')


def normalize_data_query_bounds(lower, upper, time, tz):
    """Adjust the first and last dates in the requested datetime index based on
    the provided query time and tz.

    lower : pd.Timestamp