import time
import math
import numpy as np
import pandas as pd
from tqdm import trange
from sievecomp.io.sql import load_sql
from sievecomp.core.runner import Runner

def fetch_events_by_time_range(
    runner: Runner,
    query_name: str,
    start_time: pd.Timestamp,
    end_time: pd.Timestamp
) -> pd.DataFrame:
    """
    Fetches events from a specific SeisComP database query within a dynamic time window.
    Bypasses the engine's checks and duplicates logic.
    """
    # 1. Retrieve the configured time column used for sorting and filtering
    time_col = runner._cm.get_time_column()

    # 2. Prepare the dynamic time bounds and SQL parameters
    sql_params = {}
    where_clauses = []

    if start_time:
        sql_params["start_time"] = start_time.strftime('%Y-%m-%d %H:%M:%S')
        where_clauses.append(f"base_query.{time_col} >= %(start_time)s")

    if end_time:
        sql_params["end_time"] = end_time.strftime('%Y-%m-%d %H:%M:%S')
        where_clauses.append(f"base_query.{time_col} < %(end_time)s")

    where_sql = " AND ".join(where_clauses)

    # 3. Load the base SQL text strictly as it is defined in the TOML configuration
    query_cfg = runner._cm.select_query(name=query_name)
    base_sql = load_sql(
        query_cfg=query_cfg,
        base_dir=runner._cm.base_dir,
        config_name=runner._cm.config_path.name
    )

    # 4. Wrap the base query as a subquery, applying the new WHERE clauses and sorting
    wrapped_sql = f"""SELECT * FROM ({base_sql}) AS base_query WHERE {where_sql} ORDER BY base_query.{time_col} ASC"""

    # 5. Fetch the events (skipping the check/duplicate evaluation steps for raw data retrieval)
    result = runner.run(
        query_name=query_name,
        multi_query=False,
        perform_duplicates=False,
        perform_checks=False,
        sql_text_override=wrapped_sql,
        sql_params=sql_params
    )

    # 6. Return the raw queried events DataFrame
    return result.total_events

columns = ['time_value', 'publicID', 'text', 'depth_value', 'magnitude_value', 'magnitude_type',
           'quality_standardError', 'depth_uncertainty', 'latitude_uncertainty', 'longitude_uncertainty',
           'quality_associatedPhaseCount', 'creationInfo_author', 'event_type', 'creationInfo_agencyID']

def calculate_haversine_distance(lat1: float, lon1: float, lat2: float, lon2: float) -> float:
    """
    Calculates the great-circle distance in kilometers between two points
    on the Earth's surface using the Haversine formula.
    """
    R = 6371.0  # Earth radius in kilometers

    # Convert decimal degrees to radians
    phi1 = math.radians(lat1)
    phi2 = math.radians(lat2)
    dphi = math.radians(lat2 - lat1)
    dlambda = math.radians(lon2 - lon1)

    # Haversine formula
    a = math.sin(dphi / 2.0)**2 + math.cos(phi1) * math.cos(phi2) * math.sin(dlambda / 2.0)**2
    c = 2 * math.atan2(math.sqrt(a), math.sqrt(1 - a))

    return R * c


def check_duplicates(
        events: pd.DataFrame
):
    """
    Identifies duplicate events based on time and geographical distance.

    Parameters:
    -----------
    events: pandas dataframe
        Seismic data for the given time range. Generally obtained from the connect2mysql function.

    Returns:
    --------
    duplicates: A pandas dataframe
        Table with information about duplicate events.
    """
    # Filter events with the following types
    checks = ["earthquake", "volcanic eruption", "explosion", "outside of network interest"]
    selections = events[events['event_type'].isin(checks)].reset_index(drop=True)

    # Make a pandas dataframe to store the info of the duplicates
    duplicates = pd.DataFrame()

    # Loop over al earthquakes, ordered by time
    for i in trange(len(selections) - 1, desc="Checking duplicates", unit=" events", leave=False):
        event1 = selections.iloc[i]  # Take the i-esim event on the list
        event2 = selections.iloc[i + 1]  # Compared to the next event

        # Check if the events are within 4 seconds of each other
        time_diff = abs((event2['time_value'] - event1['time_value']).total_seconds())
        if time_diff <= 4:
            # Estimate the distance between the two events using haversine formula
            distance = calculate_haversine_distance(event1['latitude_value'], event1['longitude_value'],
                                                    event2['latitude_value'], event2['longitude_value'])
            if distance <= 100:  # If both events are within 100 km and 4 seconds, they can be duplicates
                dup_1 = event1[columns].copy()
                dup_1['Observations'] = f'Possible duplicate event of {event2["publicID"]}'
                dup_2 = event2[columns].copy()
                dup_2['Observations'] = f'Possible duplicate event of {event1["publicID"]}'
                # Add the two events to the duplicates dataframe
                duplicates = pd.concat([duplicates, dup_1.to_frame().T], ignore_index=True)
                duplicates = pd.concat([duplicates, dup_2.to_frame().T], ignore_index=True)
    return duplicates

def check_duplicates_adjacent(events: pd.DataFrame) -> pd.DataFrame:
    """
    Identifies duplicate events based on time and geographical distance.
    Same adjacent-only logic as the original, but micro-optimized:
      - Pre-extracted NumPy arrays (avoids per-row iloc overhead)
      - Row accumulation in a list (avoids repeated pd.concat in loop)
      - Single pd.concat at the end

    Time complexity: O(n) comparisons — only i vs i+1
    Parameters:
    -----------
    events : pd.DataFrame
        Seismic data for the given time range.
    Returns:
    --------
    duplicates : pd.DataFrame
        Table with information about duplicate events.
    """
    checks = ["earthquake", "volcanic eruption", "explosion", "outside of network interest"]
    selections = events[events['event_type'].isin(checks)].reset_index(drop=True)

    # Edge case: empty or single-event dataset
    if len(selections) < 2:
        return pd.DataFrame(columns=columns + ['Observations'])

    # --- Pre-extract arrays for fast access ---
    times  = selections['time_value'].values          # numpy datetime64
    lats   = selections['latitude_value'].values.astype(np.float64)
    lons   = selections['longitude_value'].values.astype(np.float64)
    ids    = selections['publicID'].values

    TIME_WINDOW    = np.timedelta64(4, 's')
    DIST_THRESHOLD = 100  # km

    dup_rows = []  # Accumulate rows here, build DF once at the end

    for i in trange(len(selections) - 1, desc="Checking duplicates", unit=" events", leave=False):
        # Time check using numpy (faster than .total_seconds() on Timedelta)
        if abs(times[i + 1] - times[i]) <= TIME_WINDOW:
            # Distance check using haversine formula
            distance = calculate_haversine_distance(lats[i], lons[i], lats[i + 1], lons[i + 1])
            if distance <= DIST_THRESHOLD:
                row_i = selections.iloc[i][columns].copy()
                row_i['Observations'] = f'Possible duplicate event of {ids[i + 1]}'
                row_j = selections.iloc[i + 1][columns].copy()
                row_j['Observations'] = f'Possible duplicate event of {ids[i]}'
                dup_rows.extend([row_i, row_j])

    # Concat duplicates to one single Dataframe
    if dup_rows:
        duplicates = pd.DataFrame(dup_rows, columns=columns + ['Observations']).reset_index(drop=True)
    else:
        duplicates = pd.DataFrame(columns=columns + ['Observations'])

    return duplicates

def haversine_np(lat1, lon1, lat2, lon2):
    """
    Fully vectorized haversine distance (km) between paired lat/lon arrays.
    Replaces the scalar `hs.haversine()` calls inside the original loop.
    """
    R = 6371.0088
    lat1r, lat2r = np.radians(lat1), np.radians(lat2)
    dlat = lat2r - lat1r
    dlon = np.radians(lon2 - lon1)
    a = np.sin(dlat / 2.0) ** 2 + np.cos(lat1r) * np.cos(lat2r) * np.sin(dlon / 2.0) ** 2
    return 2 * R * np.arcsin(np.sqrt(a))

def check_duplicates_vectorized(
        events: pd.DataFrame,
        time_window: int,
        dist_threshold: float,
) -> pd.DataFrame:
    """
    Identifies duplicate events based on time proximity and geographical
    distance between temporally-adjacent events (after sorting by time).

    Time complexity: O(n) — all comparisons are vectorized NumPy array
    operations; only the (typically rare) confirmed duplicate pairs are
    iterated in Python to build per-row observation text.

    Parameters
    ----------
    events : pd.DataFrame
        Seismic data for the given time range. Must contain 'time_value',
        'latitude_value', 'longitude_value', and 'publicID' columns.
    time_window : int
        Time window in seconds to detect duplicated events.
    dist_threshold : float
        Distance threshold in kilometers to detect duplicates.

    Returns
    -------
    pd.DataFrame
        Subset of `events` (all original columns preserved) with an added
        'Observations' column, containing only the flagged duplicate rows.
    """
    if len(events) < 2:
        empty = events.copy()
        empty['Observations'] = pd.Series(dtype='object')
        return empty.iloc[0:0]

    sorted_events = events.sort_values('time_value').reset_index(drop=True)

    times = sorted_events['time_value'].to_numpy()
    lats  = sorted_events['latitude_value'].to_numpy(dtype=np.float64)
    lons  = sorted_events['longitude_value'].to_numpy(dtype=np.float64)
    ids   = sorted_events['publicID'].to_numpy()

    time_diff   = np.abs(times[1:] - times[:-1])
    within_time = time_diff <= np.timedelta64(time_window, 's')

    distances   = haversine_np(lats[:-1], lons[:-1], lats[1:], lons[1:])
    within_dist = distances <= dist_threshold

    adjacent_dup = within_time & within_dist
    dup_idx = np.where(adjacent_dup)[0]

    if dup_idx.size == 0:
        empty = events.copy()
        empty['Observations'] = pd.Series(dtype='object')
        return empty.iloc[0:0]

    obs_map: dict[int, list[str]] = {}
    for i in dup_idx:
        obs_map.setdefault(int(i), []).append(f'Possible duplicate event of {ids[i + 1]}')
        obs_map.setdefault(int(i) + 1, []).append(f'Possible duplicate event of {ids[i]}')

    flagged_positions = sorted(obs_map.keys())
    flagged = sorted_events.iloc[flagged_positions].copy()
    flagged['Observations'] = ['; '.join(obs_map[p]) for p in flagged_positions]

    return flagged.reset_index(drop=True)

def check_duplicates_2(events: pd.DataFrame) -> pd.DataFrame:
    """
        Identifies duplicate events based on time and geographical distance. Optimized using a sorted sliding time window (O(n log n)).

    Parameters:
    -----------
    events : pd.DataFrame
        Seismic data for the given time range.

    Returns:
    --------
    duplicates : pd.DataFrame
        Table with information about duplicate events.
    """
    checks = ["earthquake", "volcanic eruption", "explosion", "outside of network interest"]
    selections = (
        events[events['event_type'].isin(checks)]
        .sort_values('time_value')
        .reset_index(drop=True)
    )

    # Edge case: empty or single-event dataset
    if len(selections) < 2:
        return pd.DataFrame(columns=columns + ['Observations'])

    # Pre-extract arrays for fast access (avoids per-row iloc overhead)
    times = selections['time_value'].values
    lats  = selections['latitude_value'].values
    lons  = selections['longitude_value'].values
    ids   = selections['publicID'].values

    TIME_WINDOW   = 4    # seconds
    DIST_THRESHOLD = 100  # km

    # Track which (i, j) pairs were already flagged to avoid duplicate rows
    flagged_pairs = set()
    dup_rows = []

    left = 0  # Left pointer of the sliding window

    for right in trange(1, len(selections), desc="Checking duplicates", unit=" events", leave=False):
        # Advance left pointer: drop events outside the 4-second window
        while left < right:
            delta = (times[right] - times[left]) / pd.Timedelta('1s')
            if delta > TIME_WINDOW:
                left += 1
            else:
                break

        # Compare `right` against every event still inside the window [left, right-1]
        for i in range(left, right):
            pair = (i, right) if i < right else (right, i)
            if pair in flagged_pairs:
                continue

            time_diff = abs(
                (selections.iloc[right]['time_value'] - selections.iloc[i]['time_value'])
                .total_seconds()
            )
            if time_diff > TIME_WINDOW:
                continue

            distance = calculate_haversine_distance(lats[i], lons[i], lats[right], lons[right])
            if distance <= DIST_THRESHOLD:
                flagged_pairs.add(pair)

                row_i = selections.iloc[i][columns].copy()
                row_i['Observations'] = f'Possible duplicate event of {ids[right]}'

                row_j = selections.iloc[right][columns].copy()
                row_j['Observations'] = f'Possible duplicate event of {ids[i]}'

                dup_rows.extend([row_i, row_j])

    # Build duplicates dataframe in one shot (avoids repeated pd.concat overhead)
    if dup_rows:
        duplicates = pd.DataFrame(dup_rows, columns=columns + ['Observations']).reset_index(drop=True)
    else:
        duplicates = pd.DataFrame(columns=columns + ['Observations'])

    return duplicates

def check_duplicates_2_vectorized(
        events: pd.DataFrame,
        columns: list[str],
        time_window: int,
        dist_threshold: float,
) -> pd.DataFrame:
    """
    Identifies duplicate events based on time proximity and geographical
    distance, using a sliding time window (compares every event against
    ALL others within time_window seconds, not just its direct neighbor).

    Parameters
    ----------
    events : pd.DataFrame
        Seismic data for the given time range. Must contain the columns
        referenced by `columns`.
    columns : list[str]
        Ordered list [time_col, lat_col, lon_col, id_col] identifying which
        columns hold time, latitude, longitude, and event ID values.
    time_window : int
        Time window in seconds to detect duplicated events.
    dist_threshold : float
        Distance threshold in kilometers to detect duplicates.

    Returns
    -------
    pd.DataFrame
        Subset of `events` (all original columns preserved) with an added
        'Observations' column, one row per flagged duplicate event —
        chains of 3+ nearby events are merged into a single row per event
        rather than duplicated across multiple pair-rows.
    """
    time_col, lat_col, lon_col, id_col = columns
    selections = events.copy()

    if len(selections) < 2:
        empty = events.copy()
        empty['Observations'] = pd.Series(dtype='object')
        return empty.iloc[0:0]

    sorted_events = selections.sort_values(time_col).reset_index(drop=True)

    times = sorted_events[time_col].to_numpy()
    lats  = sorted_events[lat_col].to_numpy(dtype=np.float64)
    lons  = sorted_events[lon_col].to_numpy(dtype=np.float64)
    ids   = sorted_events[id_col].to_numpy()
    n     = len(sorted_events)

    window_delta = np.timedelta64(time_window, 's')

    # For every i, right_bound[i] = first index j where times[j] > times[i] + window
    right_bound = np.searchsorted(times, times + window_delta, side='right')

    obs_map: dict[int, list[str]] = {}

    for i in range(n - 1):
        j_end = right_bound[i]
        if j_end <= i + 1:
            continue

        window_lats = lats[i + 1:j_end]
        window_lons = lons[i + 1:j_end]
        distances = haversine_np(lats[i], lons[i], window_lats, window_lons)

        matched_offsets = np.where(distances <= dist_threshold)[0]
        if matched_offsets.size == 0:
            continue

        matched_js = i + 1 + matched_offsets
        for j in matched_js:
            obs_map.setdefault(i, []).append(f'Possible duplicate event of {ids[j]}')
            obs_map.setdefault(int(j), []).append(f'Possible duplicate event of {ids[i]}')

    if not obs_map:
        empty = events.copy()
        empty['Observations'] = pd.Series(dtype='object')
        return empty.iloc[0:0]

    flagged_positions = sorted(obs_map.keys())
    flagged = sorted_events.iloc[flagged_positions].copy()
    flagged['Observations'] = ['; '.join(obs_map[p]) for p in flagged_positions]

    return flagged.reset_index(drop=True)