classes.py

from imports import pd, Enum
import datetime
import os
import requests
import json
import pytz
import time
import random

class State(Enum):
    UNKNOW = 'X'
    WAITING = 'Y'
    A = 'A'
    B = 'B'

class BusStop(object):
    def __init__(self, id, x, y, name, set_id):
        self.x = x
        self.y = y
        self.id = id
        self.name = name
        self.set_id = set_id

    def __str__(self):
        return f'{self.id} {self.name}'

class LineRoute(object):
    def __init__(self, row, bs):
        self.line = row["line"]
        self.track = row["track"].split()
        self.direction = row["direction"]
        self.type = row["type"]
        self.code = row["code"]
        self.level = row["level"]
        self.city_code_from = row["city_code_from"]
        self.city_code_to = row["city_code_to"]
        self.trackBS = self.create_busstops_array(bs)
        self.length = len(self.track)

    def __str__(self):
        return f'Line: {self.line} Code: {self.code} Dir: {self.direction}'

    def __len__(self):
        return self.length

    def create_busstops_array(self, dico):
        array = []
        for bs_id in self.track:
            try:
                array.append(dico[int(bs_id)])
            except:
                return None
        return array

    def get_index(self, stop_id:str) -> int:
        if stop_id >= self.length:
            return -1
        
        for i in range(self.length):
            if self.track[i] == stop_id:
                return i

    def to_df(self):
        return pd.DataFrame({
            'line': self.line,
            'track_code': self.code,
            'direction': self.direction,
            'type': self.type,
            'first_stop': self.track[0],
            'last_stop': self.track[-1],
            'priority': self.level,
            'direction': self.direction,
            'city_code_from': self.city_code_from,
            'city_code_to': self.city_code_to,
        }, index=[0])

class RideHistory(object):
    def __len__(self):
        return len(self.time)

    def __init__(self, vehicle=None):
        self.ride_id = None
        self.line = None
        self.direction_state = None
        self.direction = None
        self.vehicle_nr = None
        self.brigade = None
        self.course_name = None
        self.vehicle_id = None
        self.route:LineRoute = None
        self.ride:Ride = None
        self.ride_id_timetable:str = ''
        self.vehicle:Vehicle = None
        self.day = None
        if vehicle is not None:
            self.set_vehicle(vehicle)

        self.time = []
        self.time_str = []
        self.bus_stop_nr = []
        self.bus_stop_id = []
        self.delay = []
        self.is_on_time = []
        self.activity = []
        self.sorting = []

        self.correct_first_time = None
        

    def set_vehicle(self, vehicle):
        self.line = vehicle.line
        self.direction_state = vehicle.state
        self.direction = vehicle.state.value
        self.vehicle_nr = vehicle.veh_num
        self.brigade = vehicle.brigade
        self.course_name = vehicle.get_current_course_name()
        self.vehicle_id = vehicle.veh_id
        self.route = vehicle.route
        self.vehicle = vehicle
        self.day = vehicle.day

    def set_ride(self, ride):
        self.ride = ride
        self.ride_id_timetable = ride.ride_id
        self.course_name = ride.course_name
        self.route = ride.route

    def add_row(self, time:datetime.datetime):
        if self.direction != self.vehicle.state.value:
            raise("Vehicle in a different state than current history")
        bs_index = self.vehicle.next_stop_index
        next_bs_id = self.route.track[bs_index]
        # delay = None
        # if self.ride is not None:
        #     delay = self.ride.calculate_delay(bs_index, time)
        # is_on_time = 1 if delay == 0 else 0
        act = self.vehicle.get_activity()
        activity = act[0]
        sorting = act[1]
        next_bs_id = self.route.track[bs_index]
        
        self.activity.append(activity)
        self.time.append(time)
        # self.time_str.append(time.strftime("%Y-%m-%d %H:%M:%S"))
        self.bus_stop_id.append(next_bs_id)
        self.bus_stop_nr.append(bs_index)
        self.sorting.append(sorting)

        if len(self.time) == 3:
            self.ride_id = f'{self.vehicle_id}_{self.time[0].strftime("%y%m%d-%H:%M")}_{self.direction}'
        
        # if delay is None:
        #     self.is_on_time.append(-1)
        #     self.delay.append(0)
        # else:
        #     self.is_on_time.append(is_on_time)
        #     self.delay.append(delay)

    def correct_row(self, time:datetime.datetime):
        if len(self.time) == 0:
            self.correct_first_time = time
            return
        self.time[0] = time
        #self.time_str[0] = time.strftime("%Y-%m-%d %H:%M:%S")
        #self.set_delay_values(0)

    def set_delay_values(self, ind:int):
        if ind == 0 and len(self.delay):
            raise Exception("There was already some delay values! Whoa!")
        time = self.time[ind]
        delay = None
        if self.ride is not None:
            delay = self.ride.calculate_delay(ind, time)
        is_on_time = 1 if delay == 0 else 0
        if delay is None:
            self.is_on_time.append(-1)
            self.delay.append(0)
        else:
            self.is_on_time.append(is_on_time)
            self.delay.append(delay)
    
    def check_validity(self):
        if len(self.time) == 0:
            return False
            #self.day is not None and 
        if self.time[0].date() != self.day:
            return False
        if self.vehicle is None:
            return False
        if self.route is None:
            return False
        if len(self.time) != len(self.route):
            return False
        if len(self.bus_stop_id) != len(self.route):
            return False
        if len(self.bus_stop_nr) != len(self.route):
            return False
        if len(self.sorting) != len(self.route):
            return False
        if self.sorting[0] != 0:
            return False
        if self.sorting[-1] != 2:
            return False
        for i in range(1,len(self.time)):
            if (self.time[i] - self.time[i-1]).total_seconds()/60.0 >= 23:
                return False
        if self.direction != self.route.direction:
            raise Exception("Different direction of route and history")
        # if self.ride is not None:
        #     if not all(delay <= 300 and delay >= 0 for delay in self.delay):
        #         raise Exception("Wrong delay values")
        #     if not all(indic != -1 for indic in self.is_on_time):
        #         raise Exception("Wrong is_on_time indicators")
        return True

    def merge_if_continuous(self, old_hist):
        if self.direction != old_hist.direction:
            return False
        if self.route != old_hist.route:
            return False
        if self.course_name != old_hist.course_name:
            return False
        old_last_ind = old_hist.bus_stop_nr[-1] if len(old_hist) >= 1 else -100
        new_first_ind = self.bus_stop_nr[0] if len(self) >= 1 else -10

        if new_first_ind == old_last_ind + 1:
            self.time = old_hist.time + self.time
            self.time_str = old_hist.time_str + self.time_str
            self.bus_stop_nr = old_hist.bus_stop_nr + self.bus_stop_nr
            self.bus_stop_id = old_hist.bus_stop_id + self.bus_stop_id
            self.delay = old_hist.delay + self.delay
            self.is_on_time = old_hist.is_on_time + self.is_on_time
            self.activity = old_hist.activity + self.activity
            self.sorting = old_hist.sorting + self.sorting

            if self.correct_first_time is not None:
                self.time[0] = self.correct_first_time
            return True
        return False
        

    def get_df(self):
        self.ride_id = f'{self.vehicle_id}_{self.time[0].strftime("%y%m%d-%H:%M")}_{self.direction}'
        self.time_str = [time.strftime("%Y-%m-%d %H:%M:%S") for time in self.time]
        for i in range(len(self.time)):
            self.set_delay_values(i)
        if len(self.delay) != len(self.time):
            raise Exception("Something went wrong while calculating delays.")
        df = pd.DataFrame({
            'ride_id': self.ride_id,
            'time': self.time,
            'vehicle_id': self.vehicle_id,
            'line': self.line,
            'vehicle_number': self.vehicle_nr,
            'brigade': self.brigade,
            'bus_stop_id': self.bus_stop_id,
            'bus_stop': self.bus_stop_nr,
            'direction': self.direction,
            'is_on_time': self.is_on_time,
            'delay': self.delay,
            'ride_id_timetable': self.ride_id_timetable,
            'course_name': self.course_name,
            'activity': self.activity,
            'sorting': self.sorting,
            'time_str': self.time_str
        })
        df['bus_stop'] = df['bus_stop'].astype('int')
        df['is_on_time'] = df['is_on_time'].astype('int')
        df['sorting'] = df['sorting'].astype('int')
        return df

    def append_case(self, cases):
        cases['ride_id'].append(self.ride_id)
        cases['vehicle_id'].append(self.vehicle_id)
        cases['course_name'].append(self.course_name)
        cases['line'].append(self.line)
        cases['vehicle_number'].append(self.vehicle_nr)
        cases['brigade'].append(self.brigade)
        cases['direction'].append(self.direction)
        cases['start_date'].append(self.time[0].strftime('%Y-%m-%d'))
        cases['first_stop'].append(self.route.track[0])
        cases['last_stop'].append(self.route.track[-1])
        cases['ride_type'].append(self.ride.ride_type if self.ride is not None else '')
        cases['ride_day_type'].append(self.ride.day_type if self.ride is not None else '')
        cases['ride_is_on_time'].append(self.ride.is_on_time if self.ride is not None else -1)
        cases['ride_id_timetable'].append(self.ride.ride_id if self.ride is not None else '')

class Ride(object):
    def __init__(self, row, date, df_courses, vehicle):
        self.ride_id = row.ride_name
        self.course_name = row.course_name
        self.line = row.line_number
        self.direction = row.direction
        self.vehicle = vehicle

        self.ride_type = row.ride_type
        self.day_type = row.day_type
        self.times = []
        self.time_deltas = []
        self.is_on_time = True
        self.route = None

        df_ride = df_courses[df_courses['ride_name'] == row.ride_name]
        self.process_ride(df_ride, date)

    def set_route(self, route):
        self.route = route

    def process_ride(self, df_ride, date):
        first_time = date + df_ride.iloc[0].time_actual
        first_date = date

        first_sec = df_ride.iloc[0].time_actual.total_seconds()
        self.first_stop = df_ride.iloc[0].bus_stop_id
        self.last_stop = df_ride.iloc[-1].bus_stop_id
        for _, row in df_ride.iterrows():
            if row.time_actual.total_seconds() < first_sec:
                first_date = first_date + datetime.timedelta(days=1)
                first_sec = row.time_actual.total_seconds()
            curtime = first_date + row.time_actual
            self.times.append(curtime)
            self.time_deltas.append(int((curtime - first_time).seconds/60))
        return

    def calculate_delay(self, ind, timestamp):
        actual_timestamp = timestamp.to_pydatetime()
        desired_time = self.times[ind]
        delta = (actual_timestamp - desired_time)
        delay = round(delta.total_seconds()/60.0, 2) if delta.total_seconds() >= 60 else 0

        if delay > 1000:
            print()
        
        if delay != 0:
            self.is_on_time = False
        return delay

    def __str__(self):
        return f'Ride {self.ride_id} Line {self.line}'

class Vehicle(object):
    def __init__(self, line, vehNum, brigade, df_tracks_A, df_tracks_B):
        self.line = line
        self.veh_num = vehNum
        self.brigade = brigade
        self.veh_id = f'{line}_{vehNum}_{brigade}'
        self.state = State.UNKNOW
        self.last_state = State.UNKNOW
        self.last_stop_index = -1
        self.next_stop_index = 0
        self.routesA = df_tracks_A
        self.routesB = df_tracks_B
        self.search_for_first_stop = False
        self.day = None

        self.route:LineRoute = None
        self.ride:Ride = None
        self.history:RideHistory = RideHistory()

        self.rides = []
        self.histories = [self.history]

    def __str__(self):
        return f'Line:{self.line} number:{self.veh_num}'

    def name(self):
        return f'{self.line}-{self.veh_num}-{self.bridge}'

    def update_state(self, state:State, busstop:BusStop = None):
        self.last_state = self.state
        self.state = state
        self.ride = None
        if state == State.A:
            self.route = self.routesA[0]
        elif state == State.B:
            self.route = self.routesB[0]
        else:
            self.route = None

        if self.last_state in (State.A, State.B) and self.state == State.UNKNOW:
            self.history = RideHistory()
            #self.check_last_history_valid()
            self.histories.append(self.history)

        #Changing state from UNKNOWN
        if busstop is not None:
            #TODO: CHECK IF HISTORY'S LAST STOP AND NEXT IS ON ROUTE
            self.history.set_vehicle(self)
            # if len(self.histories) >= 2:
            #     self.history.merge_if_continuous(self.histories[-2])
            for i in range(0, len(self.route)):
                if self.route.trackBS[i].id == busstop.id:
                    self.last_stop_index = i - 1
                    self.next_stop_index = i
                    break

    def get_last_stop(self):
        if self.route is not None:
            return self.route.trackBS[self.last_stop_index]
        else: 
            raise("Vehicle doesn't have a course")

    def get_next_stop(self):
        if self.route is not None:
            if self.last_stop_index + 1 < len(self.route):
                return self.route.trackBS[self.last_stop_index + 1]
            else:
                return None
        else: 
            raise("Vehicle doesn't have a course")

    def achieve_next_stop(self):
        if self.route is not None:
            if self.last_stop_index < len(self.route) - 2:
                self.last_stop_index += 1
                self.next_stop_index += 1
            else:
                self.should_achieve_next_stop = True
                self.achieve_next_stop_indexes()
                self.update_course(None)
                self.history = RideHistory(self)
                #self.check_last_history_valid()
                self.histories.append(self.history)

    # def check_last_history_valid(self):
    #     if not self.histories[-1].check_validity():
    #         self.histories.pop()
                
    def achieve_next_stop_indexes(self):
        self.last_stop_index = -1
        self.next_stop_index = 0
        self.search_for_first_stop = True
        if self.state == State.A:
            self.update_state(State.B)
        elif self.state == State.B:
            self.update_state(State.A)

    def get_current_track(self):
        return self.route

    def get_current_course_name(self):
        return self.route.code

    def get_current_bus_stops_list(self):
        return self.route.trackBS

    def get_activity(self):
        track = self.get_current_bus_stops_list()
        if self.last_stop_index == len(track) - 2: #Last stop reached
            return (f'Ride ended', 2)
            return f'Ride ended at {self.get_next_stop().name}'
        elif self.last_stop_index == -1: #First stop reached
            return (f'Ride started', 0)
            return f'Ride started at {self.get_next_stop().name}'
        else:  #Regular stop reached
            return (f'Arrival at {self.get_next_stop().name}', 1)

    def search_for_ride(self, df_rides, timestamp):
        cur_bs = self.get_current_track().track[self.next_stop_index]
        datetime = timestamp.to_pydatetime()
        date = datetime.replace(hour=0, minute=0, second=0, microsecond=0)

        dir = self.state.value
        possible_rides = df_rides[(df_rides['line_number'] == self.line) & (df_rides['bus_stop_id'] == cur_bs) 
                    & (df_rides['direction'] == dir)]
        if possible_rides.empty:
            return
        else:
            possible_rides['time_difference'] = possible_rides.apply(lambda x: self.get_difference(x, datetime), axis = 1)
            possible_rides.sort_values(by=['time_difference'], ascending=True, inplace=True)
            ind = possible_rides['time_difference'].idxmin()
            if possible_rides.iloc[0].time_difference < 7:
                self.update_ride_and_course(Ride(possible_rides.loc[ind], date, df_rides, self))

    def update_ride_and_course(self, ride):
        self.ride = ride
        self.rides.append(ride)
        route = ride.course_name
        self.update_course(route)
        self.ride.set_route(self.route)
        self.history.set_ride(ride)
        return

    def update_course(self, course_name):
        if course_name is None:
            #reset at the end of course
            self.route = self.routesA[0] if self.state == State.A else self.routesB[0]
        else:
            #search for track in tracks
            if self.state == State.A:
                for track in self.routesA:
                    if track.code == course_name:
                        self.route = track
            elif self.state == State.B:
                for track in self.routesB:
                    if track.code == course_name:
                        self.route = track

    def get_current_bus_stops_list_string(self):
        return self.route.track

    def get_difference(self, row_from_timetable, gps_timestamp):
        date = gps_timestamp.replace(hour=0, minute=0, second=0, microsecond=0)
        second_stamp = date + row_from_timetable.time_actual
        delta = max(gps_timestamp, second_stamp) - min(gps_timestamp, second_stamp)
        return delta.seconds/60.0

    def get_result(self):
        valid_histories = [hist for hist in self.histories if hist.check_validity()]
        not_valid_histories = [hist for hist in self.histories if not hist.check_validity()]
        merged_histories = []
        i = 1
        while i < len(not_valid_histories):
            new_hist:RideHistory = not_valid_histories[i]
            old_hist:RideHistory = not_valid_histories[i - 1]
            if new_hist.merge_if_continuous(old_hist):
                if new_hist.check_validity():
                    merged_histories.append(new_hist)
                i += 2
                continue
            i += 1

        all_histories = valid_histories + merged_histories
        all_histories.sort(key=lambda x: x.time[0])
        if len(all_histories) == 0:
            return None
        frames = []
        cases = {
            'ride_id': [],
            'vehicle_id': [],
            'course_name': [],
            'line': [],
            'vehicle_number': [],
            'brigade': [],
            'direction': [],
            'start_date': [],
            'first_stop': [],
            'last_stop': [],
            'ride_type': [],
            'ride_day_type': [],
            'ride_is_on_time': [],
            'ride_id_timetable': []
        }
        for hist in all_histories:
            frames.append(hist.get_df())
            hist.append_case(cases)
        
        res_df = pd.concat(frames)
        cases_df = pd.DataFrame(data=cases)
        cases_df['ride_is_on_time'] = cases_df['ride_is_on_time'].astype('int')
        return (
            res_df, 
            cases_df, 
            len(all_histories), 
            len([hist for hist in all_histories if hist.ride is not None])
            )

class APIQuery(object):
    def __init__(self, stop:tuple, stop_id:str, day_date:datetime.datetime):
        self.stop_id = stop_id
        self.base_url = "https://history.openweathermap.org/data/2.5/history/city?"
        self.api_key = f'appid={os.getenv("WEATHER_API_KEY")}'
        self.lat = f'lat={stop[1]}'
        self.lon = f'lon={stop[0]}'
        self.type = f'type=hour'
        self.start = f'start={(datetime.datetime(year=day_date.year, month=day_date.month, day=day_date.day) + datetime.timedelta(hours=-4)).timestamp()}'
        self.cnt = f'cnt=48'
        self.units = f'units=metric'

    def get_args(self):
        return '&'.join([
            self.api_key,
            self.lat,
            self.lon,
            self.type,
            self.units,
            self.start,
            self.cnt
            ])

    def get_data(self):
        complete_url = self.base_url + self.get_args()
        i = 0
        while i < 6:
            try:
                time.sleep(random.uniform(0.0, 5.0))
                response = requests.get(complete_url, timeout=5)
                res_json = response.json()

                if res_json["cod"] == '200':
                    return res_json["list"]
                else:
                    raise Exception("Wrong response code")
            except Exception as e:
                #print(f'API query for {self.stop_id} timed out, try: {i}, reason:{e}')
                i += 1
                time.sleep(random.uniform(20, 50))
        raise Exception(f"Couldn't get weather data for {self.stop_id}")

class WeatherHistory(object):
    def __init__(self, query:APIQuery):
        self.stop_id = query.stop_id
        self.data = []
        self.prepare_obj(query.get_data())

    def prepare_obj(self, items:list):
        for item in items:
            self.data.append(WeatherRow(item))

    def get_row(self, stamp:datetime.datetime):
        for row in self.data:
            if stamp >= row.start and stamp < row.end:
                return row
        raise Exception("No WeatherRow found!")

class WeatherRow(object):
    def __init__(self, item):
        self.start = datetime.datetime.fromtimestamp(item['dt'])
        self.end = self.start + datetime.timedelta(hours=1)
        self.temperature = item['main']['temp']
        self.humidity = item['main']['humidity']
        self.wind_speed = round(item['wind']['speed'] * 3.6, 2) if 'wind' in item.keys() else 0
        self.rain = item['rain']['1h'] if 'rain' in item.keys() else 0
        self.snow = item['snow']['1h'] if 'snow' in item.keys() else 0
        weather = item['weather'] if 'weather' in item.keys() else []
        self.weather_info = ','.join(wet['main'] for wet in weather)