import math
import xlrd
import itertools as it
import numpy as np
import pandas as pd
from datetime import datetime

dfincidents_0 = pd.read_excel('CCL EVAC E-brake occurrences hourly update_mod.xlsx', sheetname='Aug Sep')
dfincidents_1 = pd.read_excel('CCL EVAC E-brake occurrences hourly update_mod.xlsx', sheetname='Nov')

# Incident data for Nov had different format
dfincidents_1['Time'] = dfincidents_1['Time'].str.strip('hrs').str.strip(' ')
dfincidents_1['Time'] = pd.to_datetime(dfincidents_1['Time'], format='%H%M').dt.time

dfincidents = pd.concat([dfincidents_0, dfincidents_1])

# Reset the index because they were concatenated from two data sources
dfincidents.reset_index(inplace=True, drop=True)

def datetime_from_date_and_time(row):
    """
    Combines the date column and time column into a single column
    """
    d = row['Date']
    t = row['Time']

    return datetime(
        d.year, d.month, d.day,
        t.hour, t.minute, t.second
    )

# Add DateTime to the data for easier visualization
dfincidents['DateTime'] = dfincidents.apply(datetime_from_date_and_time, axis=1)

stations=("MRB,BFT,DBG,BBS,EPN,PMN,NCH,SDM,MBT,DKT,PYL,MPS,TSG,BLY,SER,"
          "LRC,BSH,MRM,CDT,BTN,FRR,HLV,BNV,ONH,KRG,HPV,PPJ,LBD,TLB,HBF").split(',')

def loc_id(station1, station2 = None):
    """
    Translates a 3-letter station code to a number,
    or a pair of 3-letter station codes to a number.
    Single stations are represented as whole numbers.
    Locations between stations are represented with a .5.

    Example:
    loc_id('MRB')         # 0 (Marina Bay)
    loc_id('MRB', 'BFT')  # 0.5 (Between Marina Bay and Bayfront)
    loc_id('DBG')         # 2 (Dhoby Ghaut)
    loc_id('HBF')         # 29
    loc_id('HBF', 'TLB')  # 28.5 (Between Harbourfront and Telok Blangah)
    loc_id('HBF', 'DBG')  # throws and error, because these stations are not adjacent
    """
    if station2 == None or station2 == 'nan' or (type(station2) is float and math.isnan(station2)):
        # Single stations
        return stations.index(station1)

    else: # Pairs of stations -- take the average to get the 0.5
        stn1_index = stations.index(station1)
        stn2_index = stations.index(station2)

        # Handle the branch at Promenade
        if (set(['PMN', 'EPN']) == set([station1, station2])):
            return float(stations.index('EPN')) + 0.5
        elif set(['PMN', 'BFT']) == set([station1, station2]):
            return float(stations.index('BFT')) + 0.5
        else:
            # Require station pairs to be adjacent stations
            assert(math.fabs(stn1_index - stn2_index) == 1)
            return float(stn1_index + stn2_index) / 2

def loc_id_from_stations(row):
    try:
        # This handles entries with both "Station from" and "Station to"
        # and entries with only "Station from"
        return loc_id(row['Station from'], str(row['Station to']))
    except ValueError:
        # Some entries only have "Station to" but no
        # "Station from"
        return loc_id(row['Station to'])

# Select only some columns that we are interested in
sel_dfincidents = dfincidents[['DateTime', 'PV', 'Bound', 'Station from', 'Station to', 'Event', 'Remarks']]

# Add the location ID into the dataset
sel_dfincidents['LocID'] = sel_dfincidents.apply(loc_id_from_stations, axis=1)

def same_cascade(i, j):
    """
    Given a pair of incidents (i,j), returns true if:

        t <= d * 4 mins

    where t is  the time difference between occurrences
    and d is the distance (measured by difference in location ID).

    Moreover, we consider the track direction, and only consider
    incidents that are "moving backwards".
    """

    # If trains are not travelling in the same direction
    # they cannot be due to the same "backward moving" interference
    # source.
    # (Note: This was the hypothesis when this code was written.
    # It turned out that the rogue train could affect all
    # trains in the vicinity, not just in the opposite track)
    if i["Bound"] != j["Bound"] or \
        i["Bound"] not in ['IT', 'OT']:
        return False

    # time difference in minutes
    time_difference = (i["DateTime"] - j["DateTime"]) / np.timedelta64(1, 'm')
    location_difference = i["LocID"] - j["LocID"]

    if location_difference == 0:
        return False

    ratio = time_difference / location_difference

    if i["Bound"] == 'OT':
        return ratio > 0 and ratio < 4
    elif i["Bound"] == 'IT':
        return ratio < 0 and ratio > -4

incidents = sel_dfincidents.to_records()
# (a, b, c, d, ...) --> ((a,b), (a,c), (a,d), ..., (b,c), (b,d), ..., (c,d), ...)
incident_pairs = list(it.combinations(incidents, 2))

related_pairs = [ip for ip in incident_pairs if same_cascade(*ip)]
related_pairs = [(i[0], j[0]) for i,j in related_pairs]

def pairs_to_clusters(pairs):
    """
    A quick-and-dirty disjoint-set data structure. But this works fast enough for 200+ records
    Could be better.

    Example input:
    (1,2), (2,3), (4,5)

    Output:
    1: {1,2,3}
    2: {1,2,3}
    3: {1,2,3}
    4: {4,5}
    5: {4,5}
    """
    the_clusters = dict()

    for i,j in pairs:
        if i not in the_clusters:
            if j in the_clusters:
                the_clusters[j].add(i)
                the_clusters[i] = the_clusters[j]
            else:
                the_clusters[i] = set(list([i, j]))
                the_clusters[j] = the_clusters[i]
        else:
            if j in the_clusters:
                if the_clusters[i] is not the_clusters[j]: # union the two sets
                    for k in the_clusters[j]:
                        the_clusters[i].add(k)
                        the_clusters[k] = the_clusters[i]
                else: # they are already in the same set
                    pass
            else:
                the_clusters[i].add(j)
                the_clusters[j] = the_clusters[i]

    return the_clusters

clusters = pairs_to_clusters(related_pairs)
# Show each set only once
clusters = [v for k,v in clusters.items() if min(v) == k]
clusters[0:10]
view raw

[{0, 1},
 {2, 4},
 {5, 6, 7},
 {8, 9},
 {18, 19, 20},
 {21, 22, 24, 26, 27},
 {28, 29, 30, 31, 32, 33, 34},
 {42, 44, 45},
 {47, 48},
 {51, 52, 53, 56}]

# count % of incidents occurring in a cluster
all_clustered_incidents = set()

for i,clust in enumerate(clusters):
    all_clustered_incidents |= clust

(len(all_clustered_incidents),
 len(incidents),
 float(len(all_clustered_incidents)) / len(incidents))

(189, 259, 0.7297297297297297)