# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Initial experiments
#
# In this notebook we perform some experimentations to understand the behaviour of the trajectory recovery algorithm as explained in the studied paper.

import itertools
import numpy as  np

# +
N = 10
M = 3

towers = np.array(['a', 'b', 'c'])
towers_index = {
    'a': 0,
    'b': 1,
    'c': 2
}
# -

# ## Towers users per time slot
#
# Specifies for each tower the number of users a the moments t0, t1 and t2

P = np.array([
    [3, 2, 5],  # P_t0
    [1, 3, 6],  # P_t1
    [4, 2, 4],  # P_t2
])

# ## Building distribution matrix
#
# Generates the users distribution matrix for each defined time period in P

# +
L = []

for Pt in P:
    L.append(
        np.array(list(
            itertools.chain(*(
                [towers[tower_index]] * count
                for tower_index, count in enumerate(Pt)
            ))
        )
    ))

L
# -

# ## Build random trajectories
#
# Generates the recovered trajectory for the users like a permutation of the users distribution (L) until the time t1

# +
S_t1 = np.array([
    np.random.permutation(Li)
    for Li in L[:2]
]).T

S_t1
# -

# # Night scenario
#
# ## Calculate possible next stops for each user
#
# Calculate next estimate as the last tower the users were until time t1

# +
L_t2_est = S_t1[:,1]

L_t2_est
# -

# Night probability of users switching between towers

# PROBABILITIES ARE RAW DISTANCES
prob_night = np.array([
    [0, 4, 3],
    [4, 0, 5],
    [3, 5, 0],
])

# Convert tower letters into numbers to later access the prob_night matrix easier

L_t2_est_index = [towers_index[t] for t in L_t2_est]
L_t2_index = [towers_index[t] for t in L[2]]

# Estimated tower index for users at time t2 and users distribution for time t2

list(zip(L_t2_est_index, L_t2_index))

# ## Build cost matrix
#
# Calculates the cost matrix at t1 as the probability of jumping between towers at night, taking into account the estimated towers for the time t2

C_t1 = np.zeros((N, N))

for i, l_t2_est in enumerate(L_t2_est_index):
    for j, l_t2 in enumerate(L_t2_index):
        C_t1[i, j] = prob_night[l_t2_est, l_t2]

C_t1

# +
from scipy.optimize import linear_sum_assignment

# Ref: https://docs.scipy.org/doc/scipy-0.18.1/reference/generated/scipy.optimize.linear_sum_assignment.html
row_ind, col_ind = linear_sum_assignment(C_t1)

list(zip(row_ind, col_ind))
# -

# # Daylight scenario
#
# Try building a cost matrix using the daylight probabilities specifications



StarcoderdataJupyterScriptsDedupFiltered