import numpy as np
from sklearn.base import BaseEstimator, TransformerMixin


class CalculateDistanceTransformer(BaseEstimator, TransformerMixin):
    """
    Compute the great-circle (Haversine) distance between two lat/long coordinates.

    Parameters
    ----------
    lat_1 : str
        Name of the first latitude column.
    lon_1 : str
        Name of the first longitude column.
    lat_1 : str
        Name of the second latitude column.
    lon_2 : str
        Name of the second longitude column.
    new_feature_name : str, default="distance"
        Name of the generated distance column.
    """

    def __init__(self, lat_1, lon_1, lat_2, lon_2, new_feature_name="distance_km"):
        self.lat_1 = lat_1
        self.lon_1 = lon_1
        self.lat_2 = lat_2
        self.lon_2 = lon_2
        self.new_feature_name = new_feature_name

    def fit(self, X, y=None):
        return self

    def transform(self, X):
        X = X.copy()

        dep_lat_rad = np.radians(X[self.lat_1])
        dep_lon_rad = np.radians(X[self.lon_1])
        arr_lat_rad = np.radians(X[self.lat_2])
        arr_lon_rad = np.radians(X[self.lon_2])

        lat_distance = arr_lat_rad - dep_lat_rad
        lon_distance = arr_lon_rad - dep_lon_rad

        a = np.sin(lat_distance / 2) ** 2 + np.cos(dep_lat_rad) * np.cos(arr_lat_rad) * np.sin(lon_distance / 2) ** 2
        c = 2 * np.arctan2(np.sqrt(a), np.sqrt(1 - a))
        distance = 6371  * c

        X[self.new_feature_name] = distance
        return X
