import pandas as pd
import glob
import csv
import numpy as np
import networkx as nx
from networkx.algorithms import bipartite
import pickle
import random
import collections
import sys

country = str(sys.argv[1])
typ = str(sys.argv[2])

years = range(1961,2016+1)
thresh = 0.1


if typ == 'P':
    production = pd.read_csv('data/production.csv',encoding = "ISO-8859-1")
elif typ == 'PI':
    production = pd.read_csv('data/prodPlusImp.csv',encoding = "ISO-8859-1")
elif typ == 'PIE':
    production = pd.read_csv('data/prodPlusImpMinusExp.csv',encoding = "ISO-8859-1")
else:
    print('Invalid type argument. Enter "P","PI", or "PIE".')


foodNutrients = pd.read_csv('data/foodNutrients-frac.csv')
servingSizes = pickle.load(open('data/servingSizes.pkl','rb'))
population = pd.read_csv('data/API_SP.POP.TOTL_DS2_en_csv_v2_63973.csv')

def multicurve_unweighted(H,lenL,numRuns):
    df = pd.DataFrame()
    for i in range(numRuns):
        G = H.copy()
        height = []
        height.append(len(list(set([v for u,v in list(G.edges())]))))
        order = list(G.nodes())[:lenL]
        random.shuffle(order)
        for node in order:
            G.remove_node(node)
            height.append(len(list(set([v for u,v in list(G.edges())]))))
        df[i] = height
    df['avg'] = df.mean(axis=1)
    
    #there are 17 nutrients to begin with
    return((df['avg']/17).tolist())

def countryStability(countryName):
    nutrientList = [x for x in list(foodNutrients)[4:-1] if x!='Sodium']

    stabilities = []
    for year in years:
        countryCrops = [x for x in list(set(production.loc[(production['Area']==countryName)&(production['Y'+str(year)]>0),'Item'].tolist())) if x in list(servingSizes.keys())]

        bnk = nx.OrderedGraph()
        bnk.add_nodes_from(countryCrops, bipartite=0)
        bnk.add_nodes_from(nutrientList, bipartite=0)
        edges = []
        weights = []
        for crop in countryCrops:
            newEdges = []
            for nutrient in nutrientList:
                weight = np.mean(foodNutrients.loc[foodNutrients['FAO_name']==crop,nutrient])
                weight = weight * 10**6 * 10**(-2) * (1/servingSizes[crop]) / population[population['Country Name']==countryName][str(year)].iloc[0]
                weight = weight * np.mean(production.loc[(production['Area']==countryName)&(production['Item']==crop)]['Y'+str(year)])
                if weight>0.1:
                    newEdges.append((crop, nutrient, weight))
                    weights.append(weight)
            edges.extend(newEdges)
            bnk.add_weighted_edges_from(newEdges)
            if bnk.degree[crop] == 0:
                bnk.remove_node(crop)
                countryCrops = [x for x in countryCrops if x!=crop]

        curve = multicurve_unweighted(bnk,len(countryCrops),1000)    
        stabilities.append(sum([curve[i]/len(curve) for i in range(len(curve))]))
    return stabilities


stabilities = countryStability(country)
df = pd.DataFrame()
df[country] = stabilities
df.to_csv('data/timeseries/nutriStability-timeseries-'+country+'-'+typ+'.csv')
































