from __future__ import annotations

import math

import numpy as np
import pandas as pd


JITTER_MERCADO_PASSO_METROS = 5.0
JITTER_MERCADO_MAX_RAIO_METROS = 28.0


def aplicar_jitter_dados_mercado(
    df_mapa: pd.DataFrame,
    lat_col: str,
    lon_col: str,
    lat_plot_col: str = "__mesa_lat_plot__",
    lon_plot_col: str = "__mesa_lon_plot__",
    *,
    precisao: int = 7,
    passo_metros: float = JITTER_MERCADO_PASSO_METROS,
    max_raio_metros: float = JITTER_MERCADO_MAX_RAIO_METROS,
) -> pd.DataFrame:
    """
    Aplica jitter visual determinístico apenas em pontos de mercado sobrepostos.

    As coordenadas originais permanecem nas colunas de entrada; as coordenadas
    deslocadas são gravadas em lat_plot_col/lon_plot_col.
    """
    df_plot = df_mapa.copy()
    df_plot[lat_plot_col] = pd.to_numeric(df_plot[lat_col], errors="coerce").astype(float)
    df_plot[lon_plot_col] = pd.to_numeric(df_plot[lon_col], errors="coerce").astype(float)

    if len(df_plot) <= 1:
        return df_plot

    chave_lat = df_plot[lat_col].round(precisao)
    chave_lon = df_plot[lon_col].round(precisao)
    grupos = df_plot.groupby([chave_lat, chave_lon], sort=False)

    metros_por_grau_lat = 111_320.0
    lat_plot_pos = int(df_plot.columns.get_loc(lat_plot_col))
    lon_plot_pos = int(df_plot.columns.get_loc(lon_plot_col))

    for _, idx_labels in grupos.indices.items():
        posicoes = np.asarray(idx_labels, dtype=int)
        if posicoes.size <= 1:
            continue

        base_lat = float(df_plot.iat[int(posicoes[0]), lat_plot_pos])
        base_lon = float(df_plot.iat[int(posicoes[0]), lon_plot_pos])
        if not np.isfinite(base_lat) or not np.isfinite(base_lon):
            continue

        seed_val = int((abs(base_lat) * 1_000_000.0) + (abs(base_lon) * 1_000_000.0) * 3.0) % 360
        angulo_base = math.radians(seed_val)
        cos_lat = max(abs(math.cos(math.radians(base_lat))), 1e-6)
        metros_por_grau_lon = metros_por_grau_lat * cos_lat

        for pos, pos_idx in enumerate(posicoes):
            if pos == 0:
                continue

            pos_ring = pos - 1
            ring = 1
            while pos_ring >= (6 * ring):
                pos_ring -= 6 * ring
                ring += 1

            slots_ring = max(6 * ring, 1)
            angulo = angulo_base + (2.0 * math.pi * (pos_ring / slots_ring))
            raio_m = min(ring * passo_metros, max_raio_metros)
            delta_lat = (raio_m * math.sin(angulo)) / metros_por_grau_lat
            delta_lon = (raio_m * math.cos(angulo)) / metros_por_grau_lon

            df_plot.iat[int(pos_idx), lat_plot_pos] = base_lat + delta_lat
            df_plot.iat[int(pos_idx), lon_plot_pos] = base_lon + delta_lon

    return df_plot