#!/usr/bin/env python

import matplotlib.pyplot as plt
import geopandas as gpd
import pandas as pd

import geopandas as gpd
import numpy as np

from fiona.crs import from_epsg

import psycopg2

plt.switch_backend('agg')

cen_con = psycopg2.connect(database = "census", user = "", password = "",
                                   host = "", port = 5432)

geo10 = gpd.read_postgis("SELECT geoid, geometry FROM PlotTract10();",
                         geom_col = "geometry", con = cen_con, crs = from_epsg(2163))


w2 = {0 : 1, 1 : 0.68, 2 : 0.22}
def t_to_w2(t, scale = 1.0):

    v = int(t / 10. / scale)
    if v > 2: return 0

    return w2[v]

w3 = {1 : 0.962, 2 : 0.704, 3 : 0.377, 4 : 0.042}
def t_to_w3(t, scale = 1.0):

    v = int(t / 10. / scale) + 1
    if v > 6: return 0
    if v > 4: v = 4

    return w3[v]


def map_format(ax, on = False):

    plt.subplots_adjust(top = 1, bottom = 0, right = 1, left = 0, hspace = 0, wspace = 0)
    plt.margins(0,0)
    ax.xaxis.set_major_locator(plt.NullLocator())
    ax.yaxis.set_major_locator(plt.NullLocator())

    if not on:
        ax.set_axis_off()
        ax.set_axis_on()
        for a in ["bottom", "top", "right", "left"]:
            ax.spines[a].set_linewidth(0)

    return ax



# Data directory.

data = "data/"


# Load populations.  Merge these together.

doc = pd.read_csv(data + "doc.csv", dtype = {"geoid" : int, "supply" : int, "region" : int})
doc.rename(columns = {"region" : "destination_region"}, inplace = True)

pop = pd.read_csv(data + "pop.csv", dtype = {"geoid" : int, "pop" : int, "region" : int})
pop.rename(columns = {"region" : "origin_region"}, inplace = True)

pop = pd.merge(doc, pop)
pop.rename(columns = {"pop" : "D", "supply" : "S"}, inplace = True)


# Create the times dataframe.

times = pd.read_csv(data + "times.csv", dtype = {'origin': int, 'destination': int, 'cost': float})


# Ensure self-times of 0.

loc = list(pop[pop.D > 0].geoid.unique())
self = pd.DataFrame({"origin" : loc, "destination" : loc, "cost" : [0] * len(loc)})

times = pd.concat([self, times], sort = False)


# Merge the files

times = times.merge(pop[["geoid", "D", "origin_region"]]     .rename(columns = {"geoid" : "origin"}),      how = "left")
times = times.merge(pop[["geoid", "S", "destination_region"]].rename(columns = {"geoid" : "destination"}), how = "left")

times["in_region"] = (times.origin_region == times.destination_region).astype(int)
times.drop(["origin_region", "destination_region"], axis = 1, inplace = True)


# Convert time costs to weights.

times["W2"] = times.cost.apply(t_to_w2)
times["W3"] = times.cost.apply(t_to_w3)


# Sum the weights and get the preferences; merge them.

W3sum = times[["origin", "W3"]].groupby("origin").sum().rename(columns = {"W3" : "W3sum"}).reset_index()

times = pd.merge(times, W3sum)
times["G"] = times.W3 / times.W3sum


# Get the total demand in an office location; merge it.

times["D2tot"] = times.W2 * times.D
times["D3tot"] = times.W3 * times.D * times.G

demand_at_dest = times[["destination", "D2tot", "D3tot"]].groupby("destination").sum().reset_index()
demand_at_dest = pd.merge(demand_at_dest, pop[["geoid", "S"]],
                          left_on = "destination", right_on = "geoid", how = "left")

# Get the supply to demand ratio, at the location; merge it.

demand_at_dest["R2"] = demand_at_dest.S / demand_at_dest.D2tot
demand_at_dest["R3"] = demand_at_dest.S / demand_at_dest.D3tot
times = pd.merge(times, demand_at_dest[["destination", "R2", "R3"]])


# Calculate the total supply per location.

times["fca2"]    = times.W2 * times.R2
times["fca2_ir"] = times.W2 * times.R2 * times.in_region

times["fca3"]    = times.G * times.W3 * times.R3
times["fca3_ir"] = times.G * times.W3 * times.R3 * times.in_region


# Sum the supply to get the region's access.

fca = times[["origin", "fca2", "fca2_ir", "fca3", "fca3_ir"]].groupby("origin").sum().reset_index()

fca["in_region_2"] = fca["fca2_ir"] / fca["fca2"]
fca.loc[fca["fca2"] == 0, "in_region_2"] = 0

fca["in_region_3"] = fca["fca3_ir"] / fca["fca3"]
fca.loc[fca["fca3"] == 0, "in_region_3"] = 0

fca.drop(["fca2_ir", "fca3_ir"], axis = 1, inplace = True)

fca = pd.merge(fca, pop[["geoid", "D"]].copy().rename(columns = {"geoid" : "origin", "D" : "pop"}))


# Make the normalized access

mean_access = (fca["fca2"] * fca["pop"]).sum() / fca["pop"].sum()
fca["fca2_norm"] = fca["fca2"] / mean_access

mean_access = (fca["fca3"] * fca["pop"]).sum() / fca["pop"].sum()
fca["fca3_norm"] = fca["fca3"] / mean_access


# Save for posterity...

fca.to_csv("fca.csv", index = False)

geo_fca = pd.merge(geo10, fca, left_on = "geoid", right_on = "origin")

ax = geo_fca.plot(column = "fca2_norm", cmap = "coolwarm_r", vmin = 0, vmax = 2)
map_format(ax)
ax.figure.savefig("js_2sfca.png", dpi = 1200)


ax = geo_fca.plot(column = "in_region_2", cmap = "coolwarm_r", legend = True, vmin = 0, vmax = 1)

map_format(ax)
ax.figure.savefig("js_2sfca_ir.png", dpi = 1200)


ax = geo_fca.plot(column = "fca3_norm", cmap = "coolwarm_r", vmin = 0, vmax = 2)
map_format(ax)
ax.figure.savefig("js_3sfca.png", dpi = 1200)


ax = geo_fca.plot(column = "in_region_3", cmap = "coolwarm_r", legend = True, vmin = 0, vmax = 1)

map_format(ax)
ax.figure.savefig("js_3sfca_ir.png", dpi = 1200)




hp = times[(times.origin == 17031410900) & (times.cost < 60)]
hp = pd.merge(geo10, hp, left_on = "geoid", right_on = "destination")

hpax = hp.plot(column = "W3", cmap = "viridis", legend = True)
map_format(hpax)
hpax.figure.savefig("hp_W3.pdf")

hpax = hp.plot(column = "G", cmap = "viridis", legend = True)
map_format(hpax)
hpax.figure.savefig("hp_G.pdf")

hpax = hp.plot(column = "S", cmap = "viridis", legend = True)
map_format(hpax)
hpax.figure.savefig("hp_S.pdf")

hpax = hp.plot(column = "R3", cmap = "viridis", legend = True)
map_format(hpax)
hpax.figure.savefig("hp_R3.pdf")

hpax = hp.plot(column = "fca3", cmap = "viridis", legend = True)
map_format(hpax)
hpax.figure.savefig("hp_fca3.pdf")


uc = times[(times.destination == 17031836200) & (times.cost < 60)]
uc = pd.merge(geo10, uc, left_on = "geoid", right_on = "origin")

ucax = uc.plot(column = "D", cmap = "viridis", legend = True)
map_format(ucax)
ucax.figure.savefig("uc_D.pdf")

ucax = uc.plot(column = "W3", cmap = "viridis", legend = True)
map_format(ucax)
ucax.figure.savefig("uc_W3.pdf")

ucax = uc.plot(column = "G", cmap = "viridis", legend = True)
map_format(ucax)
ucax.figure.savefig("uc_G.pdf")

ucax = uc.plot(column = "D3tot", cmap = "viridis", legend = True)
map_format(ucax)
ucax.figure.savefig("uc_D3tot.pdf")