import pandas as pd
import numpy as np
import plotly.graph_objects as go
import plotly.express as px
import datetime

routes_dataset_url = 'https://raw.githubusercontent.com/jpatokal/openflights/master/data/routes.dat'
airports_dataset_url = 'https://raw.githubusercontent.com/jpatokal/openflights/master/data/airports.dat'
countries_dataset_url = 'https://raw.githubusercontent.com/jpatokal/openflights/master/data/countries.dat'


# Load the file with the routes
routes = pd.read_csv(routes_dataset_url, names=['Airline','AirlineID','SourceAirport', 
                                                       'SourceAirportID', 'DestinationAirport', 'DestinationAirportID',
                                                      'CodeShare','Stops','Equipment'], 
                     na_values='\\N')
display(routes)


# Load the file with the airports
airports = pd.read_csv(airports_dataset_url, names=['AirportID','Name','City','Country','IATA','ICAO',
                                                         'Latitude', 'Longitude', 'Altitude', 'Timezone', 'DST',
                                                         'TzDatabaseTimeZone', 'Type', 'Source'], index_col='AirportID',
                      na_values='\\N')

display(airports)

# Covid19 data
covid_data = pd.read_csv('covid19-global-forecasting-week-2/train.csv',parse_dates=['Date'], 
                   date_parser=lambda x: datetime.datetime.strptime(x, '%Y-%m-%d'))

# We compute some intermediate dataframes that will be useful later

# Aggregate the number of flights per airport.
# Source_airport -> destination_airport, number of flights
flights_airport = routes.groupby(['SourceAirportID','DestinationAirportID']).size().reset_index(name='num_flights')

# Get the countries for the previous airport IDs
flights_country=flights_airport.join(airports.Country,on='SourceAirportID',lsuffix='Source').\
                join(airports.Country,on='DestinationAirportID',lsuffix='Destination')[['Country','CountryDestination','num_flights']]

# Aggregate the number of flights between countries
# SourceCountry -> DestinationCountry, number of flights
paths = flights_country.groupby(['Country','CountryDestination']).sum()
# Aggregate the number of flights to other countries.
paths_other_countries = flights_country[flights_country.Country != flights_country.CountryDestination].groupby(['Country','CountryDestination']).sum()

# Total number of flights per country
total_flights_country = paths.groupby(['Country']).sum()
# Total number of flights to other countries
total_flights_other_countries = paths_other_countries.groupby(['Country']).sum()

# Some code to have a nice progress bar :)
from IPython.display import clear_output

def update_progress(progress):
    bar_length = 20
    if isinstance(progress, int):
        progress = float(progress)
    if not isinstance(progress, float):
        progress = 0
    if progress < 0:
        progress = 0
    if progress >= 1:
        progress = 1

    block = int(round(bar_length * progress))
    clear_output(wait = True)
    text = 'Progress: [{0}] {1:.1f}%'.format( '#' * block + '-' * (bar_length - block), progress * 100)
    print(text)

fig = go.Figure()
total_rows = float(len(paths_other_countries.index))
num_row=-1
for index, row in paths_other_countries.iterrows():
    
    num_row = num_row + 1
    update_progress(num_row/total_rows)
    fig.add_trace(
        go.Scattergeo(
            locationmode = 'country names',
            locations = index,
            mode = 'lines',
            hoverinfo= 'skip',
            line = dict(width=1,color='red'),
            opacity = float(row['num_flights'])/float(total_flights_other_countries.loc[index[0]]),
            showlegend=False,
        )
    )

fig.update_layout(
    geo = dict(
        showcountries = True,
    ),
    showlegend = False,
    title = 'Aggregated air routes between countries excluding domestic flights'
)

fig.show(renderer='svg')

aux = paths_other_countries.reset_index()
aux = aux[aux.Country=='China']


fig = go.Figure()

total_rows = float(len(aux.index))
num_row=-1

the_min = aux.num_flights.min()
the_max = aux.num_flights.max()
the_mean = aux.num_flights.mean()
the_std = aux.num_flights.std()
scalator = float(the_max - the_min)

for index, row in aux.iterrows():
    num_row = num_row + 1
    update_progress(num_row/total_rows)
    width = ((float(row['num_flights'])-the_min)/scalator)*15
    #width = np.max([0.5, width])
   
    fig.add_trace(
        go.Scattergeo(
            locationmode = 'country names',
            locations = [row.Country,row.CountryDestination],
            mode = 'lines',
            hoverinfo= 'skip',
            line = dict(width=width,color='red'),
            showlegend=False,
        )
    )
    

fig.update_layout(
    geo = dict(
        showcountries = True,
    ),
    showlegend = False,
    title = 'China aggregated air routes between countries excluding domestic flights'
)

fig.show(renderer='svg')

# Run this to ensure that the drawing module is correctly stablished
from graph_tool.all import *
import graph_tool as gt

g = gt.Graph()

# Create the country property with the country names
countries_prop = g.new_vertex_property('string')
g.vertex_properties['country'] = countries_prop

# Total number of outgoing flights per vertex
out_flights_prop = g.new_vertex_property('int')
g.vertex_properties['out_flights'] = out_flights_prop

# Create edge property with the number of flights between countries
num_flights_prop = g.new_edge_property('double')
g.edge_properties['num_flights'] = num_flights_prop

# Get the complete list of countries
list_countries = np.union1d(paths.reset_index().Country.unique(), paths.reset_index().CountryDestination.unique())

g.add_vertex(len(list_countries))

# For a given country, get the vertex
country_vertex = {}
index=0
#for c in total_flights_country.index:
for c in list_countries:
    v = g.vertex(index)
    countries_prop[v] = c
    country_vertex[c] = v
    # skip self-loops
    try: 
        out_flights_prop[v] = total_flights_other_countries.loc[c]['num_flights']
    except:
        out_flights_prop[v] = 0.0
        
    index=index+1


# Add the edges
for index,num_flights in paths.iterrows():
    s = country_vertex[index[0]]
    d = country_vertex[index[1]]
    # No self-loops
    if s != d:
        e = g.add_edge(s,d)
        num_flights_prop[e] = num_flights

# plot with a radial distribution with China in the center
china_vertex = country_vertex['China']

pos = gt.draw.radial_tree_layout(g,china_vertex)
gv = gt.GraphView(g, efilt=lambda e : e.source()==china_vertex)

p=gt.draw.graph_draw(gv,pos)

second_step = []
for v in g.get_out_neighbors(china_vertex):
    second_step.append(v)
    
gv2 = gt.GraphView(g, efilt=lambda e : e.target!=china_vertex and e.source() in second_step)
p=gt.draw.graph_draw(gv2,pos)

aux = gv.edge_properties['num_flights'].copy()

aux.a = (aux.a - aux.a.min())/(aux.a.max()-aux.a.min()) * 10
p=gt.draw.graph_draw(gv,pos,edge_pen_width=aux)

def runSIR(state,num_iter=100):
    '''
    For a graph tool epidemic model run the state and return a dataframe with the countries infected in every step
    '''
    infections = pd.DataFrame()
    
    previous = state.get_state().fa
    initial = np.where(previous==1)[0]
    
    for t in range(num_iter):
        ret=state.iterate_sync()
        new = state.get_state().fa
        # Find new infected countries
        diff = new - previous
        already_infected = state.get_state().fa.sum()
        non_infected = len(state.get_state().fa)-already_infected
        new_infected = [countries_prop[g.vertex(v)] for v in np.where(diff==1)[0]]
        previous = new.copy()
        # collect the results in a Pandas friendly way
        infections = infections.append([{'step':t,'norm_steps':t/float(num_iter),'infected':already_infected,
                                'new_infected': len(new_infected),'non_infected': non_infected,
                                'new_infected_countries': new_infected}], ignore_index=True)
        
    return infections

# Create an edge property map that can help us to define the beta probability between vertices 
beta_prop = g.new_edge_property('double')
for e in g.edges():
    beta_prop[e] = num_flights_prop[e]/out_flights_prop[e.source()]

# Assume that vertices are not susceptible, except the seed.
susceptible_prop = g.new_vertex_property('float')
for v in g.vertices():
       susceptible_prop[v] = 0.0
susceptible_prop[china_vertex]=1.0

state = gt.dynamics.SIRState(g, beta=beta_prop,v0 = china_vertex, constant_beta=True,gamma=0,
                              r=0,s=susceptible_prop,epsilon=0,)

infections_SIR = runSIR(state,64)
display(infections_SIR[infections_SIR.new_infected>0].head(10))

# Compute the number of infected countries
# Get the first date with confirmed cases for every country
first_date = covid_data[covid_data.ConfirmedCases > 0].groupby('Country_Region', as_index=False).Date.min().sort_values(by='Date')
# Compute the number of elapsed days since the beginning of China's outbreak
first_date['step']=first_date.Date-datetime.datetime(2020,1,22)
# Convert this column into an integer with the number of days.
first_date.step = first_date.step.dt.days
# Get the number of infected countries per region
covid_infected = first_date[['step','Country_Region']].groupby('step').agg(new_infected=('step','count'))
covid_infected['infected'] = covid_infected['new_infected'].cumsum()
covid_infected = covid_infected.reset_index()

f = go.Figure()
f.add_trace(
    go.Scatter(x=covid_infected.step,y=covid_infected.infected,name='Covid19')
)
f.add_trace(
    go.Scatter(x=infections_SIR.step,y=infections_SIR.infected,name='SIR')
)
f.update_layout(
    title_text='Comparing SIR infected countries evolution with original Covid19',
    xaxis=dict(title='Step'),
    yaxis=dict(title='Infected countries')
)

f.show(renderer='svg')