#!/usr/bin/python3

import xpress as xp
import pandas as pd
import geopandas as gpd
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
import numpy as np
import folium

def main():
    gdf = gpd.read_file('knapsack_example.geojson')

    print("\n:: This is the data")
    print(gdf.shape)
    print(gdf.head())

    # decision variables from each Landscape Management Units (LMU)
    x_i = [xp.var(i, vartype=xp.binary) for i in gdf['LMU_ID'].astype(str)]

    # objective to maximize
    Beta_i = gdf['res_d_SPM']
    Z = xp.Sum(Beta_i * x_i)

    # constraint (i.e. LMUs selected are constrained by number of acres)
    alpha_i = gdf['Acres']
    c1 = xp.Sum(alpha_i * x_i) <= 100  #try changing me to 80

    # linear program
    print("\n\n:: This is some licensing stuff, can be ignored")
    p = xp.problem()
    p.addVariable(x_i)
    p.addConstraint(c1)
    p.setObjective(Z, sense=xp.maximize)

    print("\n\n:: This is the solution")
    p.solve()
    print(p.getObjVal())

    # solution
    var_solutions = [int(x.name) for x in x_i if p.getSolution(x) == 1]

    # plot all LMUs and the objective variable
    fig, ax = plt.subplots(figsize=(16,10))
    gdf.plot(color='None', edgecolor='k', ax=ax)
    gdf.plot(column='res_d_SPM', ax=ax, legend=True, scheme='quantiles', legend_kwds={'loc': 'lower right'})
    ax.set_title('res_d_SPM', fontsize=18)

    # highlight solution
    outline_solutions = gdf.loc[gdf['LMU_ID'].isin(var_solutions)]
    outline_solutions.plot(color='None', edgecolor='r', linewidth=5, ax=ax, hatch='///')

    fig.savefig('choropleth_map.png')

if __name__ == '__main__':
    main()