import numpy as np
import pandas as pd
from typing import List, Tuple, Optional, Dict, Any
from src.analysis.geospatial import haversine_distance_vectorized
from src.utils import format_datetime_ddmmyyyy_hhmmss

def build_lightning_table_rows(
	centroid_lat: float,
	centroid_lng: float,
	lightning_df: pd.DataFrame,
	distance_rings_m: List[int],
	ring_colors: List[str],
) -> Tuple[List[List[str]], List[str]]:
	rows: List[List[str]] = []
	row_colors: List[str] = []

	if len(lightning_df) == 0:
		return rows, row_colors

	outermost_km = max(distance_rings_m) / 1000.0
	rings_km = np.array(distance_rings_m, dtype=float) / 1000.0

	dists_km = haversine_distance_vectorized(
		centroid_lat, centroid_lng,
		lightning_df['lat'].values, lightning_df['lng'].values,
	) / 1000.0

	for i, rec in enumerate(lightning_df.itertuples(index=False)):
		proximity = dists_km[i]
		if proximity > outermost_km:
			continue

		ring_idx = int(np.searchsorted(rings_km, proximity, side='left'))
		if ring_idx >= len(rings_km):
			continue
		color = ring_colors[ring_idx]

		try:
			dt_val = rec.local_time if not isinstance(rec.local_time, str) else pd.to_datetime(str(rec.local_time)[:19])
			local_time = format_datetime_ddmmyyyy_hhmmss(pd.to_datetime(dt_val))
		except Exception:
			local_time = str(getattr(rec, 'local_time', ''))[:19]

		lightning_type = "cloud-to-ground" if str(rec.p_type) == '0' else "intercloud"
		rows.append([
			"",
			local_time,
			f"{rec.lat:.5f}",
			f"{rec.lng:.5f}",
			str(rec.current),
			str(getattr(rec, 'height', '')),
			lightning_type,
			f"{proximity:.2f}"
		])
		row_colors.append(color)

	return rows, row_colors


def precompute_group_distances_and_rings(
	centroid_lat: float,
	centroid_lng: float,
	lightning_df: pd.DataFrame,
	distance_rings_m: List[int],
) -> Dict[str, Any]:
	if len(lightning_df) == 0:
		return {
			'dists_km': np.array([], dtype=float),
			'ring_idx': np.array([], dtype=int),
			'mask_within': np.array([], dtype=bool),
		}
	rings_km = np.array(distance_rings_m, dtype=float) / 1000.0
	dists_km = haversine_distance_vectorized(
		centroid_lat, centroid_lng,
		lightning_df['lat'].values, lightning_df['lng'].values,
	) / 1000.0
	mask_within = dists_km <= rings_km[-1]
	# For values beyond last ring, clamp index to len(rings)
	ring_idx = np.searchsorted(rings_km, dists_km, side='left')
	return {
		'dists_km': dists_km,
		'ring_idx': ring_idx,
		'mask_within': mask_within,
	}


