passenger_assigner.py

import sys
import re
import pandas as pd
import numpy as np
from .airspace_particularities import get_nas_airport
import datetime
import multiprocessing as mp
import math as m
import pulp as p
import ast
import time
from concurrent.futures import ProcessPoolExecutor

from pyomo.environ import minimize, ConcreteModel, Var, Objective, Constraint, NonNegativeIntegers, maximize

from .lexicographic_lib import lexicographic_optimization
from .general_tools import fit


# import importlib
# importlib.reload(mysql)
# importlib.reload(ap)

#from connection_tools import read_data, write_data
import logging

# Set up basic configuration for logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


__version__ = '2.0'

############################
# PAX OPTIONS COMPUTATION  #
############################

passenger_options_version = '1.1'
# NOTE: THE PASSENGER OPTIONS COMPUTATION MIGHT NOT BE WORKING AS IT WAS SET TO WORK
#       WITH THE DATABASE OF VISTA/ DOMINO PROJECTS. THE STANDALONE FUNCTION TO COMPUTE
#       ALTERNATIVES COULD BE FINE.
#       TODO: Check/review this code to make it work standalone (if not already).

def compute_alliance_score(airline_list,alliance={'IBE': 'one','IBS': 'one','ANE': 'one'}):
    #score -1: dif airlines, 0:same airlines, 1:same alliance
    #print(airline_list)
    score = 0
    if len(airline_list)==0:
        return -1
    if airline_list.count(airline_list[0]) == len(airline_list):
        return 0
    for i,airline in enumerate(airline_list):
        #print(i,airline)
        if i==0:
            if airline not in alliance:
                return -1
            continue
        if airline == airline_list[i-1]:
            continue
        if airline not in alliance:
            return -1
        if alliance[airline]==alliance[airline_list[i-1]]:
            score = 1

    return score

def get_scenario_id(paras, connection=None):
    sql = """
		SELECT id FROM scenario WHERE scenario='{}' 
			""".format(paras['scenario'])
    df = read_data(fmt='mysql',
                   query=sql,
                   connection=connection)

    return df['id'].iloc[0]

def add_possible_options_itineraries_with_csv(d, scenario_id, PO_run, SM_run, connection=None, table_name='possible_options_itineraries'):
    df = d[['id','option_number','nid_f1','nid_f2','nid_f3','mct_leg2','mct_leg3','waiting_time_c1','waiting_time_c2','total_waiting_time',
            'total_time','PO_run','SM_run','scenario_id']].sort_values(by=['id','option_number','PO_run'])

    #write_data(df,
    #connection=connection,
    #table_name=table_name,
    #how='update',
    #keys_for_update={'scenario_id':scenario_id,
    #'PO_run':PO_run,
    #'SM_run':SM_run}
    #)
    df.to_csv(table_name)
#self.load_data_infile(, , delete_file=False)

# def insert_run_version(connection, run, model, model_version):
# 	# TODO: replace by write_data
# 	df = pd.DataFrame([{'run':run, 'model':model,'version':model_version}])
# 	write_data(connection=connection,
# 				)
# 	.to_sql('model_version',self.engine,index=False,if_exists="append")

def add_itineraries(it, scenario_id, PO_run, SM_run, connection=None, table_name='pax_itineraries'):
    file_name = "pax_it_" + str(datetime.datetime.now()) + "_" + str(np.round(np.random.random()*999999)) + ".csv"

    df = it[['it','option','leg1','leg2','leg3','distance_total','pax','avg_fare','nid',
             'ticket_type','generated_info','scenario_id','SM_run','PO_run','PG_run']].sort_values(by=['leg1','avg_fare','nid'])

    # Change labels of passengers for Mercury. UNTESTED
    df.loc[df['ticket_type']=='premium','ticket_type'] = 'flex'
    df.loc[df['ticket_type']=='standard','ticket_type'] = 'economy'
    #write_data(df,
    #connection=connection,
    #table_name=table_name,
    #how='update',
    #keys_for_update={'scenario_id':scenario_id,
    #'PO_run':PO_run,
    #'SM_run':SM_run}
    #)
    df.to_csv(table_name+'.csv')

def compute_passenger_options(paras, connection, verbose=False):
    print ('Computing options for scenario {}'.format(paras['scenario']))
    scenario_id = paras['scenario']

    pc = paras['nprocs']

    if verbose:
        print("-----------------------------------------")
        print("Computing possible options")
        print(datetime.datetime.now())


    dpf = pd.read_csv(paras['name_input_flow'], low_memory=False)
    dpf = dpf[(dpf['Origin Country Name']=='SPAIN') | (dpf['Destination Country Name']=='SPAIN')]
    airports = pd.read_csv(paras['name_airports'], low_memory=False)

    if len(dpf)==0:
        print("No passenger flows available for scenario "+str(scenario_id)+" in table "+paras['name_input_flow'])
        sys.exit(-1)

    ds = pd.read_parquet(paras['name_input_schedule'])

    if len(ds)==0:
        print("No flight schedules available for scenario "+str(scenario_id)+" in table "+paras['name_input_schedule'])
        sys.exit(-1)

    ds = ds[['origin', 'destination', 'airline', 'nid', 'sibt', 'sobt']]
    mct = pd.read_csv(paras['name_airport_static'])

    if len(mct)==0:
        print("No MCT available for scenario "+str(scenario_id)+" in table "+paras['name_airport_static'])
        sys.exit(-1)

    if verbose:
        print ("READ PASSENGER FLOWS "+str(len(dpf)))
        print ("READ Flight SCHEDULES "+str(len(ds)))
        print ("READ MCT "+str(len(mct)))


    dict_mct_dom = mct[['icao_id','MCT_domestic','MCT_international']].set_index('icao_id').to_dict()['MCT_domestic']
    dict_mct_int = mct[['icao_id','MCT_domestic','MCT_international']].set_index('icao_id').to_dict()['MCT_international']

    dpf['number_connections'] = dpf.apply(lambda row: {'NON-STOP':0,'ONE-STOP':1,'TWO-STOP':2,'THREE-STOP':3}[row['Itinerary']],axis=1)
    dpf = dpf.merge(airports[['IATA','ICAO']],left_on='Origin Airport', right_on="IATA", how='left')
    dpf = dpf.rename(columns={"ICAO": "origin"}).drop(['IATA'], axis=1)
    dpf = dpf.merge(airports[['IATA','ICAO']],left_on='Destination Airport', right_on="IATA", how='left')
    dpf = dpf.rename(columns={"ICAO": "destination"}).drop(['IATA'], axis=1)
    dpf = dpf.merge(airports[['IATA','ICAO']],left_on='Connecting Airport1', right_on="IATA", how='left')
    dpf = dpf.rename(columns={"ICAO": "connecting_airport1"}).drop(['IATA'], axis=1)
    dpf = dpf.merge(airports[['IATA','ICAO']],left_on='Connecting Airport2', right_on="IATA", how='left')
    dpf = dpf.rename(columns={"ICAO": "connecting_airport2"}).drop(['IATA'], axis=1)
    dpf = dpf.merge(airports[['IATA','ICAO']],left_on='Connecting Airport3', right_on="IATA", how='left')
    dpf = dpf.rename(columns={"ICAO": "connecting_airport3"}).drop(['IATA'], axis=1)

    dpf['airport_sequence_list'] = dpf.apply(lambda row: [row['origin'],row['connecting_airport1'],row['connecting_airport2'],row['connecting_airport3'],row['destination']],axis=1)
    dpf['airport_sequence_list'] = dpf.apply(lambda row: [x for x in row['airport_sequence_list'] if pd.isna(x)==False],axis=1)
    dpf['con_test'] = dpf.apply(lambda row: True if len(row['airport_sequence_list'])==row['number_connections']+2 else False,axis=1)
    dpf = dpf[dpf['con_test']==True] #drop unrecognised connections, very few in dataset
    dpf = dpf.dropna(subset=['origin','destination']) #dropping flows which are probably to rail station or ferry
    dpf['id'] = dpf.index
    dpf.to_csv('dfp.csv')
    print(dpf)
    #dpf['airline_sequence_list'] = dpf['airline_sequence'].str.split('_')
    dpf['scenario_id'] = scenario_id
    #dpf['economic_run'] = economic_run
    dpf['SM_run'] = paras['schedule_run']

    dpf_options_orig = dpf.copy()

    if paras['only_flights_with_trajectories']:
        #Read which flights have flight plan and then keep only those
        if verbose:
            print ("READING TRAJECTORIES")
        fwfp = mysql_vista.read_trajectories_options(scenario_id=scenario_id,fields="distinct schedule_id")
        ds_reduced = ds.merge(fwfp, how="inner", left_on="nid", right_on="schedule_id").drop(['schedule_id'],axis=1)
        if verbose:
            print ("ORIGINALLY HAVE THIS SCHEDULES "+str(len(ds)))
            print ("REDUCED TO WITH FLIGH PLAN     "+str(len(ds_reduced)))
    else:
        ds_reduced = ds#ds[['origin','destination','airline','nid','sibt','sobt']]


    #Compute options
    if verbose:
        print ("Compute parameters")
        print (datetime.datetime.now())

    #Use parallel computing
    if pc>1:
        n_compute = len(dpf)

        prev_i = 0
        n_per_section = max(1,round(len(dpf)/pc))
        if n_per_section==1:
            pc = len(dpf)
        i = n_per_section

        for nr in range(pc):
            if nr==pc-1:
                i=len(dpf)

            d = dpf.iloc[prev_i:i].copy().reset_index(drop=True)

            if nr == 0:
                n_params = [[d]+[ds_reduced]+[dict_mct_dom]+[dict_mct_int]] #[[d,ds_reduced,dict_mct_dom,dict_mct_int]]
            else:
                pass
                n_params.append([d]+[ds_reduced]+[dict_mct_dom]+[dict_mct_int])

            prev_i = i
            i = i + n_per_section

        pool = mp.Pool(processes=pc)

        if verbose:
            print("Launching parallel options computation")

        res = pool.starmap(compute_options,n_params)

        pool.close()
        pool.join()
    else:

        res = [compute_options(dpf, ds_reduced, dict_mct_dom, dict_mct_int)]

    # STORE options IN DATABASE
    #po_run = 0 #u.get_run_number(model='PO', engine=mysql_vista.engine)

    if verbose:
        print ("Store in DB results")
        print (paras['po_run'])
        print (datetime.datetime.now())

    #mysql_vista.insert_run_version(paras['po_run'],'PO',passenger_options_version)
    if verbose:
        k = 0
    for r in res:
        if verbose:
            print (k)
            k = k + 1
        r['PO_run'] = paras['po_run']
        for i in range(1, 4):
            if 'nid_f'+str(i) not in r.columns:
                r['nid_f'+str(i)] = None
                r['mct_leg'+str(i)] = None

        if verbose:
            print("CALL ADD POSSIBLE OPTIONS ITINERARIES")
        #add_possible_options_itineraries_with_csv(r)

        #print(r)


        add_possible_options_itineraries_with_csv(r,
                                                  scenario_id,
                                                  paras['po_run'],
                                                  paras['schedule_run'],
                                                  connection=connection,
                                                  table_name=paras['name_output_pax_option'])

    # finally:
    #     if ssh_parameters is not None:
    #         mysql_vista.close(close_server=True)
    #     else:
    #         mysql_vista.close()

    if verbose:
        print("Done")
        print(datetime.datetime.now())

def compute_options(dpf, ds, dict_mct_dom, dict_mct_int):
    dpf_options = dpf.copy()
    dpf_options_orig = dpf.copy()

    for i in range(1,5):
        connection = i-1
        dpf_leg = dpf_options_orig[['id','airport_sequence_list']][dpf_options_orig['number_connections']>=connection].copy()
        if len(dpf_leg)>0:
            dpf_leg['origin_leg'+str(i)] = dpf_leg['airport_sequence_list'].apply(lambda x: x[i-1])
            dpf_leg['origin_nas'+str(i)] = dpf_leg['airport_sequence_list'].apply(lambda x: get_nas_airport(x[i-1]))
            dpf_leg['destination_leg'+str(i)] = dpf_leg['airport_sequence_list'].apply(lambda x: x[i])
            dpf_leg['destination_nas'+str(i)] = dpf_leg['airport_sequence_list'].apply(lambda x: get_nas_airport(x[i]))
            dpf_leg['airline_leg'+str(i)] = 0

            if i>1:
                #print(dict_mct_int)
                dpf_leg['mct_leg'+str(i)] = dpf_leg['airport_sequence_list'].apply(lambda x: dict_mct_dom.get(x[i-1])
                if get_nas_airport(x[i-1])==get_nas_airport(x[i])
                else dict_mct_int.get(x[i-1]))
                dpf_leg['mct_leg'+str(i)] = dpf_leg.apply(lambda x: 120 if (x['mct_leg'+str(i)]=='') or (pd.isna(x['mct_leg'+str(i)])) else x['mct_leg'+str(i)], axis=1)
                #dpf_leg.to_csv('dpf_leg.csv')
                dpf_leg['mct_leg'+str(i)] = dpf_leg.apply(lambda x: datetime.timedelta(minutes=x['mct_leg'+str(i)]), axis=1)
            else:
                dpf_leg['mct_leg'+str(i)] = 0

            dpf_options = pd.merge(dpf_options, dpf_leg[['id','origin_leg'+str(i),'destination_leg'+str(i),'airline_leg'+str(i),
                                                         'origin_nas'+str(i),'destination_nas'+str(i),'mct_leg'+str(i)]],
                                   left_on=['id'], right_on=['id'], how='left')

            dpf_options = pd.merge(dpf_options, ds,
                                   left_on=['origin_leg'+str(i),'destination_leg'+str(i)],
                                   right_on=['origin','destination'], suffixes=('','_f'+str(i)), how='left')

            dpf_options = dpf_options.rename(columns={'sibt':'SIBT_time_f'+str(i), 'sobt':'SOBT_time_f'+str(i),
                                                      'airline':'airline_f'+str(i),'nid':'nid_f'+str(i)})

            if i>1:
                dpf_options = dpf_options[(dpf_options['SOBT_time_f'+str(i)]>(dpf_options['SIBT_time_f'+str(i-1)]+
                                                                              dpf_options['mct_leg'+str(i)])) |
                                          (dpf_options['number_connections']<(i-1))]
    #for connection in range(1,4):
    #df = dpf_options[dpf_options['number_connections']==connection]
    #df
    dpf_options['airline_list'] = dpf_options.apply(lambda x: [x['airline_f1'],x['airline_f2'],x['airline_f3'],x['airline_f4']], axis=1)
    dpf_options['airline_list'] = dpf_options.apply(lambda row: [x for x in row['airline_list'] if pd.isna(x)==False],axis=1)
    dpf_options['alliance_score'] = dpf_options.apply(lambda row: compute_alliance_score(row['airline_list']),axis=1)

    #only keep possible connections with same airline or alliance
    dpf_options = dpf_options[dpf_options['alliance_score']>-1]

    dpf_options.to_csv('dpf_options.csv')
    dpf_options['total_time'] = dpf_options.apply(lambda x: (x['SIBT_time_f'+str(x['number_connections']+1)]-x['SOBT_time_f1']).total_seconds()/60,
                                                  axis=1)

    if 'SOBT_time_f2' in dpf_options.columns:
        dpf_options['waiting_time_c1'] = dpf_options.apply(lambda x: (x['SOBT_time_f2']-x['SIBT_time_f1']).total_seconds()/60
        if not pd.isnull(x['SOBT_time_f2']-x['SIBT_time_f1'])
        else 0,axis=1)
    else:
        dpf_options['waiting_time_c1'] = 0

    if 'SOBT_time_f3' in dpf_options.columns:
        dpf_options['waiting_time_c2'] = dpf_options.apply(lambda x: (x['SOBT_time_f3']-x['SIBT_time_f2']).total_seconds()/60 if not pd.isnull(x['SOBT_time_f3']-x['SIBT_time_f2'])
        else 0,axis=1)
    else:
        dpf_options['waiting_time_c2'] = 0

    for i in range(2,4):
        if 'mct_leg'+str(i) in dpf_options.columns:
            dpf_options['mct_leg'+str(i)] = dpf_options['mct_leg'+str(i)].apply(lambda x: x.total_seconds()/60)

    dpf_options['total_waiting_time'] = dpf_options['waiting_time_c1'] + dpf_options['waiting_time_c2']
    dpf_options['option_number'] = dpf_options.groupby(['id']).cumcount()+1

    return dpf_options


########################################
# ITINERARY GENERATOR -- PAX ASSIGMENT #
########################################

itinerary_generator_version = '2.1'
# Version 2.0 which allows use of lexicographic optimisation too
# Version 2.1 improves lexicographic to consider trains with their stops

# NOTE: SOME OF THE OPTIONS ARE ALSO RELATED TO THE DOMINO/VISTA ARCHITECTURE (E.G. READ OPTIONS)
#       TODO: Review code to make it standalone (might break backcompatibility with Vista/Domino project

def assign_passengers(paras, connection, verbose=False, dry_run=True):
    logging.info('Assigning passenger for scenario {}'.format(paras['scenario']))

    scenario_id = paras['scenario']

    if verbose:
        logging.info("-----------------------------------------")
        logging.info("Scenario: "+str(paras['scenario']))
        logging.info("PO_run: "+str(paras['po_run']))
        logging.info("schedule_run: "+str(paras['schedule_run']))
        logging.info("")
    now = datetime.datetime.now()
    now_ini = now

    #READ DATA FROM DATABASE ON FLIGHT SCHEDULES TO KNOW SEATS AND ITINERARIES OPTIONS
    ds = pd.read_parquet(paras['name_input_schedule'])
    ds_trains = pd.read_parquet(paras['trains_schedule'])
    ds = pd.concat([ds,ds_trains])
    ds['nid'] = ds['nid'].astype('str')

    ds = ds[['nid', 'max_seats', 'gcdistance']]

    pax_options = pd.read_csv(paras['external_options'])
    ds = ds[(ds['nid'].isin(pax_options['nid_f1'])) | ds['nid'].isin(pax_options['nid_f2'])]
    #pax_options = pax_options[pax_options['type'].isin(['flight', 'flight_flight'])]

    #pax_options = pd.read_csv(paras['name_output_pax_option'])
    dfp = pd.read_csv(paras['name_input_flows'])
    dfp['Passengers'] = dfp['Passengers']/30.0 #monthly flow divided into daily
    dfp = dfp.rename(columns={'Passengers':'volume','Avg. Total Fare(USD)':'fare'})
    #options = read_passenger_options(connection=connection,
    #scenario_id=scenario_id,
    #po_run=paras['po_run'],
    #possible_itineraries_table=paras['name_output_pax_option'],
    #pax_flow_table=paras['name_input_flow']
    #)
    options = pd.merge(pax_options,
                       dfp[['id','volume','fare']],
                       on=['id'],
                       how='left',
                       suffixes=('','_x'))

    if len(options)==0:
        if paras['allow_pax_options_computation_pax_assignment']:
            # TODO: this part of the code doesn't work as read_passenger_options won't work
            compute_passenger_options(paras=paras,
                                      connection=connection,
                                      verbose=paras['verbose'])

            options = read_passenger_options(connection=connection,
                                             scenario_id=scenario_id,
                                             po_run=paras['po_run'],
                                             possible_itineraries_table=paras['name_output_pax_option'],
                                             pax_flow_table=paras['name_input_flow']
                                             )
            if len(options)==0:
                logging.error("No options after computing them for scenario"+str(scenario_id))
                sys.exit(-1)

        else:
            logging.error("No options available for scenario"+str(scenario_id)+"and not allow compute options")
            sys.exit(-1)

    if verbose:
        logging.info("ASSIGNING PAX WITH OPTIMISATION")
    it_gen, d_seats_max, options = assign_passengers_options_solver(ds, options, paras, verbose, now, pg_run = paras['pg_run'])

    if verbose:
        logging.info("GENERATING FILLERS AS NEEDED")
    it_gen = fill_flights_too_empy(it_gen, ds, d_seats_max, options, paras, verbose)

    logging.info("ADD FARE TYPE AND DISTANCE")
    print(it_gen.head())
    print(it_gen.generated_info.drop_duplicates())
    print(it_gen[it_gen.generated_info=='filler'])

    it_gen = compute_standard_premium(it_gen, ds, verbose)

    print(it_gen.head())
    print(it_gen.generated_info.drop_duplicates())
    print(it_gen[it_gen.generated_info=='filler'])

    it_gen['scenario_id'] = scenario_id
    it_gen['SM_run'] = paras['schedule_run']
    it_gen['PO_run'] = paras['po_run']

    # TODO saving results
    if not dry_run:
        if verbose:
            logging.info("INSERTING ITINERARIES"+str((datetime.datetime.now() - now)))
        it_gen['PG_run'] = paras['pg_run']

        add_itineraries(it_gen,
                        scenario_id,
                        paras['po_run'],  # PO_run,
                        paras['schedule_run'],  # -1,#SM_run,
                        connection=connection,
                        table_name=paras['name_output'])

    # if verbose:
    # 	print("INSERTING RUN VERSION")
    # 	print(datetime.datetime.now())
    # 	print(datetime.datetime.now()-now)
    # 	now = datetime.datetime.now()
    # mysql_vista.insert_run_version(pg_run,'PG',itinerary_generator_version)

    if verbose:
        logging.info("TOTAL TIME "+str(datetime.datetime.now() - now_ini))


def assign_passengers_options_solver(ds, options, paras, verbose=False, now=None, pg_run=0):
    if now is None:
        now = datetime.datetime.now()

    d_seats_max = ds[['nid','max_seats', 'mode']].copy()
    d_seats_max.rename(columns={'mode':'mode_transport'}, inplace=True)
    d_seats_max['nid'] = d_seats_max['nid'].astype('str')

    nid_columns_in_final = [col for col in options.columns if col.startswith('nid_f')]
    it = options[['cluster_id', 'option'] + nid_columns_in_final].rename(
        columns={'cluster_id': 'id'}).copy()

    # Compute volume of pax to be assigned per cluster
    options['volume_ceil'] = options['volume'].apply(m.ceil)
    dict_volume = options[['cluster_id','volume_ceil']].drop_duplicates().set_index('cluster_id').to_dict()['volume_ceil']

    # Randomise how much flights are full - target load factor
    def compute_target_load_factor(d_seats_max, distribution, mode_transport=None):
        if 'target_load_factor' not in d_seats_max.columns:
            d_seats_max['target_load_factor'] = None

        if distribution['type'] == 'fix':
            if mode_transport is not None:
                d_seats_max.loc[d_seats_max.mode_transport==mode_transport, 'target_load_factor'] = distribution['param']
            else:
                d_seats_max.loc[d_seats_max.target_load_factor.isna(), 'target_load_factor'] = distribution['param']

        elif distribution['type'] == 'triangular':
            # If you want a target load factor < 100
            if mode_transport is not None:
                mask = d_seats_max.mode_transport == mode_transport
            else:
                mask = d_seats_max.target_load_factor.isna()

            d_seats_max.loc[mask, 'target_load_factor'] = np.random.triangular(distribution['param'][0],
                                                                     distribution['param'][1],
                                                                     distribution['param'][2],
                                                                     len(d_seats_max.loc[mask, 'target_load_factor'] ))
        return d_seats_max

    if 'flight' in paras['target_load_factor'].keys():
        d_seats_max = compute_target_load_factor(d_seats_max, paras['target_load_factor']['flight'],
                                                 mode_transport='flight')
    if 'rail' in paras['target_load_factor'].keys():
        d_seats_max = compute_target_load_factor(d_seats_max, paras['target_load_factor']['rail'],
                                                 mode_transport='rail')

    # Put 1 in all the remaining
    d_seats_max = compute_target_load_factor(d_seats_max, {'type':'fix', 'param':1})

    d_seats_max['target_seats'] = d_seats_max['target_load_factor'] * d_seats_max['max_seats']
    d_seats_max['target_seats'] = d_seats_max['target_seats'].apply(m.ceil)

    dict_fc = d_seats_max.set_index('nid').to_dict()['target_seats']
    if paras.get('train_seats_per_segment') == 'combined':
        dict_type_service = None
    else:
        dict_type_service = d_seats_max.set_index('nid').to_dict()['mode_transport']

    # Keep only the capacities restrictions for the services which appear in the itineraries.
    unique_ids = set(it[nid_columns_in_final].stack())
    filtered_dict_fc = {k: v for k, v in dict_fc.items() if k in unique_ids}

    if paras['type_of_optimisation'] == 'max_assinged_only':
        # Might not be working in current version (not used in any case)
        it = optimise_assigned_pax_only_gurobi_pulp(it, filtered_dict_fc, dict_volume, paras[paras['type_of_optimisation']],
                                                    verbose=verbose, now=now, pg_run=pg_run)

    elif paras['type_of_optimisation'] == 'lexicographic':
        #print("it", it)
        #print("filtered_dict_fc", filtered_dict_fc)
        #print("dict_volume", dict_volume)
        #print("type_optimisation", paras[paras['type_of_optimisation']])
        #print("solver", paras[paras['type_of_optimisation']].get('solver'))

        it = optimise_lexicographic_pyomo(it, filtered_dict_fc, dict_type_service,
                                          dict_volume, paras[paras['type_of_optimisation']],
                                          solver=paras[paras['type_of_optimisation']].get('solver'))

    else:
        print("Type of optimisation not valid ", paras['type_of_optimisation'])
        sys.exit(0)

    if verbose:
        print ("COMPUTING FLIGHT OCCUPANCY")
        print (datetime.datetime.now())
        print (datetime.datetime.now()-now)
        now = datetime.datetime.now()


    it_gen = it[it['pax']>0].merge(options[['cluster_id','option','fare']],
                                   left_on=['id', 'option'],
                                   right_on=['cluster_id','option'],
                                   how='left').drop(['id'],
                                                    axis=1) \
        .reset_index(drop=True)


    # Create a dictionary to rename 'nid_fx' to 'legx'
    rename_dict = {col: 'leg' + str(re.search(r'_f(\d+)$', col).group(1))
                   for col in it_gen.columns if col.startswith('nid_f')}
    rename_dict['fare'] = 'avg_fare'  # Rename 'fare' to 'avg_fare'
    # Apply the renaming
    it_gen = it_gen.rename(rename_dict, axis=1)

    it_gen['generated_info'] = 'flow'

    return it_gen, d_seats_max, options


def optimise_assigned_pax_only_gurobi_pulp(it, dict_fc, dict_volume, paras, verbose=False, now=None, pg_run=None):
    if now is None:
        now = datetime.datetime.now()
    if pg_run is None:
        pg_run = 'None'

    if (paras['compute_leg2_plus_first'] == 1) and ('nid_f2' in it.columns):
        # Control for nid_f2 to make sure we have 2 leg pax in the itineraries
        if verbose:
            print("DOING LEG 2 AND + OPTIMISATION")
            print(datetime.datetime.now())
            print(datetime.datetime.now() - now)
            now = datetime.datetime.now()

        if paras['problem_file'] is not None:
            prob_file = paras['problem_file'] + "_" + str(pg_run) + "_2+.lp"
        else:
            prob_file = None

        prob_23leg = create_problem_gurobi_pulp(it[it['nid_f2'].notnull()], dict_fc, dict_volume, prob_file)
        dict_var_23leg = solve_problem_gurobi_pulp(prob_23leg, verbose=verbose)
    else:
        dict_var_23leg = {}

    if verbose:
        print("DOING ALL OPTIMISATION")
        print(datetime.datetime.now())
        print(datetime.datetime.now() - now)
        now = datetime.datetime.now()

    if paras['problem_file'] is not None:
        prob_file = paras['problem_file'] + "_" + str(pg_run) + "_all.lp"
    else:
        prob_file = None

    # Saved into  pickle to do some test on data format, etc.
    #import pickle
    #with open("problem_initialisation_variables.pkl", "wb") as file:
    #	pickle.dump({'it': it, 'dict_fc': dict_fc, 'dict_volume': dict_volume, 'prob_file': prob_file,
    #				 'dict_var_23leg': dict_var_23leg}, file)

    prob_all = create_problem_gurobi_pulp(it, dict_fc, dict_volume, prob_file, dict_var_23leg)
    dict_var_all = solve_problem_gurobi_pulp(prob_all, verbose=verbose)

    it['ind'] = it.index
    it['pax'] = it['ind'].apply(lambda x: dict_var_all.get(x, 0))
    #with open("problem_results.pkl", "wb") as file:
    #	pickle.dump({'dict_var_all': dict_var_all, 'it_f': it}, file)

    it.drop(['ind'], axis=1, inplace=True)

    return it


def optimise_lexicographic_pyomo(it, dict_fc,dict_mode_transport, dict_volume, paras, solver=None):
    if solver is None:
        solver = 'gurobi'

    objectives_create_model = [o[0] for o in paras['objectives']]

    model = create_model_passenger_assigner_pyomo(it, dict_fc, dict_volume, pc=paras['nprocs'],
                                                  objectives=objectives_create_model,
                                                  dict_mode_transport=dict_mode_transport)

    # Count variables and constraints
    num_vars = len([v for v in model.component_data_objects(Var)])
    num_constraints = len([c for c in model.component_data_objects(Constraint)])
    print(f"Number of variables: {num_vars}")
    print(f"Number of constraints: {num_constraints}")

    objectives = [(o[0], maximize) if o[1]=='maximize' else (o[0], minimize) for o in paras['objectives']]
    thresholds = paras['thresholds']

    results = lexicographic_optimization(model, objectives, solver=solver, thresholds=thresholds,
                                         num_threads=paras.get('nprocs',1))

    # Process results

    # Create a dictionary of (it, opt) to pax values from the model
    assigned_pax_dict = {(id, opt): model.x[id, opt].value for id, opt in model.x}

    # Convert the dictionary to a DataFrame for efficient merging
    assigned_pax_df = pd.DataFrame(
        list(assigned_pax_dict.items()), columns=["it_opt", "pax"]
    )

    # Split the tuple column into separate 'it' and 'opt' columns
    assigned_pax_df[["id", "option"]] = pd.DataFrame(assigned_pax_df["it_opt"].tolist(), index=assigned_pax_df.index)

    # Drop the tuple column as it's no longer needed
    assigned_pax_df = assigned_pax_df.drop(columns=["it_opt"])

    # Merge with the original it_data
    it = it.merge(assigned_pax_df, on=["id", "option"], how="left")
    it["pax"] = it["pax"].apply(lambda x: 0 if x == -0.0 else x)

    # Fill missing pax values with 0
    it["pax"] = it["pax"].fillna(0)

    print("Optimization Results: " + str(results))
    print("Total num pax assigned: " + str(it.pax.sum()))
    if 'nid_f2' in it.columns:
        print("Total num pax connecting: " + str(it[~it.nid_f2.isna()].pax.sum()))

    logging.info("Optimization Results: " + str(results))
    logging.info("Total num pax assigned: " + str(it.pax.sum()))
    if 'nid_f2' in it.columns:
        logging.info("Total num pax connecting: " + str(it[~it.nid_f2.isna()].pax.sum()))

    return it


def fill_flights_too_empy(it, ds, d_seats_max, options, paras, verbose=False, now=None):
    if now is None:
        now = datetime.datetime.now()

    # Compute current flight occupancy before putting fillers
    print("")
    print("")
    print("")
    print("")
    print(it.head())
    print(d_seats_max)

    rename_dict = {col: 'nid_f' + str(re.search(r'leg(\d+)$', col).group(1))
                   for col in it.columns if col.startswith('leg')}
    fo = compute_flight_occupancy(it.copy().rename(rename_dict, axis=1), d_seats_max)

    if verbose:
        print("*****************")
        print("PRE FILLERS")
        print_analysis_assigment(it.copy().rename(rename_dict, axis=1), options, fo)
        print("*****************")

    if paras['fillers_in_not_used_flights'] == 1:
        #Get flights from d_seats which don't have any passenger assigned, i.e., do not appear in fo
        #add them to the list of fo with 0 pax and 0 load factor used

        ds_wo_fo = d_seats_max.merge(fo[['f_id']],left_on='nid',right_on='f_id',how='left')
        ds_wo_fo=ds_wo_fo[ds_wo_fo['f_id'].isnull()].reset_index(drop=True)
        ds_wo_fo['f_id']=ds_wo_fo['nid']
        ds_wo_fo=ds_wo_fo.drop('nid',axis=1)
        ds_wo_fo['pax']=0
        ds_wo_fo['load_factor']=0
        ds_wo_fo['ava_seats']=ds_wo_fo['max_seats']

        if verbose:
            print("*****************")
            print(len(ds_wo_fo),"FLIGHTS/TRAINS WITHOUT PAX TO BE FILL WITH FILLERS")
            print("*****************")

        fo = pd.concat([fo,ds_wo_fo])

    #FILL FLIGHTS THAT ARE TOO EMPTY
    if verbose:
        print ("REVIEW FLIGHT OCCUPANCY AND CREATE EXTRA ITINERARIES")
        print (datetime.datetime.now())
        print (datetime.datetime.now()-now)
        now = datetime.datetime.now()

    #print("fo",len(fo))
    #print(fo.head())

    fo['load_factor_new_target'] = fo['load_factor']

    if paras['minimum_load_factor']['type']=='fix':
        fo['load_factor_new_target'] = fo['load_factor'].apply(lambda x:
                                                               max(x,paras['minimum_load_factor']['param']))

    elif paras['minimum_load_factor']['type']=='triangular':
        fo['load_factor_new_target'] = fo['load_factor'].apply(lambda x:
                                                               max(x,np.random.triangular(paras['minimum_load_factor']['param'][0],
                                                                                          paras['minimum_load_factor']['param'][1],
                                                                                          paras['minimum_load_factor']['param'][2])))
    else:
        sys.exit(-1)

    fo['extra_pax'] = fo['load_factor_new_target']*fo['max_seats']-fo['pax']
    fo['extra_pax'] = fo['extra_pax'].apply(round)

    it_new = fo.loc[fo['extra_pax']>0,['f_id','extra_pax']].copy().reset_index(drop=True).rename({'f_id':'leg1','extra_pax':'pax'},axis=1)
    it_new['it'] = None
    it_new['option'] = None

    it_leg_columns = [col for col in it.columns if col.startswith('leg')]
    for l in it_leg_columns:
        print("AAA", l)
        if l != 'leg1':
            it_new[l] = None

    it_new['generated_info'] = 'filler'
    it_new = it_new.merge(ds[['nid','gcdistance']],left_on='leg1', right_on='nid',suffixes=('','_s')).drop(['nid'],axis=1)

    #Linear approximation of cost of flight as a function of distance leg1
    it_w_dist = it[it['leg2'].isnull()].merge(ds[['nid','gcdistance']],left_on='leg1',right_on='nid',suffixes=('','_s')).drop(['nid'],axis=1)
    popt, pcov, cf = fit(it_w_dist['gcdistance'], it_w_dist['avg_fare'])

    it_new['avg_fare'] = it_new['gcdistance'].apply(cf)
    it_new['avg_fare'] = it_new['avg_fare'].apply(lambda x: max(np.random.normal(x,100),cf(300)))

    it = pd.concat([it, it_new])
    it = it.drop(['gcdistance'],axis=1)

    if verbose:
        print ("RECOMPUTE FLIGHT OCCUPANCY AFTER FILLERS")
        print (datetime.datetime.now())
        print (datetime.datetime.now()-now)
        now = datetime.datetime.now()

    fo_w_fillers = compute_flight_occupancy(it.copy().rename({'leg1':'nid_f1','leg2':'nid_f2','leg3':'nid_f3'},
                                                             axis=1),d_seats_max)

    if verbose:
        print_analysis_assigment(it.copy().rename({'leg1':'nid_f1','leg2':'nid_f2','leg3':'nid_f3'},
                                                  axis=1),options, fo_w_fillers)
    return it


def compute_standard_premium(it, ds, verbose=False, now=None):
    if now is None:
        now = datetime.datetime.now()

    if verbose:
        print ("COMPUTING PAX STANDARD / PREMIUM")
        print (datetime.datetime.now())
        print (datetime.datetime.now()-now)
        now = datetime.datetime.now()

    it['ticket_type'] = 'standard'
    it['avg_fare'] = it['avg_fare'].apply(lambda x: round(x,2))

    it.loc[it['leg2'].isnull(),'leg2'] = -1
    it.loc[it['leg3'].isnull(),'leg3'] = -1
    ds['nid'] = ds['nid'].astype('str')
    it['leg1'] = it['leg1'].astype('str')
    it['leg2'] = it['leg2'].astype('str')
    it['leg3'] = it['leg3'].astype('str')
    print(it)
    it = it.merge(ds[['nid','gcdistance']], left_on='leg1', right_on='nid', suffixes=('','_s')).rename({'gcdistance':'dist_leg1'}, axis=1) \
        .merge(ds[['nid','gcdistance']], how="left", left_on='leg2', right_on='nid', suffixes=('','_s')).rename({'gcdistance':'dist_leg2'}, axis=1) \
        .merge(ds[['nid','gcdistance']], how="left", left_on='leg3', right_on='nid', suffixes=('','_s')).rename({'gcdistance':'dist_leg3'}, axis=1)
    print(it)
    it.loc[it['leg2']==-1,'leg2'] = None
    it.loc[it['leg3']==-1,'leg3'] = None

    it['dist_leg2'] = it['dist_leg2'].fillna(0)
    it['dist_leg3'] = it['dist_leg3'].fillna(0)
    it['distance_total'] = it['dist_leg1']+it['dist_leg2']+it['dist_leg3']

    popt, pcov, cf = fit(it['distance_total'],it['avg_fare'])
    it['fare_estimated'] = it['distance_total'].apply(lambda x: cf(x))
    it['extra_price'] = it['avg_fare']-it['fare_estimated']
    it['extra_price_rel'] = (it['avg_fare']-it['fare_estimated'])/it['fare_estimated']

    it_premium = it[it['extra_price_rel']>-0.25].copy().reset_index(drop=True)
    min_extra = min(it_premium['extra_price_rel'])
    max_extra = m.ceil(max(it_premium['extra_price_rel']))
    it_premium['pax_premium'] = it_premium.apply(lambda x:
                                                 min(round(30*np.random.rand()+25),round(x['pax']*max(0,(np.random.triangular(min_extra-4,x['extra_price_rel'],max_extra)-min_extra)/(max_extra-min_extra))))
                                                 ,axis=1)
    it_premium['ticket_type'] = 'premium'

    it = it.merge(it_premium[['it','avg_fare','leg1','leg2','leg3','pax','pax_premium']], on=['it','avg_fare','leg1','leg2','leg3','pax'],
                  how='left')
    it['pax_premium'] = it['pax_premium'].fillna(0)
    it['pax'] = it['pax']-it['pax_premium']
    it = it.drop(['pax_premium','fare_estimated','extra_price','extra_price_rel','dist_leg1','dist_leg2','dist_leg3'],axis=1)
    it = it[it['pax']>0].reset_index(drop=True)
    it_premium = it_premium.drop(['dist_leg1','dist_leg2','dist_leg3','fare_estimated','extra_price_rel','extra_price','pax'],axis=1).rename({'pax_premium':'pax'},axis=1)
    it_premium = it_premium[it_premium['pax']>0].reset_index(drop=True)
    it = pd.concat([it[['it','option','leg1','leg2','leg3','distance_total','pax','avg_fare','ticket_type','generated_info']],it_premium[['it','option','leg1','leg2','leg3','distance_total','pax','avg_fare','ticket_type','generated_info']]])

    it.reset_index(drop=True, inplace=True)
    it['nid'] = it.index
    it['avg_fare'] = it['avg_fare'].apply(lambda x: round(x,2))

    return it


def create_problem_gurobi_pulp(it, dict_fc, dict_volume, f_prob, dict_var_min={}):
    git = it.groupby(['it'])
    i_num_var = {}

    # Create the 'prob' variable to contain the problem data
    prob = p.LpProblem("Pax_assignment", p.LpMaximize)


    for g in git:
        dg = g[1]
        for v in dg.index:
            i_num_var[v]=p.LpVariable('opt'+'_'+str(v),lowBound=dict_var_min.get(v,0),upBound=dict_volume.get(dg.loc[v,'it']),cat=p.LpInteger)

        prob += p.lpSum([i_num_var[i] for i in dg.index]) <= dict_volume.get(dg.loc[v,'it']), "pax_in_it_"+str(dg.loc[v,'it'])


    if 'nid_f2' not in it.columns:
        it = it.copy()
        it['nid_f2'] = None
    if 'nid_f3' not in it.columns:
        it = it.copy()
        it['nid_f3'] = None

    for f in dict_fc:
        i_f=it.loc[(it['nid_f1']==f) | (it['nid_f2']==f) | (it['nid_f3']==f)].index
        prob += p.lpSum([i_num_var[i] for i in i_f]) <= dict_fc.get(f), "capacity_flight_in_"+str(f)
        #print(i_f)
    # The objective function is added to 'prob' first
    prob += p.lpSum([i_num_var[i] for i in i_num_var]), "Total_number_pax_assigned"


    if f_prob is not None:
        # The problem data is written to an .lp file
        prob.writeLP(f_prob)

    return prob


def solve_problem_gurobi_pulp(prob, verbose=False):
    # The problem is solved using PuLP and GUROBI

    #p.GUROBI(msg=verbose).solve(prob) #TODO add  OutputFlag
    p.GUROBI_CMD(msg=verbose).solve(prob) #TODO add  OutputFlag

    if verbose:
        # The status of the solution is printed to the screen
        print("Status:", p.LpStatus[prob.status])

        # The optimised objective function value is printed to the screen
        print("Total number pax assigned = ", p.value(prob.objective))


    # Each of the variables is printed with it's resolved optimum value
    dict_var={}
    for v in prob.variables():
        if v.name!='__dummy':
            dict_var[int(v.name.replace('opt_',''))]=v.varValue

    return dict_var


def compute_flight_occupancy(it, d_seats_max):
    flight_occupancy = pd.concat([it.groupby(['nid_f1'], as_index=False)['pax'].sum().rename(columns={'nid_f1':'f_id'}),
                                  it.groupby(['nid_f2'], as_index=False)['pax'].sum().rename(columns={'nid_f2':'f_id'}),
                                  it.groupby(['nid_f3'], as_index=False)['pax'].sum().rename(columns={'nid_f3':'f_id'})])


    flight_occupancy = flight_occupancy.groupby(['f_id'], as_index=False)['pax'].sum()

    flight_occupancy = flight_occupancy.merge(d_seats_max, left_on=['f_id'], right_on=['nid'], how='left')
    #print('flight_occupancy', flight_occupancy)
    #print(d_seats_max)
    x = flight_occupancy[flight_occupancy['f_id']=='0']
    y = d_seats_max[d_seats_max['nid']=='0']
    #print(x,y)
    flight_occupancy['load_factor'] = flight_occupancy['pax']/flight_occupancy['max_seats']
    flight_occupancy['ava_seats'] = flight_occupancy['max_seats']-flight_occupancy['pax']
    flight_occupancy.drop('nid', axis=1, inplace=True)

    return flight_occupancy


def print_analysis_assigment(it, options, fo):
    print("PAX ASSIGMENT RESULTS")
    print("--------")
    print("VOLUMES")
    print("Volume pax: ", options[['id', 'volume_ceil']].drop_duplicates()['volume_ceil'].sum())
    print("Assigned: ", it['pax'].sum())
    print("Perc: ", round(it['pax'].sum() / options[['id', 'volume_ceil']].drop_duplicates()['volume_ceil'].sum(), 2))
    print("--------")
    print("ITINERARIES")

    # Total itineraries
    all_with_pax = len(it.loc[it['pax'] > 0, ['it']].drop_duplicates())
    all_total = len(it[['it']].drop_duplicates())
    print("total itineraries: ", all_total)
    print("it with pax: ", all_with_pax)
    print("per with pax: ", round(100 * all_with_pax / all_total, 2))
    print("---")

    # Determine the maximum number of legs by checking 'nid_f' columns
    max_legs = max([int(col.split('_')[1][1:]) for col in it.columns if col.startswith('nid_f')])

    # Iterate over each leg dynamically
    for leg in range(1, max_legs + 1):
        leg_col = f'nid_f{leg}'
        next_leg_col = f'nid_f{leg + 1}' if leg + 1 <= max_legs else None

        if next_leg_col:
            leg_with_pax = len(it.loc[(it[leg_col].notnull()) & (it[next_leg_col].isnull()) & (it['pax'] > 0), [
                'it']].drop_duplicates())
            leg_total = len(it.loc[(it[leg_col].notnull()) & (it[next_leg_col].isnull()), ['it']].drop_duplicates())
        else:
            # Handle the last leg (it doesn't have a next leg)
            leg_with_pax = len(it.loc[(it[leg_col].notnull()) & (it['pax'] > 0), ['it']].drop_duplicates())
            leg_total = len(it.loc[it[leg_col].notnull(), ['it']].drop_duplicates())

        print(f"it {leg} leg: ", leg_total)
        print(f"it {leg} leg with pax: ", leg_with_pax)
        if leg_total > 0:
            print(f"per {leg} leg :", round(100 * leg_with_pax / leg_total, 2))
        print("---")

    print("-----")
    print("LOAD FACTOR")
    print(fo['load_factor'].describe())
    print("")
    print("")


def compute_service_assigment_chunk(chunk):
    flight_assignments_local = {}
    nid_columns = [col for col in chunk.columns if col.startswith("nid_f")]
    for _, row in chunk.iterrows():
        for f in nid_columns: #["nid_f1", "nid_f2", "nid_f3"]:
            flight = row[f]
            if pd.notna(flight):  # Ensure flight is not NaN
                if flight not in flight_assignments_local:
                    flight_assignments_local[flight] = []
                flight_assignments_local[flight].append((row["id"], row["option"]))
    return flight_assignments_local


def create_model_passenger_assigner_pyomo(it_data, dict_fc, dict_volume, pc=23, objectives=None, dict_mode_transport=None):
    if dict_mode_transport is None:
        dict_mode_transport = {}

    # Pyomo model
    model = ConcreteModel()

    ########
    # Sets #
    ########
    start_time_1 = time.time()

    # Get the unique itineraries
    itineraries = it_data["id"].unique()

    # Get all columns that start with "nid_f" (e.g., nid_f1, nid_f2, nid_f3, ...)
    nid_columns = [col for col in it_data.columns if col.startswith("nid_f")]

    # Create a set of all unique services by iterating over the nid_columns
    services = set()
    for col in nid_columns:
        services.update(it_data[col].dropna())

    services = {s for s in services}  # Ensure services are strings

    # Print the time taken for this part
    end_time_1 = time.time()
    logging.info(f"Time taken for creating itineraries and services: {end_time_1 - start_time_1:.4f} seconds")


    # Precompute flight assignments
    start_time_1 = time.time()

    # parallel_flight_assignments to gain speed
    def parallel_service_assignments(it_data, num_workers=24):
        # Combine results from multiple chunks
        def combine_results(results):
            combined = {}
            for result in results:
                for service, assignments in result.items():
                    if service not in combined:
                        combined[service] = []
                    combined[service].extend(assignments)
            return combined

        # Split data into chunks
        chunks = np.array_split(it_data, num_workers)

        # Use ProcessPoolExecutor for parallel processing
        with ProcessPoolExecutor(max_workers=num_workers) as executor:
            results = list(executor.map(compute_service_assigment_chunk, chunks))

        # Combine results
        service_assignments = combine_results(results)
        return service_assignments

    service_assignments = parallel_service_assignments(it_data, num_workers=pc)

    end_time_1 = time.time()
    logging.info(f"Time taken for Precompute service assigments: {end_time_1 - start_time_1:.4f} seconds")


    #############
    # Variables #
    #############


    start_time_1 = time.time()
    model.x = Var(((id, opt) for id, opt in it_data[["id", "option"]].values), domain=NonNegativeIntegers)
    end_time_1 = time.time()

    # Print the time taken for this part
    logging.info(f"Time taken for creating variables: {end_time_1 - start_time_1:.4f} seconds")


    ##############
    # Objectives #
    ##############
    # **Objective 1: Maximize total passengers on connecting itineraries**
    if (((objectives is None) or ('total_pax_connecting' in objectives)) and ('nid_f2' in it_data.columns)):
        # Check that there're at least itineraries with one connection, i.e., nid_f2 exist
        start_time_1 = time.time()

        # Precompute connecting options directly from the DataFrame
        connecting_options = [
            (row["id"], row["option"])
            for _, row in it_data.iterrows()
            if pd.notna(row["nid_f2"])  # Check if 'nid_f2' is not NaN
        ]

        # Define the objective using the precomputed list
        model.total_pax_connecting = Objective(
            expr=sum(model.x[it, opt] for it, opt in connecting_options),
            sense=maximize
        )

        end_time_1 = time.time()
        logging.info(f"Time taken for creating objective total_pax_connecting: {end_time_1 - start_time_1:.4f} seconds")


    # **Objective 2: Maximize total passengers overall**
    if (objectives is None) or ('total_pax' in objectives):
        start_time_1 = time.time()
        model.total_pax = Objective(
            expr=sum(model.x[id, opt] for id, opt in model.x),
            sense=maximize
        )
        end_time_1 = time.time()
        logging.info(f"Time taken for creating objective total_pax: {end_time_1 - start_time_1:.4f} seconds")


    # **Objective 3: Minimize the difference in load factors across flights**
    if (objectives is None) or ('minimize_load_factor_diff' in objectives):
        start_time_1 = time.time()
        # If the objective already exists, remove it first to avoid the warning
        #if hasattr(model, "minimize_load_factor_diff"):
        #    model.del_component(model.minimize_load_factor_diff)


        def load_factor_difference(model):
            service_load_factors = {}
            for service, assignments in service_assignments.items():
                total_pax_for_service = sum(
                    model.x[it, opt] for it, opt in assignments
                )
                service_load_factors[service] = total_pax_for_service / dict_fc[str(service)]

            load_factor_diff = sum(
                (service_load_factors[f] - sum(service_load_factors.values()) / len(service_load_factors)) ** 2
                for f in service_load_factors
            )
            return load_factor_diff

        model.minimize_load_factor_diff = Objective(
            expr=load_factor_difference(model),
            sense=minimize
        )

        end_time_1 = time.time()
        logging.info(f"Time taken for creating objective minimize_load_factor_diff: {end_time_1 - start_time_1:.4f} seconds")


    # **Objective 4: Penalise underutilisation of services**
    if (objectives is None) or ('minimize_underutilization' in objectives):
        start_time_1 = time.time()
        # Define a penalty for underutilization (adjust as necessary)
        underutilization_penalty = 3

        model.minimize_underutilization = Objective(
            expr=sum(
                underutilization_penalty * (
                        dict_fc[service] - sum(model.x[id, opt] for id, opt in service_assignments[service]))
                for service in services
            ),
            sense=minimize
        )
        end_time_1 = time.time()
        logging.info(f"Time taken for creating objective minimize_underutilization: {end_time_1 - start_time_1:.4f} seconds")


    # **Objective 5: Maximise combination of passengers assinged (Objective 2) and
    #                minimise underutilisation of services (Objective 4)**
    if (objectives is None) or ('total_pax_combined_utilisation' in objectives):

        start_time_1 = time.time()
        # Define the penalty weight for load factor imbalance (adjust as necessary)
        penalty_weight = 0.3

        # Combined Objective: Maximize passengers while minimizing load factor imbalance
        model.total_pax_combined_utilisation = Objective(
            expr=(
                    sum(model.x[it, opt] for it, opt in model.x) -
                    penalty_weight * sum(
                (sum(model.x[id, opt] for id, opt in service_assignments[s]) / dict_fc[s]) ** 2
                for s in services
            )
            ),
            sense=maximize
        )

        end_time_1 = time.time()
        logging.info(f"Time taken for creating objective total_pax_combined_utilisation: {end_time_1 - start_time_1:.4f} seconds")

    ###############
    # Constraints #
    ###############

    # 1. Itinerary passenger volume constraints
    start_time_1 = time.time()
    # Precompute the mapping of itineraries to their options
    itinerary_options = it_data.groupby('id')['option'].apply(list).to_dict()

    # Create the constraints using the precomputed mapping
    def itinerary_volume_rule(model, it):
        return sum(model.x[it, opt] for opt in itinerary_options[it]) <= dict_volume[it]

    model.itinerary_volume_constraints = Constraint(itineraries, rule=itinerary_volume_rule)
    end_time_1 = time.time()
    logging.info(f"Time taken for creating constraint 1: {end_time_1 - start_time_1:.4f} seconds")


    # 2. Services capacity constraints
    start_time_1 = time.time()

    def service_capacity_rule_prev(model, service):
        service_demand = sum(
            model.x[it, opt]
            for it, opt in service_assignments[service]
        )
        return service_demand <= dict_fc[str(service)]


    def is_overlapping(seg, origin, destination):
        _, seg_origin, seg_destination = seg.split("_")
        seg_origin, seg_destination = int(seg_origin), int(seg_destination)
        return not (seg_destination <= origin or seg_origin >= destination)


    def flight_service_capacity_rule(model, flight_service):
        flight_demand = sum(
            model.x[it, opt]
            for it, opt in service_assignments[flight_service]
        )
        # print("F", service, flight_demand, "\n")
        return flight_demand <= dict_fc[str(flight_service)]


    def rail_capacity_rule(model, rail_constraint):
        # We have in rail_constraint the capacity and the it and option
        # The rail_constraint is encoded in a slightly weird way:
        # 295|['ES511_ES300_0+10']
        # 295|['ES512_ES511_0+14', 'ES512_ES300_0+0']
        # This is because rail_capacity_rule will be called by
        # Constraint from pyomo environment and this only allows
        # to pass one variable, not a tuple/list or anything like that

        capacity = int(rail_constraint.split("|")[0])
        it_opts = ast.literal_eval(rail_constraint.split("|")[1])
        # Total demand across overlapping segments
        it_opts = [(id_opt.split("+")[0], int(id_opt.split("+")[1])) for id_opt in it_opts]
        train_demand = sum(
            model.x[id, opt] for id, opt in it_opts
        )

        return train_demand <= capacity


    # Divide the services between flight and rail services
    flight_services = [s for s in services if dict_mode_transport.get(s,'flight')=='flight']
    rail_services = [s for s in services if dict_mode_transport.get(s,'flight')=='rail']

    if len(flight_services) > 0:
        # Add constraints flights
        model.flight_capacity_constraints = Constraint(flight_services, rule=flight_service_capacity_rule)

    if len(rail_services) >0:
        # Add constraints rail
        # Dictionary of rail services (only the rail not between which stops, i.e., 123_5_6 becomes 123
        # Then for each service which stops are served
        dict_rail_services = {}
        # Capacity per service, i.e. capacity for 123, not for 123_5_6
        dict_capacity_rail_services = {}

        for rs in rail_services:
            # For each rail_service (e.g. 123_5_6) get list of stops and add them to the set of stops by that service
            rs_s = rs.split("_")
            rs_s[1] = int(rs_s[1])
            rs_s[2] = int(rs_s[2])
            rail_service_id = rs_s[0]  # Extract train ID
            service_stops = dict_rail_services.setdefault(rail_service_id, set())  # Use a set to store unique stops
            service_stops.update(rs_s[1:])  # add new stops
            dict_rail_services[rail_service_id] = service_stops

            # Keep most restrictive capacity of all the services for a given train
            dict_capacity_rail_services[rail_service_id] = min(dict_fc[rs],
                                                               dict_capacity_rail_services.get(rail_service_id, 999999999))


        # Convert sets back to sorted lists to have list of stops ordered per service id
        dict_rail_services = {k: sorted(list(v)) for k, v in dict_rail_services.items()}


        # Create rail constraints
        def build_constraints_rail_services(dict_rail_services, dict_capacity_rail_services):
            # This function builds the constraint for ail services as 295|['ES511_ES300_0+10'] or
            # 295|['ES512_ES511_0+14', 'ES512_ES300_0+0']
            # It iterates throw the dictionary of services (rs) and the stops they use
            # Then it goes between each succesive stops and check which rails use that
            # Then it checks whith id clusters and options use those service and adds them
            # to the list of services that if pax are assigned to them,then they should be
            # considered for the capacity of that segment/train.

            constraints_rail = []
            for (rs, stops) in dict_rail_services.items():
                for i in range(len(stops)-1):
                    # Find all segments of the same train
                    relevant_segments = [
                        seg for seg in dict_fc if (seg.startswith(rs + "_") and
                                                   dict_mode_transport.get(seg, 'flight') == 'rail')
                    ]

                    # Find all overlapping segments
                    overlapping_segments = [
                        seg for seg in relevant_segments
                        if is_overlapping(seg, stops[i], stops[i+1])
                    ]

                    id_opt_overlapping = [service_assignments[seg] for seg in overlapping_segments]
                    id_opt_overlapping_no_duplicates = list(
                        set([item for sublist in id_opt_overlapping for item in sublist]))

                    if len(id_opt_overlapping_no_duplicates) > 0:
                        #constraints_rail += [(rs+'_'+str(stops[i])+'_'+str(stops[i+1]), id_opt_overlapping_no_duplicates)]
                        # Keep in constraint only id of rail not between stops
                        constraints_rail += [(rs, id_opt_overlapping_no_duplicates)]


            # Remove duplicate constraints, as in some case it could be that the same o-d pair goes several
            # times over the same segment (if contained)
            constraints_rail = list({(k, tuple(v)) for k, v in constraints_rail})

            # Reformat into str to be able to pass it to Constraint function later
            # replacing the flight by their capacity
            constraints_rail = [str(dict_capacity_rail_services[s[0]])+"|"+str([str(t[0])+"+"+str(t[1]) for t in s[1]])
                                for s in constraints_rail]

            return constraints_rail

        constraints_rail = build_constraints_rail_services(dict_rail_services, dict_capacity_rail_services)

        model.rail_capacity_constraints = Constraint(constraints_rail, rule=rail_capacity_rule)

    end_time_1 = time.time()

    logging.info(f"Time taken for creating constraint 2: {end_time_1 - start_time_1:.4f} seconds")

    return model


def create_model_passenger_reassigner_pyomo(it_data, dict_volume, dict_sc, dict_mode_transport,
                                            objectives=None, pc=23):

    if dict_mode_transport is None:
        dict_mode_transport = {}

    # Pyomo model
    model = ConcreteModel()

    ########
    # Sets #
    ########
    start_time_1 = time.time()

    # Get the unique itineraries
    it_data.rename(columns={'pax_group_id': 'id'}, inplace=True)
    itineraries = it_data["id"].unique()

    # Get all columns that start with "service_id_" (e.g., service_id_0, service_id_1, service_id_2, ...)
    # Not get service_id_pax_ as those are the 'original' ones
    nid_columns = [col for col in it_data.columns if
                    col.startswith("service_id_") and not col.startswith("service_id_pax_")]

    # change servive_id_1, servive_id_n to nid_f1, nid_fn so that I can reuse
    # previous code. TODO: to make more generic instead of nid_fn use service_id_n even in previours case.
    dict_rename_service_id = {x: 'nid_f' + x.split("_")[-1] for x in nid_columns}
    it_data.rename(columns=dict_rename_service_id, inplace=True)

    # Create a set of all unique services by iterating over the nid_columns
    services = set()
    for col in dict_rename_service_id.values():
        services.update(it_data[col].dropna())

    services = {s for s in services}  # Ensure services are strings

    # Print the time taken for this part
    end_time_1 = time.time()
    logging.info(f"Time taken for creating itineraries and services: {end_time_1 - start_time_1:.4f} seconds")

    # Precompute service assignments
    start_time_1 = time.time()


    # parallel_flight_assignments to gain speed
    def parallel_service_assignments(it_data, num_workers=24):
        # Combine results from multiple chunks
        def combine_results(results):
            combined = {}
            for result in results:
                for service, assignments in result.items():
                    if service not in combined:
                        combined[service] = []
                    combined[service].extend(assignments)
            return combined

        # Split data into chunks
        chunks = np.array_split(it_data, num_workers)

        # Use ProcessPoolExecutor for parallel processing
        with ProcessPoolExecutor(max_workers=num_workers) as executor:
            results = list(executor.map(compute_service_assigment_chunk, chunks))

        # Combine results
        service_assignments = combine_results(results)
        return service_assignments

    service_assignments = parallel_service_assignments(it_data, num_workers=pc)

    end_time_1 = time.time()
    logging.info(f"Time taken for Precompute service assigments: {end_time_1 - start_time_1:.4f} seconds")


    #############
    # Variables #
    #############

    start_time_1 = time.time()
    model.x = Var(((id, opt) for id, opt in it_data[["id", "option"]].values), domain=NonNegativeIntegers)
    end_time_1 = time.time()

    # Print the time taken for this part
    logging.info(f"Time taken for creating variables: {end_time_1 - start_time_1:.4f} seconds")


    ##############
    # Objectives #
    ##############
    # **Objective 1: Assign maximum number of people**
    if (objectives is None) or ('total_pax' in objectives):
        start_time_1 = time.time()
        model.total_pax = Objective(
            expr=sum(model.x[id, opt] for id, opt in model.x),
            sense=maximize
        )
        end_time_1 = time.time()
        logging.info(f"Time taken for creating objective total_pax: {end_time_1 - start_time_1:.4f} seconds")

    it_data['opt_key'] = it_data.apply(lambda x: (x['id'], x['option']), axis=1)

    # **Objective 2: Assign passengers with the lowest delay possible w.r.t. arrival at home
    if 'arrival_door_difference' in objectives:
        start_time_1 = time.time()
        dict_delay_pos_arrival_home = it_data[it_data.delay_arrival_home>0].set_index('opt_key')['delay_arrival_home'].to_dict()

        if len(dict_delay_pos_arrival_home) > 0:
             # We have some options that will make the pax arrive later than planned home
            model.arrival_door_difference = Objective(
                expr=sum(model.x[id, opt] * dict_delay_pos_arrival_home.get((id, opt), 0)  for id, opt in model.x),
                    sense=minimize
                )

        end_time_1 = time.time()
        logging.info(f"Time taken for creating objective arrival_door_difference: {end_time_1 - start_time_1:.4f} seconds")

    # **Objective 3: Keep the same starting node (airport/rail station)
    if 'same_starting_node' in objectives:
        start_time_1 = time.time()

        it_data['same_initial_node_binary'] = 0
        it_data.loc[it_data.same_initial_node, 'same_initial_node_binary'] = 1
        dict_same_starting_node = it_data.set_index('opt_key')['same_initial_node_binary'].to_dict()

        model.same_starting_node = Objective(
                expr=sum(model.x[id, opt] * dict_same_starting_node.get((id, opt), 0) for id, opt in model.x),
                sense=maximize
        )

        end_time_1 = time.time()
        logging.info(
            f"Time taken for creating objective same_starting_node: {end_time_1 - start_time_1:.4f} seconds")

    # **Objective 4: Keep the same end node (airport/rail station)
    if 'same_arrival_node' in objectives:
        start_time_1 = time.time()

        it_data['same_final_node_binary'] = 0
        it_data.loc[it_data.same_final_node, 'same_final_node_binary'] = 1
        dict_delay_same_final_node = it_data.set_index('opt_key')['same_final_node_binary'].to_dict()

        model.same_arrival_node = Objective(
                expr=sum(model.x[id, opt] * dict_delay_same_final_node.get((id, opt), 0) for id, opt in model.x),
                sense=maximize
        )

        end_time_1 = time.time()
        logging.info(
            f"Time taken for creating objective same_arrival_node: {end_time_1 - start_time_1:.4f} seconds")


    # **Objective 4: Keep the same path
    if 'same_path' in objectives:
        start_time_1 = time.time()

        it_data['same_path_node_binary'] = 0
        it_data.loc[it_data.same_path, 'same_path_node_binary'] = 1
        dict_same_path = it_data.set_index('opt_key')['same_path_node_binary'].to_dict()

        model.same_path = Objective(
                expr=sum(model.x[id, opt] * dict_same_path.get((id, opt), 0) for id, opt in model.x),
                sense=maximize
        )

        end_time_1 = time.time()
        logging.info(
            f"Time taken for creating objective same_path: {end_time_1 - start_time_1:.4f} seconds")


    ###############
    # Constraints #
    ###############

    # 1. Itinerary passenger volume constraints
    start_time_1 = time.time()
    # Precompute the mapping of itineraries to their options
    itinerary_options = it_data.groupby('id')['option'].apply(list).to_dict()

    # Create the constraints using the precomputed mapping
    def itinerary_volume_rule(model, it):
        return sum(model.x[it, opt] for opt in itinerary_options[it]) <= dict_volume[it]

    model.itinerary_volume_constraints = Constraint(itineraries, rule=itinerary_volume_rule)
    end_time_1 = time.time()
    logging.info(f"Time taken for creating constraint 1: {end_time_1 - start_time_1:.4f} seconds")


    # 2. Services capacity constraints
    start_time_1 = time.time()

    def is_overlapping(seg, origin, destination):
        _, seg_origin, seg_destination = seg.split("_")
        seg_origin, seg_destination = int(seg_origin), int(seg_destination)
        return not (seg_destination <= origin or seg_origin >= destination)


    def flight_service_capacity_rule(model, flight_service):
        flight_demand = sum(
            model.x[it, opt]
            for it, opt in service_assignments[flight_service]
        )
        # print("F", service, flight_demand, "\n")
        return flight_demand <= dict_sc[str(flight_service)]


    def rail_capacity_rule(model, rail_constraint):
        # We have in rail_constraint the capacity and the it and option
        # The rail_constraint is encoded in a slightly weird way:
        # 295|['ES511_ES300_0+10']
        # 295|['ES512_ES511_0+14', 'ES512_ES300_0+0']
        # This is because rail_capacity_rule will be called by
        # Constraint from pyomo environment and this only allows
        # to pass one variable, not a tuple/list or anything like that

        capacity = int(float(rail_constraint.split("|")[0]))
        it_opts = ast.literal_eval(rail_constraint.split("|")[1])
        # Total demand across overlapping segments
        it_opts = [(int(id_opt.split("+")[0]), int(id_opt.split("+")[1])) for id_opt in it_opts]
        train_demand = sum(
            model.x[id, opt] for id, opt in it_opts
        )

        return train_demand <= capacity


    # Divide the services between flight and rail services
    flight_services = [s for s in services if dict_mode_transport.get(s,'flight')=='flight']
    rail_services = [s for s in services if dict_mode_transport.get(s,'flight')=='rail']

    if len(flight_services) > 0:
        # Add constraints flights
        model.flight_capacity_constraints = Constraint(flight_services, rule=flight_service_capacity_rule)

    if len(rail_services) > 0:
        # Add constraints rail
        # Dictionary of rail services (only the rail not between which stops, i.e., 123_5_6 becomes 123
        # Then for each service which stops are served
        dict_rail_services = {}

        for rs in rail_services:
            # For each rail_service (e.g. 123_5_6) get list of stops and add them to the set of stops by that service
            rs_s = rs.split("_")
            rs_s[1] = int(rs_s[1])
            rs_s[2] = int(rs_s[2])
            rail_service_id = rs_s[0]  # Extract train ID
            service_stops = dict_rail_services.setdefault(rail_service_id, set())  # Use a set to store unique stops
            service_stops.update(rs_s[1:])  # add new stops
            dict_rail_services[rail_service_id] = service_stops

        # Convert sets back to sorted lists to have list of stops ordered per service id
        dict_rail_services = {k: sorted(list(v)) for k, v in dict_rail_services.items()}

        # Create rail constraints
        def build_constraints_rail_services(dict_rail_services, dict_sc):
            # This function builds the constraint for ail services as 295|['ES511_ES300_0+10'] or
            # 295|['ES512_ES511_0+14', 'ES512_ES300_0+0']
            # It iterates throw the dictionary of services (rs) and the stops they use
            # Then it goes between each succesive stops and check which rails use that
            # Then it checks whith id clusters and options use those service and adds them
            # to the list of services that if pax are assigned to them,then they should be
            # considered for the capacity of that segment/train.

            constraints_rail = []
            for (rs, stops) in dict_rail_services.items():
                for i in range(len(stops) - 1):
                    # Find all segments of the same train
                    relevant_segments = [
                        seg for seg in dict_sc if (seg.startswith(rs + "_") and
                                                   dict_mode_transport.get(seg, 'flight') == 'rail')
                    ]

                    # Find all overlapping segments
                    overlapping_segments = [
                        seg for seg in relevant_segments
                        if is_overlapping(seg, stops[i], stops[i + 1])
                    ]

                    # Could be None, i.e. seg not in service_assigments, if that service_s_s is not used
                    # by any option
                    id_opt_overlapping = [service_assignments[seg] for seg in overlapping_segments if
                                          seg in service_assignments]

                    if any(id_opt_overlapping):

                        id_opt_overlapping_no_duplicates = list(
                            set([item for sublist in id_opt_overlapping for item in sublist]))

                        if len(id_opt_overlapping_no_duplicates) > 0:
                            # Keep in constraint only id of rail not between stops
                            constraints_rail += [(str(rs)+"_"+str(stops[i])+"_"+str(stops[i+1]), id_opt_overlapping_no_duplicates)]

            # Remove duplicate constraints, as in some case it could be that the same o-d pair goes several
            # times over the same segment (if contained)
            constraints_rail = list({(k, tuple(v)) for k, v in constraints_rail})

            # Reformat into str to be able to pass it to Constraint function later
            # replacing the flight by their capacity
            constraints_rail = [
                str(dict_sc[s[0]]) + "|" + str([str(t[0]) + "+" + str(t[1]) for t in s[1]])
                for s in constraints_rail]

            return constraints_rail

        constraints_rail = build_constraints_rail_services(dict_rail_services, dict_sc)

        model.rail_capacity_constraints = Constraint(constraints_rail, rule=rail_capacity_rule)

    end_time_1 = time.time()

    logging.info(f"Time taken for creating constraint 2: {end_time_1 - start_time_1:.4f} seconds")


    return model


def reassign_pax_option_solver(pax_reassigning, dict_volume, dict_sc, dict_mode_transport, objectives, pc,
                               thresholds, solver):

    objectives_names = [o[0] for o in objectives]

    model = create_model_passenger_reassigner_pyomo(it_data=pax_reassigning.copy(),
                                                    dict_volume=dict_volume,
                                                    dict_sc=dict_sc,
                                                    dict_mode_transport=dict_mode_transport,
                                                    objectives=objectives_names,
                                                    pc=pc)

    # Count variables and constraints
    num_vars = len([v for v in model.component_data_objects(Var)])
    num_constraints = len([c for c in model.component_data_objects(Constraint)])
    print(f"Number of variables: {num_vars}")
    print(f"Number of constraints: {num_constraints}")

    objectives = [(o[0], maximize) if o[1] == 'maximize' else (o[0], minimize) for o in objectives]


    results = lexicographic_optimization(model, objectives, solver=solver, thresholds=thresholds,
                                         num_threads=pc)

    # Process results

    # Create a dictionary of (it, opt) to pax values from the model
    assigned_pax_dict = {(id, opt): model.x[id, opt].value for id, opt in model.x}

    # Convert the dictionary to a DataFrame for efficient merging
    assigned_pax_df = pd.DataFrame(
        list(assigned_pax_dict.items()), columns=["it_opt", "pax"]
    )

    # Split the tuple column into separate 'it' and 'opt' columns
    assigned_pax_df[["id", "option"]] = pd.DataFrame(assigned_pax_df["it_opt"].tolist(),
                                                     index=assigned_pax_df.index)

    # Drop the tuple column as it's no longer needed
    assigned_pax_df = assigned_pax_df.drop(columns=["it_opt"])

    assigned_pax_df.rename(columns={'id': 'pax_group_id', 'pax': 'pax_assigned'}, inplace=True)


    # Merge with the original it_data
    pax_reassigning = pax_reassigning.merge(assigned_pax_df, on=["pax_group_id", "option"], how="left")
    pax_reassigning["pax_assigned"] = pax_reassigning["pax_assigned"].apply(lambda x: 0 if x == -0.0 else x)

    # Fill missing pax values with 0
    pax_reassigning["pax_assigned"] = pax_reassigning["pax_assigned"].fillna(0)

    # Check how many passengers didn't fit
    pax_demand_assigned = pax_reassigning.groupby(['pax_group_id', 'pax'])['pax_assigned'].sum().reset_index()
    pax_demand_assigned['unfulfilled'] = pax_demand_assigned['pax'] - pax_demand_assigned['pax_assigned']
    pax_demand_assigned.rename(columns={'pax': 'demand_to_assign'}, inplace=True)

    print("Optimization Results: " + str(results))
    print("Total num pax assigned: " + str(pax_reassigning.pax_assigned.sum()))

    logging.info("Optimization Results: " + str(results))
    logging.info("Total num pax assigned: " + str(pax_reassigning.pax.sum()))

    return pax_reassigning, pax_demand_assigned