from __future__ import annotations

import math

import numpy as np
import pandas as pd


INPUT_XLSX = "task1/01_clean.xlsx"
OUTPUT_XLSX = "task1/02_neighbor.xlsx"

RHO_MILES_LIST = [10.0, 20.0, 30.0]


def haversine_miles(lat1: float, lon1: float, lat2: float, lon2: float) -> float:
    r_miles = 3958.7613
    phi1 = math.radians(lat1)
    phi2 = math.radians(lat2)
    dphi = math.radians(lat2 - lat1)
    dlambda = math.radians(lon2 - lon1)

    a = math.sin(dphi / 2.0) ** 2 + math.cos(phi1) * math.cos(phi2) * math.sin(dlambda / 2.0) ** 2
    c = 2.0 * math.atan2(math.sqrt(a), math.sqrt(1.0 - a))
    return r_miles * c


def main() -> None:
    sites = pd.read_excel(INPUT_XLSX, sheet_name="sites")

    lat = sites["lat"].to_numpy(float)
    lon = sites["lon"].to_numpy(float)
    mu = sites["mu_clients_per_visit"].to_numpy(float)

    n = len(sites)
    dist = np.zeros((n, n), dtype=float)
    for i in range(n):
        for j in range(i + 1, n):
            d = haversine_miles(lat[i], lon[i], lat[j], lon[j])
            dist[i, j] = d
            dist[j, i] = d

    out = sites.copy()
    out["neighbor_demand_mu_self"] = mu.copy()
    for rho in RHO_MILES_LIST:
        w = np.exp(-(dist**2) / (2.0 * rho * rho))
        # D_i(rho) = sum_j mu_j * exp(-dist(i,j)^2 / (2 rho^2))
        out[f"neighbor_demand_mu_rho_{int(rho)}mi"] = w @ mu

    dist_df = pd.DataFrame(dist, columns=sites["site_id"].tolist())
    dist_df.insert(0, "site_id", sites["site_id"])

    with pd.ExcelWriter(OUTPUT_XLSX, engine="openpyxl") as writer:
        out.to_excel(writer, index=False, sheet_name="sites")
        dist_df.to_excel(writer, index=False, sheet_name="dist_miles")


if __name__ == "__main__":
    main()