import torch
import numpy as np
import pandas as pd
from tqdm import tqdm

def assign_events_to_cities(
    city_df, event_df, distance_threshold_km, device,
    city_lat='Latitude', city_lon='Longitude', city_country='Country',
    event_lat='Latitude', event_lon='Longitude', event_country='Country',
    batch_size=10000
):
    """
    Assign each event to the nearest city (within distance_threshold_km and same country),
    and return city_df with a new 'Number_of_Events' column.
    """

    def to_rad_tensor(df, lat_col, lon_col):
        return torch.deg2rad(torch.tensor(df[[lat_col, lon_col]].values, dtype=torch.float32, device=device))

    def build_country_mask(src_countries, dst_countries):
        src = np.array(src_countries)
        dst = np.array(dst_countries)
        return torch.tensor((src[:, None] == dst[None, :]), dtype=torch.bool, device=device)

    def haversine_distance_matrix(src_rad, dst_rad):
        src_lat, src_lon = src_rad[:, 0:1], src_rad[:, 1:2]
        dst_lat, dst_lon = dst_rad[:, 0], dst_rad[:, 1]
        dlat = dst_lat - src_lat
        dlon = dst_lon - src_lon
        a = torch.sin(dlat / 2) ** 2 + torch.cos(src_lat) * torch.cos(dst_lat) * torch.sin(dlon / 2) ** 2
        c = 2 * torch.arcsin(torch.sqrt(a))
        return 6371 * c  # km

    # Prepare coordinate tensors
    event_coords = to_rad_tensor(event_df, event_lat, event_lon)
    city_coords = to_rad_tensor(city_df, city_lat, city_lon)
    country_mask = build_country_mask(event_df[event_country].values, city_df[city_country].values)

    # Assignment arrays
    min_dists, min_idxs = [], []

    for start in tqdm(range(0, len(event_df), batch_size), desc="📍 Assigning events to cities"):
        end = min(start + batch_size, len(event_df))
        batch_coords = event_coords[start:end]
        mask_batch = country_mask[start:end]

        dist_matrix = haversine_distance_matrix(batch_coords, city_coords)
        dist_matrix[~mask_batch] = 1e6  # mask out different countries

        min_dist, min_idx = torch.min(dist_matrix, dim=1)
        min_dists.append(min_dist)
        min_idxs.append(min_idx)

    min_dists = torch.cat(min_dists)
    min_idxs = torch.cat(min_idxs)
    valid_mask = min_dists <= distance_threshold_km

    # Build result DataFrame
    assigned_idxs = min_idxs[valid_mask].cpu().numpy()
    result_df = city_df.copy()
    result_df['Number_of_Events'] = 0

    # Safe event assignment using .iloc to align correctly
    counts = pd.Series(assigned_idxs).value_counts()
    for idx, count in counts.items():
        result_df.iloc[idx, result_df.columns.get_loc('Number_of_Events')] = count

    return result_df


def calculate_served_population(airports_df, cities_df, distance_threshold_km, device, batch_size=2000):
    """
    Assign each city's population to the airport with the highest score:
        score = Departures / (Distance_km + 1)
    Only airports within 50 km and same country are considered.
    """

    # Tensors
    airport_coords = torch.tensor(airports_df[['Latitude', 'Longitude']].values, dtype=torch.float32).to(device)
    city_coords = torch.tensor(cities_df[['Latitude', 'Longitude']].values, dtype=torch.float32).to(device)
    city_pops = torch.tensor(cities_df['Population'].values, dtype=torch.float32).to(device)
    airport_departures = torch.tensor(airports_df['Number_of_Flights'].values, dtype=torch.float32).to(device)

    airport_countries = airports_df['Country'].values
    city_countries = cities_df['Country'].values

    def haversine_batch(lat1, lon1, lat2, lon2):
        lat1_rad = torch.deg2rad(lat1)
        lon1_rad = torch.deg2rad(lon1)
        lat2_rad = torch.deg2rad(lat2)
        lon2_rad = torch.deg2rad(lon2)
        dlon = lon2_rad - lon1_rad
        dlat = lat2_rad - lat1_rad
        a = torch.sin(dlat / 2) ** 2 + torch.cos(lat1_rad) * torch.cos(lat2_rad) * torch.sin(dlon / 2) ** 2
        c = 2 * torch.asin(torch.sqrt(a))
        return 6371 * c

    assigned_pop = torch.zeros(len(airport_coords), device=device)

    for start in tqdm(range(0, len(city_coords), batch_size), desc="📊 Assigning using weighted score"):
        end = min(start + batch_size, len(city_coords))
        batch_coords = city_coords[start:end]
        batch_pops = city_pops[start:end]
        batch_countries = city_countries[start:end]

        for i in range(batch_coords.shape[0]):
            country = batch_countries[i]
            matching_idx = [j for j, c in enumerate(airport_countries) if c == country]

            if not matching_idx:
                continue

            subset_coords = airport_coords[matching_idx]
            subset_deps = airport_departures[matching_idx]
            lat1 = batch_coords[i, 0].unsqueeze(0)
            lon1 = batch_coords[i, 1].unsqueeze(0)

            dists = haversine_batch(
                lat1, lon1,
                subset_coords[:, 0], subset_coords[:, 1]
            ).squeeze()

            in_range_mask = dists <= distance_threshold_km
            if torch.any(in_range_mask):
                dists_in_range = dists[in_range_mask]
                deps_in_range = subset_deps[in_range_mask]
                scores = deps_in_range / (dists_in_range + 1)
                best_idx = torch.argmax(scores)
                chosen_airport = torch.tensor(matching_idx)[in_range_mask][best_idx]
                assigned_pop[chosen_airport] += batch_pops[i]

        torch.cuda.empty_cache()

    return assigned_pop.cpu().numpy()