"""Dynamic Time Warping utilities for event alignment."""

from typing import Optional, Tuple, Union, List

import numpy as np
from scipy.interpolate import interp1d

try:
    from fastdtw import fastdtw
    from scipy.spatial.distance import euclidean
    HAS_FASTDTW = True
except ImportError:
    HAS_FASTDTW = False


def dtw_distance(
    x: np.ndarray,
    y: np.ndarray,
    radius: int = 1,
) -> Tuple[float, List[Tuple[int, int]]]:
    """
    Calculate Dynamic Time Warping distance between two sequences.

    Parameters
    ----------
    x : np.ndarray
        First sequence.
    y : np.ndarray
        Second sequence.
    radius : int, default 1
        Radius for fastdtw approximation. Higher values give more accurate
        results but are slower.

    Returns
    -------
    distance : float
        DTW distance between sequences.
    path : List[Tuple[int, int]]
        Warping path as list of (i, j) index pairs.

    Raises
    ------
    ImportError
        If fastdtw is not installed.

    Examples
    --------
    >>> x = np.array([1, 2, 3, 4, 5])
    >>> y = np.array([1, 2, 2, 4, 5])
    >>> distance, path = dtw_distance(x, y)
    """
    if not HAS_FASTDTW:
        raise ImportError(
            "fastdtw is required for DTW calculations. "
            "Install with: pip install fastdtw"
        )

    # Reshape to 2D if needed (fastdtw expects 2D arrays)
    x = np.atleast_2d(x).T if x.ndim == 1 else x
    y = np.atleast_2d(y).T if y.ndim == 1 else y

    distance, path = fastdtw(x, y, radius=radius, dist=euclidean)

    return distance, path


def align_sequences(
    source: np.ndarray,
    target_length: int,
    method: str = "dtw",
) -> np.ndarray:
    """
    Align a source sequence to a target length.

    This is useful for the "Copy" model when the historical event
    has a different duration than the target event.

    Parameters
    ----------
    source : np.ndarray
        Source sequence to align.
    target_length : int
        Desired length of output sequence.
    method : {'dtw', 'linear', 'nearest'}, default 'dtw'
        Alignment method:
        - 'dtw': Dynamic Time Warping (preserves shape)
        - 'linear': Linear interpolation
        - 'nearest': Nearest neighbor interpolation

    Returns
    -------
    aligned : np.ndarray
        Source sequence aligned to target length.

    Examples
    --------
    >>> source = np.array([100, 150, 200, 150, 100])  # 5-day event
    >>> aligned = align_sequences(source, target_length=7)  # Map to 7 days
    >>> print(len(aligned))
    7
    """
    source = np.asarray(source)
    source_length = len(source)

    if target_length == source_length:
        return source.copy()

    if method == "linear":
        return _interpolate_linear(source, target_length)

    elif method == "nearest":
        return _interpolate_nearest(source, target_length)

    elif method == "dtw":
        return _align_dtw(source, target_length)

    else:
        raise ValueError(f"Unknown method: {method}")


def _interpolate_linear(source: np.ndarray, target_length: int) -> np.ndarray:
    """Align using linear interpolation."""
    source_indices = np.linspace(0, 1, len(source))
    target_indices = np.linspace(0, 1, target_length)

    interpolator = interp1d(source_indices, source, kind="linear")
    return interpolator(target_indices)


def _interpolate_nearest(source: np.ndarray, target_length: int) -> np.ndarray:
    """Align using nearest neighbor interpolation."""
    source_length = len(source)

    # Map each target index to nearest source index
    target_indices = np.arange(target_length)
    source_positions = target_indices * (source_length - 1) / (target_length - 1)
    source_indices = np.round(source_positions).astype(int)
    source_indices = np.clip(source_indices, 0, source_length - 1)

    return source[source_indices]


def _align_dtw(source: np.ndarray, target_length: int) -> np.ndarray:
    """
    Align using DTW-based warping.

    Creates a reference sequence of desired length and finds the
    optimal warping path, then interpolates the source along that path.
    """
    if not HAS_FASTDTW:
        # Fall back to linear interpolation
        return _interpolate_linear(source, target_length)

    source_length = len(source)

    # Create reference sequence with target length
    # Use normalized indices as the reference
    reference = np.linspace(0, 1, target_length)
    source_normalized = np.linspace(0, 1, source_length)

    # Get DTW path
    _, path = dtw_distance(source_normalized, reference)

    # Build warped sequence
    # For each target position, find corresponding source position(s)
    target_to_source = {}
    for src_idx, tgt_idx in path:
        if tgt_idx not in target_to_source:
            target_to_source[tgt_idx] = []
        target_to_source[tgt_idx].append(src_idx)

    # Interpolate values
    warped = np.zeros(target_length)
    for tgt_idx in range(target_length):
        if tgt_idx in target_to_source:
            src_indices = target_to_source[tgt_idx]
            # Average the source values for this target position
            warped[tgt_idx] = np.mean([source[i] for i in src_indices])
        else:
            # Interpolate from neighbors
            prev_idx = max(0, tgt_idx - 1)
            next_idx = min(target_length - 1, tgt_idx + 1)
            warped[tgt_idx] = (warped[prev_idx] + warped[next_idx]) / 2

    return warped


def warp_sequence(
    source: np.ndarray,
    target: np.ndarray,
) -> Tuple[np.ndarray, List[Tuple[int, int]]]:
    """
    Warp source sequence to align with target sequence.

    Uses DTW to find optimal alignment, then interpolates source
    to match target.

    Parameters
    ----------
    source : np.ndarray
        Source sequence to warp.
    target : np.ndarray
        Target sequence to align to.

    Returns
    -------
    warped : np.ndarray
        Source sequence warped to align with target.
    path : List[Tuple[int, int]]
        DTW warping path.

    Examples
    --------
    >>> source = np.array([100, 200, 300, 200, 100])
    >>> target = np.array([90, 180, 270, 280, 190, 100, 50])
    >>> warped, path = warp_sequence(source, target)
    >>> print(len(warped) == len(target))
    True
    """
    if not HAS_FASTDTW:
        raise ImportError(
            "fastdtw is required for DTW calculations. "
            "Install with: pip install fastdtw"
        )

    # Get DTW path
    _, path = dtw_distance(source, target)

    target_length = len(target)

    # Build warped sequence
    target_to_source = {}
    for src_idx, tgt_idx in path:
        if tgt_idx not in target_to_source:
            target_to_source[tgt_idx] = []
        target_to_source[tgt_idx].append(src_idx)

    warped = np.zeros(target_length)
    for tgt_idx in range(target_length):
        if tgt_idx in target_to_source:
            src_indices = target_to_source[tgt_idx]
            warped[tgt_idx] = np.mean([source[i] for i in src_indices])
        else:
            # Use nearest mapped value
            for offset in range(1, target_length):
                if tgt_idx - offset >= 0 and (tgt_idx - offset) in target_to_source:
                    warped[tgt_idx] = warped[tgt_idx - offset]
                    break
                if tgt_idx + offset < target_length and (tgt_idx + offset) in target_to_source:
                    src_indices = target_to_source[tgt_idx + offset]
                    warped[tgt_idx] = np.mean([source[i] for i in src_indices])
                    break

    return warped, path


def scale_and_warp(
    source: np.ndarray,
    target_length: int,
    scale_factor: float,
    method: str = "dtw",
) -> np.ndarray:
    """
    Scale and warp a sequence for the Copy model.

    Combines alignment and scaling in one operation.

    Parameters
    ----------
    source : np.ndarray
        Source sequence (historical event values).
    target_length : int
        Desired length for target event.
    scale_factor : float
        Growth factor to apply (T_target / T_source).
    method : {'dtw', 'linear', 'nearest'}, default 'dtw'
        Alignment method.

    Returns
    -------
    forecast : np.ndarray
        Scaled and warped forecast for target event.

    Notes
    -----
    This implements the core of the Analog/Copy model:
    Y_forecast = Y_source_warped × G

    where G = T_target / T_source is the growth factor.

    Examples
    --------
    >>> # Historical 5-day event with 100k-200k impressions
    >>> source = np.array([100000, 150000, 200000, 150000, 100000])
    >>> # Target event is 7 days, with 20% trend growth
    >>> forecast = scale_and_warp(source, target_length=7, scale_factor=1.2)
    """
    # First align to target length
    aligned = align_sequences(source, target_length, method=method)

    # Then apply scale factor
    scaled = aligned * scale_factor

    return scaled


def find_similar_events(
    event_library: List[np.ndarray],
    target_shape: np.ndarray,
    top_k: int = 3,
) -> List[Tuple[int, float, np.ndarray]]:
    """
    Find most similar events in a library based on DTW distance.

    Parameters
    ----------
    event_library : List[np.ndarray]
        Library of historical event patterns.
    target_shape : np.ndarray
        Target event shape to match (can be partial).
    top_k : int, default 3
        Number of similar events to return.

    Returns
    -------
    similar_events : List[Tuple[int, float, np.ndarray]]
        List of (index, distance, event) tuples, sorted by distance.

    Examples
    --------
    >>> library = [
    ...     np.array([100, 200, 300, 200, 100]),
    ...     np.array([150, 250, 350, 250, 150]),
    ...     np.array([80, 120, 160, 200, 160, 120, 80]),
    ... ]
    >>> target = np.array([90, 180, 270])  # Partial event
    >>> similar = find_similar_events(library, target, top_k=2)
    """
    if not HAS_FASTDTW:
        raise ImportError(
            "fastdtw is required for DTW calculations. "
            "Install with: pip install fastdtw"
        )

    # Normalize target shape
    target_norm = (target_shape - target_shape.mean()) / (target_shape.std() + 1e-8)

    distances = []
    for idx, event in enumerate(event_library):
        # Normalize event
        event_norm = (event - event.mean()) / (event.std() + 1e-8)

        # Calculate DTW distance
        dist, _ = dtw_distance(event_norm, target_norm)
        distances.append((idx, dist, event))

    # Sort by distance
    distances.sort(key=lambda x: x[1])

    return distances[:top_k]


def create_event_pattern_library(
    events: List[Tuple[np.ndarray, dict]],
) -> "EventPatternLibrary":
    """
    Create a library of event patterns for similarity matching.

    Parameters
    ----------
    events : List[Tuple[np.ndarray, dict]]
        List of (values, metadata) tuples for each event.

    Returns
    -------
    library : EventPatternLibrary
        Event pattern library for similarity matching.
    """
    return EventPatternLibrary(events)


class EventPatternLibrary:
    """
    Library of historical event patterns for analog forecasting.

    Stores normalized event shapes and metadata for similarity matching.

    Parameters
    ----------
    events : List[Tuple[np.ndarray, dict]], optional
        Initial events to add to the library.

    Examples
    --------
    >>> library = EventPatternLibrary()
    >>> library.add_event(
    ...     values=np.array([100, 200, 300, 200, 100]),
    ...     metadata={"name": "Super Bowl 2023", "type": "sports"}
    ... )
    >>> similar = library.find_similar(target_shape, top_k=3)
    """

    def __init__(self, events: Optional[List[Tuple[np.ndarray, dict]]] = None):
        """Initialize the library."""
        self.events = []
        self.normalized = []
        self.metadata = []

        if events:
            for values, meta in events:
                self.add_event(values, meta)

    def add_event(self, values: np.ndarray, metadata: Optional[dict] = None) -> int:
        """
        Add an event to the library.

        Parameters
        ----------
        values : np.ndarray
            Event time series values.
        metadata : dict, optional
            Metadata about the event (name, type, date, etc.).

        Returns
        -------
        idx : int
            Index of the added event.
        """
        values = np.asarray(values)
        self.events.append(values)

        # Store normalized version for matching
        normalized = (values - values.mean()) / (values.std() + 1e-8)
        self.normalized.append(normalized)

        self.metadata.append(metadata or {})

        return len(self.events) - 1

    def find_similar(
        self,
        target_shape: np.ndarray,
        top_k: int = 3,
        event_type: Optional[str] = None,
    ) -> List[Tuple[int, float, np.ndarray, dict]]:
        """
        Find similar events in the library.

        Parameters
        ----------
        target_shape : np.ndarray
            Target shape to match.
        top_k : int, default 3
            Number of similar events to return.
        event_type : str, optional
            Filter by event type in metadata.

        Returns
        -------
        similar : List[Tuple[int, float, np.ndarray, dict]]
            List of (index, distance, values, metadata) tuples.
        """
        if not self.events:
            return []

        # Filter by event type if specified
        if event_type:
            indices = [
                i for i, m in enumerate(self.metadata)
                if m.get("type") == event_type
            ]
        else:
            indices = list(range(len(self.events)))

        if not indices:
            return []

        # Normalize target
        target_norm = (target_shape - target_shape.mean()) / (target_shape.std() + 1e-8)

        distances = []
        for idx in indices:
            if HAS_FASTDTW:
                dist, _ = dtw_distance(self.normalized[idx], target_norm)
            else:
                # Fall back to Euclidean distance with interpolation
                aligned = _interpolate_linear(self.normalized[idx], len(target_norm))
                dist = np.sqrt(np.mean((aligned - target_norm) ** 2))

            distances.append((idx, dist, self.events[idx], self.metadata[idx]))

        distances.sort(key=lambda x: x[1])
        return distances[:top_k]

    def __len__(self) -> int:
        """Return number of events in library."""
        return len(self.events)

    def __getitem__(self, idx: int) -> Tuple[np.ndarray, dict]:
        """Get event by index."""
        return self.events[idx], self.metadata[idx]
