import numpy as np
import pandas as pd

import trilateration


def process(dataset, cells, k):
    """Prepare data for training.

    :param dataset: parsed measurement data
    :param cells: cell tower database
    :param k: desired number of records per one measurement
    :return: processed dataset
    """

    data = dataset.copy()

    # Drop irrelevant columns.
    data.drop(columns=["RAT", "ARFCN", "ASU", "TA", "BSIC"], inplace=True)

    # Remove rows with missing identification of the cell.
    data.dropna(inplace=True)

    # Take relevant columns from the cell tower database and rename them.
    cells = cells[["cell", "mcc", "net", "area", "lat", "lon"]]
    cells.columns = ["cell_id", "mcc", "mnc", "lac", "cell_lat", "cell_lon"]

    # Left outer join of 'data' and 'cells' effectively adding GPS locations
    # of the cells.
    data = data.merge(cells, how='left', on=["cell_id", "mcc", "mnc", "lac"])

    # Drop rows with missing latitude and longitude information.
    data.dropna(inplace=True)

    # Zip latitude and longitude columns into tuples (lat, lon).
    data["meas_coord"] = list(zip(data.meas_lat, data.meas_lon))
    data["cell_coord"] = list(zip(data.cell_lat, data.cell_lon))

    data.drop(columns=["meas_lat", "meas_lon", "cell_lat", "cell_lon"],
              inplace=True)

    # Calculate great circle distance from cell tower to measurement location.
    data["distance"] = np.vectorize(trilateration.haversine_distance)\
                                   (data.meas_coord, data.cell_coord)

    # Make all measurements length k - drop shorter ones, shorten longer ones.
    # Take values closest to the measurement location as these are less prone to
    # noise. Then sort by signal strength as it is in original data.
    data = data.groupby("meas_id")\
               .filter(lambda x: len(x) >= k)\
               .groupby("meas_id")\
               .apply(lambda x: x.sort_values("distance")\
                                 .head(k)\
                                 .sort_values("signal"))\
               .reset_index(drop=True)

    # Reshape data to one line per measurement.
    columns = data.columns.values
    data = data.groupby('meas_id')\
             .apply(lambda x: x.values[:,:].reshape(1,-1)[0])\
             .apply(pd.Series)\
             .reset_index(drop=True)
    data.columns = ['{}_{}'.format(col, ix) for ix in range(k) for col in columns]

    return data
