"""DTW-based similarity metrics — exact replica of legacy env computations.

Provides ``calc_lag`` (cross-correlation lag), ``calc_sim`` (DTW similarity),
and ``compute_similarity`` to match legacy env reward computation.

Legacy envs used:
- Karman cloak: lag from sensor1 Uy, then DTW on 6 sensor channels
- Illusion: lag from target[:,3] vs state[:,1], then DTW on 6 sensor channels (offset +2)
- Vortex: no lag, roll by current_step+1
- Erase: lag from force channels, uses enhanced calc_sim
"""
from __future__ import annotations

from typing import Optional

import numpy as np


def calc_lag(target: np.ndarray, state: np.ndarray) -> int:
    """Cross-correlation lag between target and state sequences.

    Args:
        target: shape ``(N,)`` reference signal.
        state: shape ``(M,)`` observed signal.

    Returns:
        Integer lag (positive = state is ahead of target).
    """
    t_mean = np.mean(target)
    s_mean = np.mean(state)
    correlation = np.correlate(target - t_mean, state - s_mean, mode="full")
    lags = np.arange(-len(target) + 1, len(target))
    return int(lags[np.argmax(correlation)])


def calc_dtw_sim(target: np.ndarray, state: np.ndarray) -> float:
    """Standard DTW similarity (used by cloak, illusion, vortex, reduce_obs).

    Args:
        target: shape ``(N,)`` reference.
        state: shape ``(M,)`` observed.

    Returns:
        Similarity in [0, 1], where 1 = perfect match.
    """
    n, m = len(target), len(state)
    dtw = np.full((n + 1, m + 1), np.inf)
    dtw[0, 0] = 0.0
    for i in range(1, n + 1):
        for j in range(1, m + 1):
            cost = abs(float(target[i - 1]) - float(state[j - 1]))
            last_min = min(dtw[i - 1, j], dtw[i, j - 1], dtw[i - 1, j - 1])
            dtw[i, j] = cost + last_min
    return float(1.0 - dtw[n, m] / float(n))


def calc_dtw_sim_enhanced(target: np.ndarray, state: np.ndarray) -> float:
    """Enhanced DTW similarity with amplitude-ratio and mean components.

    Used by the legacy erase env. Combines:
    - 80% standard DTW (max-cost normalised)
    - 10% amplitude ratio (min_std/max_std)
    - 10% mean similarity (1/(1 + diff/scale*10))

    Args:
        target: shape ``(N,)`` reference.
        state: shape ``(M,)`` observed.

    Returns:
        Combined similarity in [0, 1].
    """
    target_arr = np.asarray(target, dtype=np.float64)
    state_arr = np.asarray(state, dtype=np.float64)
    n, m = len(target_arr), len(state_arr)

    # Amplitude ratio component
    t_std = max(np.std(target_arr), 1e-8)
    s_std = max(np.std(state_arr), 1e-8)
    amplitude_ratio = float(min(t_std, s_std) / max(t_std, s_std))

    # Mean similarity component
    mean_diff = abs(np.mean(target_arr) - np.mean(state_arr))
    max_scale = max(abs(np.mean(target_arr)), abs(np.mean(state_arr)), 1e-8)
    mean_similarity = 1.0 / (1.0 + mean_diff / max_scale * 10.0)

    # DTW with max-possible-cost normalisation
    dtw = np.full((n + 1, m + 1), np.inf)
    dtw[0, 0] = 0.0
    for i in range(1, n + 1):
        for j in range(1, m + 1):
            cost = abs(target_arr[i - 1] - state_arr[j - 1])
            last_min = min(dtw[i - 1, j], dtw[i, j - 1], dtw[i - 1, j - 1])
            dtw[i, j] = cost + last_min

    max_possible_cost = max(np.max(np.abs(target_arr)), np.max(np.abs(state_arr)), 1e-8)
    dtw_distance = dtw[n, m] / (n * max_possible_cost)
    dtw_sim = max(0.0, 1.0 - dtw_distance)

    return float(0.8 * dtw_sim + 0.1 * amplitude_ratio + 0.1 * mean_similarity)


def compute_similarity_karman_cloak(
    target_states: np.ndarray,
    fifo_states: np.ndarray,
    conv_len: int = 30,
) -> float:
    """Compute DTW similarity for Karman cloak (standard pattern).

    Matches legacy code:
    1. Compute lag from middle sensor (index 1) Uy component
    2. For all 6 sensor channels, roll target by lag, compute DTW, average

    Args:
        target_states: shape ``(FIFO_LEN, 6)`` target sensor data.
        fifo_states: shape ``(FIFO_LEN, 6)`` current FIFO sensor data.
        conv_len: Convergence window length (default 30).

    Returns:
        Average similarity over 6 channels in [0, 1].
    """
    target = np.asarray(target_states, dtype=np.float64)
    state = np.asarray(fifo_states, dtype=np.float64)

    id_sens = 1  # middle sensor
    target_seq = target[conv_len:2 * conv_len, id_sens]
    state_seq = state[-conv_len:, id_sens]
    lag = calc_lag(target_seq, state_seq)

    similarities = 0.0
    for i in range(6):
        t_seq = np.roll(target[:, i], -lag)[conv_len:2 * conv_len]
        s_seq = state[-conv_len:, i]
        similarities += calc_dtw_sim(t_seq, s_seq)
    return float(similarities / 6.0)


def compute_similarity_vortex(
    target_states: np.ndarray,
    fifo_states: np.ndarray,
    current_step: int,
    conv_len: int = 30,
) -> float:
    """Compute DTW similarity for vortex (no lag, roll by current_step+1).

    Matches legacy vortex env.

    Args:
        target_states: shape ``(FIFO_LEN, 6)`` target sensor data.
        fifo_states: shape ``(FIFO_LEN, 6)`` current FIFO data.
        current_step: The current simulation step index.
        conv_len: Convergence window length (default 30).

    Returns:
        Average similarity over 6 channels.
    """
    target = np.asarray(target_states, dtype=np.float64)
    state = np.asarray(fifo_states, dtype=np.float64)

    similarities = 0.0
    for i in range(6):
        t_seq = np.roll(target[-conv_len:, i], -current_step - 1)
        s_seq = state[-conv_len:, i]
        similarities += calc_dtw_sim(t_seq, s_seq)
    return float(similarities / 6.0)


def compute_similarity_illusion(
    target_states: np.ndarray,
    fifo_states: np.ndarray,
    conv_len: int = 36,
) -> float:
    """Compute DTW similarity for illusion.

    Matches legacy imit env:
    1. lag from target[:, id_sens+2] vs state[:, id_sens] (offset by 2)
    2. For 6 channels, target uses [:, i+2] offset

    Args:
        target_states: shape ``(FIFO_LEN, 8)`` (2 force + 6 sensor channels).
        fifo_states: shape ``(FIFO_LEN, 6)`` current FIFO (6 sensors only).
        conv_len: Convergence window length (default 36).

    Returns:
        Average similarity over 6 channels.
    """
    target = np.asarray(target_states, dtype=np.float64)
    state = np.asarray(fifo_states, dtype=np.float64)

    id_sens = 1
    t_seq_ref = target[conv_len:2 * conv_len, id_sens + 2]
    s_seq_ref = state[-conv_len:, id_sens]
    lag = calc_lag(t_seq_ref, s_seq_ref)

    similarities = 0.0
    for i in range(6):
        t_seq = np.roll(target[:, i + 2], -lag)[conv_len:2 * conv_len]
        s_seq = state[-conv_len:, i]
        similarities += calc_dtw_sim(t_seq, s_seq)
    return float(similarities / 6.0)


# ---------------------------------------------------------------------------
# Harmonics analysis (used by illusion)
# ---------------------------------------------------------------------------
def analyze_harmonics(
    states: np.ndarray,
    n_harmonics: int = 5,
) -> list:
    """FFT-based harmonic analysis of multi-channel time series.

    Matches legacy ``analyze_harmonics()``.

    Args:
        states: shape ``(N, D)`` time-series data.
        n_harmonics: Number of harmonics to extract per channel.

    Returns:
        List of D dicts, each with keys:
            dc: float (DC component)
            amps: (n_harmonics,) array
            freqs: (n_harmonics,) array
            phases: (n_harmonics,) array
    """
    N, D = states.shape
    result = []
    for d in range(D):
        y = states[:, d]
        fft_coef = np.fft.rfft(y)
        freqs = np.fft.rfftfreq(N, d=1)
        amps = 2.0 * np.abs(fft_coef) / N
        phases = np.angle(fft_coef)
        idx = np.argsort(amps[1:])[::-1][:n_harmonics] + 1
        harmonics = {
            "dc": float(np.real(fft_coef[0]) / N),
            "amps": np.array(amps[idx], dtype=np.float32),
            "freqs": np.array(freqs[idx], dtype=np.float32),
            "phases": np.array(phases[idx], dtype=np.float32),
        }
        result.append(harmonics)
    return result


def gen_target_states_at(t, harmonics) -> np.ndarray:
    """Reconstruct target state at time step t from harmonics.

    Matches legacy ``gen_target_states_at()``.

    Args:
        t: Integer step index.
        harmonics: Output from ``analyze_harmonics()``.

    Returns:
        shape ``(D,)`` reconstructed state vector.
    """
    D = len(harmonics)
    result = np.zeros(D, dtype=np.float32)
    for d, h in enumerate(harmonics):
        val = float(h["dc"])
        for amp, freq, phase in zip(h["amps"], h["freqs"], h["phases"]):
            val += amp * np.cos(2.0 * np.pi * freq * t + phase)
        result[d] = val
    return result
