import os
import pandas as pd
from numpy import NaN
from datetime import datetime
import geopandas as gpd
from shapely.geometry import LineString
from .gtfs_schema import GTFS_Schema
import time


class Service_Utils(object):
    """
    Main class to store GTFS data and provde helper
    functions.
    """

    def __init__(self, gtfs_dir, service_date):
        """
        Instantiate class with directory of GTFS Files and a service_date for
        which to get service.
        """
        self.gtfs_dir = gtfs_dir
        self.service_date = service_date
        self.int_service_date = int(service_date)
        self._crs_epsg = 4326

        # gtfs properties:
        self.calendar_dates = self.__get_calendar_dates()
        self.calendar = self.__get_calendar()
        self.service_ids = self.__get_service_ids()
        self.trips = self.__get_trips()
        self.stop_times = self.__get_stop_times()
        # deal with frequencies here:
        if os.path.exists(os.path.join(self.gtfs_dir, 'frequencies.txt')):
            self.trips, self.stop_times = self.frequencies_to_trips()

        # self._df_all_stops_by_trips = self.__get_trips_stop_times()
        self.routes = self.__get_routes()
        self.stop_list = self.stop_times['stop_id'].unique()
        self.stops = self.__get_stops()
        self.shapes = self.__get_shapes()

        # derived DataFrames
        self._df_all_stops_by_trips = self.__get_trips_stop_times()
        self._schedule_pattern_dict = self.__get_schedule_pattern()
        self.schedule_pattern_df = self.__get_schedule_pattern_df()
        self._df_all_stops_by_trips = self._df_all_stops_by_trips.merge(
            self.schedule_pattern_df[['orig_trip_id', 'rep_trip_id']],
            how='left', left_on='trip_id',
            right_on='orig_trip_id').drop(columns=['orig_trip_id'])
        self.__trip_list = self.__get_trip_list()
        self.__rep_trip_list = list(
            self.schedule_pattern_df.rep_trip_id.unique())

    def __get_calendar(self):
        """
        Returns calendar.txt as a DataFrame.
        """
        calendar_df = pd.read_csv(os.path.join(self.gtfs_dir, 'calendar.txt'))
        return GTFS_Schema.Calendar.validate(calendar_df)

    def __get_calendar_dates(self):
        """
        Returns calendar_dates.txt as a DataFrame.
        """
        if os.path.exists(os.path.join(
                self.gtfs_dir, 'calendar_dates.txt')) is False:
            calendar_dates = pd.DataFrame(
                    columns=GTFS_Schema.calendar_dates_columns)
        else:
            calendar_dates = pd.read_csv(
                os.path.join(self.gtfs_dir, 'calendar_dates.txt'))

        return GTFS_Schema.Calendar_Dates.validate(calendar_dates)

    def __get_trips(self):
        """
        Gets records in trips.txt for the service_ids that represent
        the service_date passed into the constructor. Returns a DataFrame.
        """
        trips_df = pd.read_csv(os.path.join(self.gtfs_dir, 'trips.txt'))
        trips_df = GTFS_Schema.Trips.validate(trips_df)
        trips_df = trips_df[trips_df['service_id'].isin(self.service_ids)]

        return trips_df

    def __get_routes(self) -> pd.DataFrame:
        """
        Gets records in routes.txt for the trips that represent the
        service_date passed into the constructor. Returns a DataFrame.
        """
        routes_df = pd.read_csv(os.path.join(self.gtfs_dir, 'routes.txt'))
        routes_df = GTFS_Schema.Routes.validate(routes_df)
        routes_df = routes_df[routes_df['route_id'].isin(
            self.trips['route_id'])]
        return routes_df

    def __get_stop_times(self):
        """
        Gets records in stop_times.txt for the trips that represent
        the service_date passed into the constructor. Returns a DataFrame.
        """
        stop_times_df = pd.read_csv(os.path.join(
            self.gtfs_dir, 'stop_times.txt'))
        stop_times_df = GTFS_Schema.Stop_Times.validate(stop_times_df)
        stop_times_df = stop_times_df[stop_times_df['trip_id'].isin(
            self.trips['trip_id'])]
        return stop_times_df

    def __get_stops(self):
        """
        Gets records in stops.txt for the stops used by trips represented in
        the service_date passed into the constructor. Returns a DataFrame.
        """
        stops_gdf = pd.read_csv(os.path.join(self.gtfs_dir, 'stops.txt'))
        stops_gdf = GTFS_Schema.Stops.validate(stops_gdf)
        stops_gdf = stops_gdf[stops_gdf['stop_id'].isin(self.stop_list)]
        stops_gdf = gpd.GeoDataFrame(stops_gdf, geometry=gpd.points_from_xy(
            stops_gdf['stop_lon'], stops_gdf['stop_lat']))
        stops_gdf = stops_gdf.set_crs(epsg=self._crs_epsg)
        return stops_gdf

    def __get_shapes(self):
        """
        Gets records in shapes.txt for the shape_ids in trips that represent
        the service_date passed into the constructor. The sequence of points
        are converted to line geometry and returned as a GeoDataFrame.
        """
        gdf = pd.read_csv(os.path.join(self.gtfs_dir, 'shapes.txt'), dtype={
            'shape_id': str})
        gdf = GTFS_Schema.Shapes.validate(gdf)
        gdf = gdf[gdf['shape_id'].isin(self.trips['shape_id'])]
        gdf = gpd.GeoDataFrame(gdf, geometry=gpd.points_from_xy(
            gdf['shape_pt_lon'], gdf['shape_pt_lat']))
        gdf = gpd.GeoDataFrame(gdf.groupby('shape_id')['geometry'].apply(
            lambda x: LineString(x.tolist())))
        gdf.reset_index(inplace=True)
        gdf = gdf.set_crs(epsg=self._crs_epsg)
        return gdf

    def __get_service_ids(self):
        """
        Returns a list of valid service_id(s) from each feed using the user
        specified service_date.
        """
        my_date = datetime(int(self.service_date[0:4]), int(
            self.service_date[4:6]), int(self.service_date[6:8]))
        day_of_week = self.__get_weekday(my_date)

        regular_service_dates = self.calendar[(
            self.calendar['start_date'] <= self.int_service_date) & (
                self.calendar['end_date'] >= self.int_service_date) & (
                    self.calendar[day_of_week] == 1)]['service_id'].tolist()

        exceptions_df = self.calendar_dates[
            self.calendar_dates['date'] == self.int_service_date]

        add_service = exceptions_df.loc[
            exceptions_df['exception_type'] == 1]['service_id'].tolist()

        remove_service = exceptions_df[
            exceptions_df['exception_type'] == 2]['service_id'].tolist()

        service_id_list = [x for x in (
            add_service + regular_service_dates) if x not in remove_service]

        assert service_id_list, "No service found in feed."
        return service_id_list

    def __get_weekday(self, my_date):
        """
        Gets the day of week from user parameter service date.
        """
        week_days = ['monday', 'tuesday', 'wednesday',
                     'thursday', 'friday', 'saturday', 'sunday']
        return week_days[my_date.weekday()]

    def __make_sequence_col(
            self, data_frame, sort_list, group_by_col, seq_col):

        """
        Sorts a pandas dataframe using sort_list, then creates a column of
        sequential integers (1,2,3, etc.) for groups, using the group_by_col.
        Then drops the existing sequence column and re-names the new sequence
        column.
        """
        # sort on tripId, sequence
        data_frame = data_frame.sort_values(sort_list, ascending=[1, 1])
        # create a new field, set = to the position of each record in a group,
        # grouped by tripId
        data_frame['temp'] = data_frame.groupby(group_by_col).cumcount() + 1
        # drop the old sequence column
        data_frame = data_frame.drop(columns=[seq_col], axis=1)
        # rename new column:
        data_frame = data_frame.rename(columns={'temp': seq_col})

        return data_frame

    def __get_trip_list(self):
        """
        Returns a list of trip_ids used by the service_date passed into the
        constructor.
        """
        trips = self._df_all_stops_by_trips.trip_id.tolist()
        trips = list(set(trips))
        return trips

    def __get_trips_stop_times(self):
        """
        Creates a merged dataframe consisting of trips & stop_ids for the
        start time, end time and service_id (from GTFS Calender.txt). This
        can include partial itineraries as only stops within the start and
        end time are included.
        """
        stop_times_df = self.__make_sequence_col(
            self.stop_times, ['trip_id', 'stop_sequence'],
            'trip_id', 'stop_sequence')

        stop_times_df = stop_times_df[stop_times_df['trip_id'].isin(
            self.trips['trip_id'])]

        # Add columns for arrival/departure in decimal minutes and hours:
        # Some schedules only have arrival/departure times for time points,
        # not all stops:
        if stop_times_df['departure_time'].isnull().any():
            stop_times_df['departure_time'].fillna('00:00:00', inplace=True)
            stop_times_df['departure_time_mins'] = stop_times_df.apply(
                self.__convert_to_decimal_minutes, axis=1, args=(
                    'departure_time',))
            stop_times_df['departure_time_mins'].replace(0, NaN, inplace=True)
            stop_times_df['departure_time_mins'].interpolate(inplace=True)
        else:
            stop_times_df['departure_time_mins'] = stop_times_df.apply(
                self.__convert_to_decimal_minutes, axis=1, args=(
                    'departure_time',))
        stop_times_df['departure_time_hrs'] = stop_times_df[
            'departure_time_mins']/60
        stop_times_df['departure_time_hrs'] = stop_times_df[
            'departure_time_hrs'].astype(int)
        stop_times_df = stop_times_df.merge(
            self.trips, 'left', left_on=["trip_id"], right_on=["trip_id"])

        return stop_times_df

    def __convert_to_decimal_minutes(self, row, field):
        """
        Convert HH:MM:SS to seconds since midnight
        for comparison purposes.
        """

        H, M, S = row[field].split(':')
        seconds = float(((float(H) * 3600) + (float(M) * 60) + float(S))/60)

        return seconds

    def __convert_to_seconds(self, row, field):
        """
        Converts from hhmmss format to seconds.
        """
        h, m, s = row[field].split(':')
        return int(h) * 3600 + int(m) * 60 + int(s)

    def __to_hhmmss(self, row, field):
        """
        Converts to hhmmss format.
        """
        return time.strftime('%H:%M:%S', time.gmtime(row[field]))

    def __get_schedule_pattern(self, route_field='route_id'):
        """
        Returns a nested diciontary where the first level key is route_id and
        values are respresentative trip_ids that have unique stop sequences.
        These are are used as keys for the second level where each value is a
        dictinary that includes a list of trips_id's that share this stop
        pattern and a list of ordered stops.

        {route_id : trip_id {trips_ids : [list of trip ids], stops :
        [list of stops]}}

        """
        stop_sequence_dict = {k: list(v) for k, v in
                              self._df_all_stops_by_trips.groupby(
                ['trip_id', route_field])['stop_id']}

        # Empty dictionary to store unique stop sequences
        my_dict = {}
        for key, value in stop_sequence_dict.items():
            trip_id = key[0]
            route_id = key[1]
            # Handle some branching later on
            found = False
            # If this is the first trip for this route, just add it
            if route_id not in my_dict.keys():
                my_dict[route_id] = {trip_id: {
                    'stops': value, 'trip_ids': [trip_id]}}
            # Otherwise check to see if this stop sequence has already been
            # added for this route
            else:
                for k, v in my_dict[route_id].items():
                    if value == v['stops']:
                        # This stop sequence has already been added for this
                        # route, add the trip_id to the list of trip_ids that
                        # have this sequence in common.
                        my_dict[route_id][k]['trip_ids'].append(trip_id)
                        found = True
                        break
                if not found:
                    # Add the stop sequence and route, trip info
                    my_dict[route_id][trip_id] = {
                        'stops': value, 'trip_ids': [trip_id]}
            # Set back to False for next iteration
            found = False
        return my_dict

    def __get_schedule_pattern_df(self):
        """
        Returns a DataFrame with a field for each trip_id used
        to represent a unique stop_pattern (rep_trip_id) and a
        column with the other trip_ids that share the same stop
        pattern
        """
        rows = []
        for route_id, trips in self._schedule_pattern_dict.items():
            for trip_id, data in trips.items():
                for trip in data['trip_ids']:
                    rows.append(
                        {'route_id': route_id, 'trip_id1': trip_id,
                            'trip_id2': trip})
        df2 = pd.DataFrame(rows)
        df = self._df_all_stops_by_trips.drop_duplicates(['trip_id'])
        df2 = df2.merge(df[['trip_id', 'shape_id']], how='right',
                        left_on=["trip_id2"], right_on=["trip_id"])
        df2 = df2.rename(
            columns={'trip_id1': 'rep_trip_id', 'trip_id': 'orig_trip_id'})
        df2 = df2.drop(columns=['trip_id2'])
        return df2

    def frequencies_to_trips(self):
        """
        For each trip_id in frequencies.txt, calculates the number
        of trips and creates records for each trip in trips.txt and
        stop_times.txt. Deletes the original represetative trip_id
        in both of these files. 
        """
        frequencies = pd.read_csv(
            os.path.join(self.gtfs_dir, 'frequencies.txt'))
        frequencies = frequencies[frequencies['trip_id'].isin(
            self.trips['trip_id'])]

        # some feeds will use the same trip_id for multiple rows
        # need to create a unique id for each row
        frequencies['frequency_id'] = frequencies.index

        frequencies['start_time_secs'] = frequencies.apply(
            self.__convert_to_seconds, axis=1, args=('start_time',))

        frequencies['end_time_secs'] = frequencies.apply(
            self.__convert_to_seconds, axis=1, args=('end_time',))

        # following is coded so the total number of trips
        # does not include a final one that leaves the first
        # stop at end_time in frequencies. I think this is the
        # correct interpredtation of the field description:
        # 'Time at which service changes to a different headway
        # (or ceases) at the first stop in the trip.'

        # Rounding total trips to make sure all trips are counted
        # when end time is in the following format: 14:59:59,
        # instead of 15:00:00.

        frequencies['total_trips'] = ((
            (frequencies['end_time_secs']-frequencies[
                'start_time_secs']) / frequencies[
                    'headway_secs'])).round(0).astype(int)

        trips_update = self.trips.merge(frequencies, on='trip_id')
        trips_update = trips_update.loc[trips_update.index.repeat(
            trips_update['total_trips'])].reset_index(drop=True)
        trips_update['counter'] = trips_update.groupby(
            'trip_id').cumcount() + 1
        trips_update['trip_id'] = trips_update['trip_id'].astype(
            str) + '_' + trips_update['counter'].astype(str)

        stop_times_update = frequencies.merge(
            self.stop_times, on='trip_id', how='left')

        stop_times_update['arrival_time_secs'] = stop_times_update.apply(
            self.__convert_to_seconds, axis=1, args=('arrival_time',))
        stop_times_update['departure_time_secs'] = stop_times_update.apply(
            self.__convert_to_seconds, axis=1, args=('departure_time',))

        stop_times_update['elapsed_time'] = stop_times_update.groupby(
            ['trip_id', 'start_time'])['arrival_time_secs'].transform('first')
        stop_times_update['elapsed_time'] = stop_times_update[
            'arrival_time_secs'] - stop_times_update['elapsed_time']
        stop_times_update['arrival_time_secs'] = stop_times_update[
            'start_time_secs'] + stop_times_update['elapsed_time']

        # for now assume departure time is the same as arrival time.
        stop_times_update['departure_time_secs'] = stop_times_update[
            'start_time_secs'] + stop_times_update['elapsed_time']

        stop_times_update = stop_times_update.loc[
            stop_times_update.index.repeat(
                stop_times_update['total_trips'])].reset_index(drop=True)

        # handles case of repeated trip_ids
        stop_times_update['counter'] = stop_times_update.groupby(
            ['frequency_id', 'stop_id']).cumcount()
        stop_times_update['departure_time_secs'] = stop_times_update[
            'departure_time_secs'] + (stop_times_update[
                'counter'] * stop_times_update['headway_secs'])
        stop_times_update['arrival_time_secs'] = stop_times_update[
            'arrival_time_secs'] + (stop_times_update[
                'counter'] * stop_times_update['headway_secs'])

        # now we want to get the cumcount based on trip_id
        stop_times_update['counter'] = stop_times_update.groupby(
            ['trip_id', 'stop_id']).cumcount() + 1
        stop_times_update['departure_time'] = stop_times_update.apply(
            self.__to_hhmmss, axis=1, args=('departure_time_secs',))
        stop_times_update['arrival_time'] = stop_times_update.apply(
            self.__to_hhmmss, axis=1, args=('arrival_time_secs',))
        stop_times_update['trip_id'] = stop_times_update[
            'trip_id'].astype(str) + '_' + stop_times_update[
                'counter'].astype(str)

        # remove trip_ids that are in frequencies
        stop_times = self.stop_times[~self.stop_times['trip_id'].isin(
            frequencies['trip_id'])]

        trips = self.trips[~self.trips['trip_id'].isin(frequencies['trip_id'])]

        # get rid of some columns
        stop_times_update = stop_times_update[stop_times.columns]
        trips_update = trips_update[trips.columns]

        # add new trips/stop times
        trips = pd.concat([trips, trips_update])
        stop_times = pd.concat([stop_times, stop_times_update])

        return trips, stop_times

    def get_tph_by_line(self):
        """
        Returns a DataFrame with records for each rep_trip_id and
        columns with the number of trips for each hour after midnight
        with service. For example 2:00-3:00 AM is called hour_2 and
        3:00-4:00 PM is called hour_15.
        """
        # get the first stop for every trip
        first_departure = self._df_all_stops_by_trips.sort_values(
            'stop_sequence', ascending=True).groupby(
                'trip_id', as_index=False).first()
        # this may not be necessary
        first_departure = first_departure.loc[(
            first_departure.stop_sequence == 1)]
        first_departure = first_departure.groupby(
            ['rep_trip_id', 'departure_time_hrs']
            )['departure_time_hrs'].count()
        first_departure_df = pd.DataFrame(first_departure)
        first_departure_df.reset_index(level=0, inplace=True)
        first_departure_df = first_departure_df.rename(
            columns={'departure_time_hrs': 'frequency'})
        first_departure_df.reset_index(level=0, inplace=True)
        t = pd.pivot_table(
            first_departure_df, values='frequency', index=['rep_trip_id'],
            columns=['departure_time_hrs'])
        t = t.fillna(0)
        for col in t.columns:
            if not col == 'rep_trip_id':
                t = t.rename(columns={col: 'hour_' + str(col)})
        t.reset_index(inplace=True)
        t = t.merge(self.trips[['route_id', 'trip_id',
                    'direction_id']], how='left', left_on='rep_trip_id',
                    right_on='trip_id')
        t.drop(columns=['trip_id'], axis=1, inplace=True)

        return t

    def get_tph_at_stops(self):
        """
        Returns a DataFrame with records for each stop_id and
        columns with the number of trips for each hour after midnight
        with service. For example 2:00-3:00 AM is called hour_2 and
        3:00-4:00 PM is called hour_15.
        """
        df = self._df_all_stops_by_trips.groupby(
            ['stop_id', 'departure_time_hrs'])['departure_time_hrs'].count()

        df = pd.DataFrame(df)
        df.reset_index(level=0, inplace=True)
        df = df.rename(columns={'departure_time_hrs': 'frequency'})
        df.reset_index(level=0, inplace=True)
        t = pd.pivot_table(df, values='frequency', index=['stop_id'], columns=[
            'departure_time_hrs'])
        t = t.fillna(0)
        for col in t.columns:
            if not col == 'rep_trip_id':
                t = t.rename(columns={col: 'hour_' + str(col)})
        t.reset_index(inplace=True)
        return t

    def get_lines_gdf(self):
        """
        Returns a GeoDataFrame with records for each rep_trip_id and
        line geomery for the shape_id used by the trip_id. Useful GTFS
        columns inlcude route_id, direction_id, route_type, route_short_name,
        route_long_name, and route_desc.
        """
        rep_trips = self.trips[self.trips['trip_id'].isin(
            self.__rep_trip_list)]
        rep_trips = rep_trips.merge(self.routes, how='left', on='route_id')
        rep_trips = self.shapes.merge(rep_trips, how='right', on='shape_id')
        rep_trips.rename(columns={'trip_id': 'rep_trip_id'}, inplace=True)
        assert rep_trips.geometry.hasnans==False
        return rep_trips

    def get_line_stops_gdf(self):
        """
        Returns a GeoDataFrame with records for each stop for each
        rep_trip_id.
        """
        route_stops = self._df_all_stops_by_trips[
            self._df_all_stops_by_trips['trip_id'].isin(self.__rep_trip_list)]
        route_stops = route_stops.merge(self.stops, how='left', on='stop_id')
        route_stops = gpd.GeoDataFrame(
            route_stops, geometry=route_stops['geometry'])
        route_stops.drop(
            columns=['shape_id', 'arrival_time', 'departure_time',
                     'departure_time_mins', 'departure_time_hrs',
                     'block_id'], axis=1, inplace=True)
        route_stops = route_stops.set_crs(epsg=self._crs_epsg)
        return route_stops

    def get_line_time(self):
        """
        Returns a DataFrame with records for each rep_trip_id
        and their total service time.
        """
        first = self._df_all_stops_by_trips.groupby(
            ['trip_id'])['departure_time_mins'].first()
        first.rename('first', inplace=True)

        last = self._df_all_stops_by_trips.groupby(
            ['trip_id'])['departure_time_mins'].last()
        last.rename('last', inplace=True)

        route_id = self._df_all_stops_by_trips.groupby(
            ['trip_id'])['rep_trip_id', 'route_id'].first()

        df = pd.concat([route_id, first, last], axis=1).reset_index()
        df['total_line_time'] = df['last'] - df['first']
        return df

    def get_service_hours_by_line(self):
        """
        Returns a DataFrame with records for each rep_trip_id and columns with
        the number of service hours for each hour after midnight with service.
        For example 2:00-3:00 AM is called hour_2 and 3:00-4:00 PM is called
        hour_15.
        """
        df = self.get_line_time()
        df = df.groupby('rep_trip_id')['total_line_time'].sum().to_frame()
        df.reset_index(inplace=True)
        df = df.merge(self.trips[['route_id', 'trip_id', 'direction_id']],
                      how='left', left_on='rep_trip_id', right_on='trip_id')

        return df[['rep_trip_id', 'total_line_time', 'route_id',
                  'direction_id']]

    def get_routes_by_stops(self):
        """
        Returns a DataFrame with records for each rep_trip_id and a column
        holding a list of stops for each line.
        """
        df = self._df_all_stops_by_trips[self._df_all_stops_by_trips[
            'trip_id'] == self._df_all_stops_by_trips['rep_trip_id']]
        df = pd.DataFrame(df.groupby('stop_id')['route_id'].apply(
            lambda x: list(set(x.tolist()))))
        df.reset_index(inplace=True)
        return df

    def get_total_trips_by_line(self):
        """
        Returns a DataFrame with records for each rep_trip_id and a column
        holding the total number of trips for each line.
        """
        df = self.schedule_pattern_df.groupby('rep_trip_id').count(
                 ).reset_index()[['rep_trip_id', 'orig_trip_id']]
        df = df.rename(columns={'orig_trip_id': 'total_trips'})
        df = df.merge(self.trips, how='left', left_on='rep_trip_id',
                      right_on='trip_id')
        return df[['rep_trip_id', 'route_id', 'direction_id', 'total_trips']]
