# python-3

# Population Estimator Implementation on predicted shared dataset. 
# Author: Sreejith Menon (smenon8@uic.edu)
# Script creation date: Sept 30, 2016

# Help for using the script via command line:
# For running 
# 1. Population estimation when top/bottom k scored images shared using classifiers with IBIES features: 
#    python PopulationEstimatorAPI.py -cr clf 
# 2. Population estimation when random k scored images shared using classifiers with IBIES features: 
#    python PopulationEstimatorAPI.py -cr clf -r True
# 3. Population estimation when top/bottom k scored images shared using regressions with IBIES features: 
#    python PopulationEstimatorAPI.py -cr rgr 
# 4. Population estimation when random k scored images shared using regressions with IBIES features: 
#    python PopulationEstimatorAPI.py -cr rgr -r True

# 5. Population estimation when top/bottom k scored images shared using classifiers with beauty features: 
#    python PopulationEstimatorAPI.py -cr clf -b True
# 6. Population estimation when random k scored images shared using classifiers with beauty features: 
#    python PopulationEstimatorAPI.py -cr clf -r True -b True
# 7. Population estimation when top/bottom k scored images shared using regressions with beauty features: 
#    python PopulationEstimatorAPI.py -cr rgr -b True
# 8. Population estimation when random k scored images shared using classifiers with beauty features: 
#    python PopulationEstimatorAPI.py -cr rgr -r True -b True

# 9. Population estimation when images above a certain threshold are shared using regressions with IBEIS features: 
#    python PopulationEstimatorAPI.py -cr clf -t True
# 10. Population estimation when images above a certain threshold are shared using regressions with beauty features: 
#    python PopulationEstimatorAPI.py -cr clf -t True -b True

import MarkRecapHelper as MR
import pandas as pd
import json, importlib, sys, random, argparse
import ClassiferHelperAPI as CH
import DeriveFinalResultSet as DRS
import cufflinks as cf # this is necessary to link pandas to plotly
cf.go_online()
import plotly.graph_objs as go
importlib.reload(MR)
import htmltag as HT
import numpy as np, re, sys
import ClassifierCapsuleClass as ClfCls
import RegressionCapsuleClass as RgrCls

'''
    There are mainly two modes defined, one for each dataset. The schema is slightly different for both the datasets.
    This mainly applies to the SD card logic, GGR stores 
'''
MODE = 'GGR' 

if mode == 'GZC':
    inExifFl,inGidAidMapFl,inAidFtrFl = "../data/imgs_exif_data_full.json","../data/full_gid_aid_map.json","../data/full_aid_features.json"
    g_truth_population_z = 3468
    g_truth_population_g = 177
    g_truth_population_all = 3620
    days = {'2015-03-01' : 1,
        '2015-03-02' : 2 }
else: # For GGR
    inExifFl,inGidAidMapFl,inAidFtrFl = "../data/ggr_gid_uuid_exif_ftr_map.json","../data/ggr_uuid_annot_uuid_map.json","../data/ggr_annot_uuid_ftr_map.json"
    g_truth_population_z = 2620
    days = {'2016-01-29' : 1,
        '2016-01-30' : 2 }


layout = go.Layout(
    title="Number of images shared(k) versus estimated population",
    titlefont = dict(size=15),
    xaxis=dict(
        title="Number of images shared(k)",
        titlefont = dict(
        size=15),
        showticklabels=True,
        tickangle=35,
        tickfont=dict(
        size=9,
        color='black')
        ),
    yaxis=dict(
        title="Estimated Population",
        titlefont = dict(
        size=15),
        showticklabels=True,
        tickfont=dict(
        size=9,
        color='black')
        )
    )


# Simulates flipping a biased coin with a bias of p. Heads mean 1 and tail mean 0.
def biasedCoinFlipper(p):
    return 1 if random.random() < p else 0

# This method is used to actually estimate the population. 
# This method utilizes methods from MarkRecapHelper to calculate the population of all animals, zebras and giraffes 
# on the basis of the prediction results that are generated by the classifiers or regressor predictions.
def estimatePopulation(prediction_results,inExifFl,inGidAidMapFl,inAidFtrFl, days = days):
    df = pd.DataFrame(prediction_results,index=['share']).transpose().reset_index()
    df.columns = ['GID','share']
    df.to_csv("/tmp/PredictionResults.csv",index=False)

    try:
        assert MR.MODE == MODE
    except AssertionError as e:
        print("Mode mismatch, Modes should be GGR or GZC in both script files")
        print(e)
        sys.exit(-2)
        
    # Entire population estimate (includes giraffes and zebras)
    if MODE == 'GZC':
        nidMarkRecapSet = MR.genNidMarkRecapDict(inExifFl,
                    inGidAidMapFl,
                    inAidFtrFl,
                    "/tmp/PredictionResults.csv",
                    days,shareData='classifier')


        marks_all,recaptures_all,population_all,confidence_all = MR.applyMarkRecap(nidMarkRecapSet)
        population_all = None if population_all == 0 else population_all
        zebra_spec = "zebra_plains"
        nidMarkRecapSet = MR.genNidMarkRecapDict(inExifFl,
                    inGidAidMapFl,
                    inAidFtrFl,
                    "/tmp/PredictionResults.csv",
                    days,
                    filterBySpecies='giraffe_masai',
                    shareData='classifier')
        marks_g,recaptures_g,population_g,confidence_g = MR.applyMarkRecap(nidMarkRecapSet) 
        population_g = None if population_g == 0 else population_g
    else:
        zebra_spec = "zebra_grevys"

    nidMarkRecapSet = MR.genNidMarkRecapDict(inExifFl,
                    inGidAidMapFl,
                    inAidFtrFl,
                    "/tmp/PredictionResults.csv",
                    days,
                    filterBySpecies=zebra_spec,
                    shareData='classifier')

    marks_z,recaptures_z,population_z,confidence_z = MR.applyMarkRecap(nidMarkRecapSet)

    # 0 recaptures using Lincoln-Petersen estimates essentially mean an infinite population and hence marked as NULL
    population_z = None if population_z == 0 else population_z

    if MODE == "GZC":
        return {'all' : population_all , 
            'zebras' : population_z , 'giraffes' : population_g, "zebra_confidence" : confidence_z }
    else:
        return {'zebras' : population_z, "zebra_confidence" : confidence_z }

contrib_extractor = lambda x : re.findall(r'GGR,([0-9]+,[A-Z]+),.*', x)[0]

# This method is used to actually estimate the population when each contributor shares top k images. 
# The images are ranked on the basis of the likelihood estimates obtained from the regression models. 
# Since the regression outputs are a likelihood, to make a choice of whether or not each photo will be 
# shared a coin with bias equal to likelihood of the picture being shared is flipped. 
def kSharesPerContribAfterCoinFlip(prediction_results,inExifFl,inGidAidMapFl,inAidFtrFl,genk,randomShare = False):
    if MODE == 'GGR':
        with open(inExifFl, "r") as exif_fl:
            exif_obj = json.load(exif_fl)
        gidContribDct = {gid_uuid : [contrib_extractor(exif_obj[gid_uuid]['contributor'])] for gid_uuid in exif_obj.keys()} # -- the ggr dataset stores the contributor information differently (unlike GZC)
    else:  
        gidContribDct = DRS.getCountingLogic(inGidAidMapFl,inAidFtrFl,'CONTRIBUTOR',False)

    sdCards = {}
    for key in gidContribDct.keys():
        sdCards[gidContribDct[key][0]] = sdCards.get(gidContribDct[key][0],[]) + [key]

    sdCardSorted = {}
    for contrib in sdCards.keys():
        sdCardSorted[contrib] = sorted(sdCards[contrib],key=lambda x : prediction_results.get(x,0),reverse=False)

    if randomShare:
        for contrib in sdCardSorted.keys():
            sdCardSorted[contrib] = random.sample(sdCardSorted[contrib],len(sdCardSorted[contrib]))

    k = genk()
    
    predictions_k = {}
    for contrib in sdCardSorted.keys():
        pctCnt = 0
        exitCntr = 0
        while pctCnt < k and pctCnt < len(sdCardSorted[contrib]) and exitCntr <= 500:
            nextImg = sdCardSorted[contrib][pctCnt]
            exitCntr += 1
            flip = biasedCoinFlipper(prediction_results.get(nextImg,0)/100)
            if flip == 1:
                predictions_k[nextImg] = 1
                pctCnt += 1

    return estimatePopulation(predictions_k,inExifFl,inGidAidMapFl,inAidFtrFl)

# synthetic experiment #1 and #2
# This method is used to actually estimate the population when each contributor shares top k images. 
# The images are ranked on the basis of the likelihood estimates obtained from the regression models. 
# Since the regression outputs are a likelihood, to make a choice of whether or not each photo will be shared a coin 
# with bias equal to likelihood of the picture being shared is flipped. 
def kSharesPerContributor(prediction_probabs,inExifFl,inGidAidMapFl,inAidFtrFl,genk, randomShare=False):
    if MODE == 'GGR':
        with open(inExifFl, "r") as exif_fl:
            exif_obj = json.load(exif_fl)

        gidContribDct = {gid_uuid : [contrib_extractor(exif_obj[gid_uuid]['contributor'])] for gid_uuid in exif_obj.keys()} # -- the ggr dataset stores the contributor information differently (unlike GZC)
    else:  
        gidContribDct = DRS.getCountingLogic(inGidAidMapFl,inAidFtrFl,'CONTRIBUTOR',False, MODE)

    sdCards = {}
    for key in gidContribDct.keys():
        sdCards[gidContribDct[key][0]] = sdCards.get(gidContribDct[key][0],[]) + [key]

    sdCardSorted = {}
    for contrib in sdCards.keys():
        sdCardSorted[contrib] = sorted(sdCards[contrib],key=lambda x : prediction_probabs.get(x,0),reverse=False)

    if randomShare:
        for contrib in sdCardSorted.keys():
            sdCardSorted[contrib] = random.sample(sdCardSorted[contrib],len(sdCardSorted[contrib]))

    # selecting the top k scores as shared
    predictions_k = {gid : 1 for contrib in sdCardSorted.keys() for gid in sdCardSorted[contrib][:genk()]}

    return estimatePopulation(predictions_k,inExifFl,inGidAidMapFl,inAidFtrFl)

def shareAbvThreshold(prediction_probabs,inExifFl,inGidAidMapFl,inAidFtrFl,threshold,randomShare=None):
    if MODE == 'GGR':
        with open(inExifFl, "r") as exif_fl:
            exif_obj = json.load(exif_fl)
        gidContribDct = {gid_uuid : [exif_obj[gid_uuid]['contributor']] for gid_uuid in exif_obj.keys()} # -- the ggr dataset stores the contributor information differently (unlike GZC)
    else:  
        gidContribDct = DRS.getCountingLogic(inGidAidMapFl,inAidFtrFl,'CONTRIBUTOR',False)

    sdCards = {}
    for key in gidContribDct.keys():
        sdCards[gidContribDct[key][0]] = sdCards.get(gidContribDct[key][0],[]) + [key]

    sdCardSorted = {}
    for contrib in sdCards.keys():
        sdCardSorted[contrib] = sorted(sdCards[contrib],key=lambda x : prediction_probabs.get(x,0),reverse=True)

    threshold = threshold()
    for contrib in sdCardSorted.keys():
        sdCardSorted[contrib] = list(filter(lambda x : prediction_probabs.get(x,0) >= threshold,sdCards[contrib]))

    # selecting the top k scores as shared
    predictions_k = {gid : 1 for contrib in sdCardSorted.keys() for gid in sdCardSorted[contrib]}

    return estimatePopulation(predictions_k,inExifFl,inGidAidMapFl,inAidFtrFl)

def runSyntheticExpts(isClf, methTypes, attribTypes, krange, methArgs, thresholdMeth = False, randomShare = False):
    if len(attribTypes) == 1 and attribTypes[0] == 'beauty':
        train_fl = "../feature_files/EXPT2_FEATURE_VECTOR_TRAINING.csv"
        test_fl = "../feature_files/GGR_FEATURE_VECTOR.csv"
    else:
        train_fl = "../FinalResults/ImgShrRnkListWithTags.csv"
        test_fl = "../data/full_gid_aid_ftr_agg.csv"

    if isClf:
        trainTestMeth = CH.trainTestClf
    else:
        trainTestMeth = CH.trainTestRgrs

    for meth in methTypes:
        for attrib in attribTypes:
            print("Starting to run %s on test data\nAttribute Selection Method : %s" %(meth,attrib))

            methObj,predResults = trainTestMeth(train_fl,
                                test_fl,
                                meth,
                                attrib,
                                infoGainFl="../data/infoGainsExpt2.csv",
                                methArgs = methArgs
                                )
            if thresholdMeth:
                flNm = str("../FinalResults/"+ meth + "_" + attrib + "_thresholded")
            else:
                if randomShare:
                    flNm = str("../FinalResults/"+ meth + "_" + attrib + "_kSharesRandom")
                else:
                    flNm = str("../FinalResults/"+ meth + "_" + attrib + "_bottom_kShares")

            flNm += MODE
            print(flNm)
            if isClf:
                predictions = {list(methObj.test_x.index)[i] : methObj.predProbabs[i] for i in range(len(methObj.test_x.index))}
                sharesMethod = kSharesPerContributor
            else:
                predictions = predResults
                if thresholdMeth:
                    sharesMethod = shareAbvThreshold
                else:
                    sharesMethod = kSharesPerContribAfterCoinFlip
            if randomShare:
                hdr = "Population Estimates when contributor shares k random photos using %s and attribute selection method %s" %(meth,attrib)
            else:
                if thresholdMeth:
                    hdr = "Population Estimates when contributors share all images above threshold using %s and attribute selection method %s" %(meth,attrib)
                else:
                    hdr = "Population Estimates with top k shares per contributor using %s and attribute selection method %s" %(meth,attrib)

            print("Starting population estimation experiments")
            
            fixedK = {k : sharesMethod(predictions,inExifFl,inGidAidMapFl,inAidFtrFl,lambda : k,randomShare=randomShare) for k in krange}
            print("Population estimation experiments complete")
            
            df = pd.DataFrame(fixedK).transpose().reset_index()

            if not thresholdMeth:
                if MODE == "GZC":
                    df.columns = ['num_images','all','giraffes','zebras', 'zebra_confidence']
                else:
                    df.columns = ['num_images','zebras', 'zebra_confidence']
                df.index = df['num_images']
                df.drop(['num_images'],1,inplace=True)
            else:
                # df.columns = ['threshold','all','giraffes','zebras']
                df.columns = ['threshold', 'zebras', 'confidence']
                df.index = df['threshold']
                df.drop(['threshold'],1,inplace=True)

            df.to_csv(str(flNm+".csv"))
            # df_html = df.to_html(index=True)
            randomized = sharesMethod(predictions,inExifFl,inGidAidMapFl,inAidFtrFl,lambda : random.randint(min(krange),max(krange)))
            if MODE == "GZC":
                df['Randomized_giraffe'] = randomized['giraffes']
                df['Randomized_all'] = randomized['all']
            df['Randomized_zebras'] = randomized['zebras']
            
            # return df

            # if thresholdMeth:
            #     fig1 = df.iplot(kind='line',layout=layout,filename=str(meth + "_" + attrib + "_thresholded"))
            # else:
            #     if randomShare:
            #         fig1 = df.iplot(kind='line',layout=layout,filename=str(meth + "_" + attrib + "_kSharesRandom"))
            #     else:
            #         fig1 = df.iplot(kind='line',layout=layout,filename=str(meth + "_" + attrib + "_kShares"))


            # fullFl = HT.HTML(HT.body(HT.h2(hdr),
            #                 HT.HTML(df_html),
            #                 HT.HTML(fig1.embed_code)         
            #         ))

            # outputFile = open(str(flNm+".html"),"w")
            # outputFile.write(fullFl)
            # outputFile.close()
            # print("Testing complete %s : %s" %(meth,attrib))
            # print()

def runSyntheticExptsRgr(inExifFl, inGidAidMapFl, inAidFtrFl, krange, thresholdMeth=False, randomShare=False, beautyFtrs = False):
    rgrTypes = ['linear', 'ridge', 'lasso', 'svr', 'dtree_regressor', 'elastic_net']
    rgrTypes = ['linear',  'elastic_net']
    if beautyFtrs:
        attribTypes = ['beauty']
    else:
        attribTypes = ['sparse','non_sparse','non_zero','abv_mean']

    regrArgs = {'linear' : {'fit_intercept' : True},
            'ridge' : {'fit_intercept' : True},
            'lasso' : {'fit_intercept' : True},
            'elastic_net' : {'fit_intercept' : True},
            'svr' : {'fit_intercept' : True},
            'dtree_regressor' : {'fit_intercept' : True}}

    return runSyntheticExpts(False, rgrTypes, attribTypes, krange, regrArgs, thresholdMeth=thresholdMeth, randomShare=randomShare)


def runSyntheticExptsClf(inExifFl, inGidAidMapFl, inAidFtrFl, krange, randomShare=False, beautyFtrs = False):
    clfTypes = ['dummy', 'bayesian', 'logistic', 'svm', 'dtree', 'random_forests', 'ada_boost']
    clfTypes = ['logistic', 'ada_boost']
    if beautyFtrs:
        attribTypes = ['beauty']
    else:
        attribTypes = ['sparse','non_sparse','non_zero','abv_mean']
    
    clfArgs = {'dummy' : {'strategy' : 'most_frequent'},
            'bayesian' : {'fit_prior' : True},
            'logistic' : {'penalty' : 'l2'},
            'svm' : {'kernel' : 'rbf','probability' : True},
            'dtree' : {'criterion' : 'entropy'},
            'random_forests' : {'n_estimators' : 10 },
            'ada_boost' : {'n_estimators' : 50 }}

    runSyntheticExpts(True, clfTypes, attribTypes, krange, clfArgs, randomShare=randomShare)

def buildErrPlots(clfOrRgr, thresholdMeth=False, randomShare=False):
    if clfOrRgr == 'clf':
        algTypes = ['bayesian','logistic','svm','dtree','random_forests','ada_boost']
    else:
        algTypes = ['linear','ridge','lasso','svr','dtree_regressor','elastic_net']
    attribTypes = ['sparse','non_sparse','non_zero','abv_mean', 'beauty']
    flNms = [str(alg + "_" + attrib) for alg in algTypes for attrib in attribTypes]

    if thresholdMeth:
        suffix = "_thresholded.csv"
        hdr = "threshold"
        if clfOrRgr == 'clf':
            titleSuffix = "classifiers thresholded"
        else:
            titleSuffix = "regressors thresholded"
    else:
        hdr = "num_images"
        if randomShare:
            suffix = "_kSharesRandom.csv"
            if clfOrRgr == 'clf':
                titleSuffix = "classifiers random choices"
            else:
                titleSuffix = "regressors random choices"
        else:
            suffix = "_kShares.csv"
            if clfOrRgr == 'clf':
                titleSuffix = "classifiers top k choices"
            else:
                titleSuffix = "regressors top k choices"

    df = pd.DataFrame.from_csv(str("../FinalResults/"+flNms[0]+suffix)).reset_index()
    df.columns = list(map(lambda x : str(x + "_" + flNms[0]) if x != hdr else x,list(df.columns)))
    for i in range(1,len(flNms)):
        df1 = pd.DataFrame.from_csv(str("../FinalResults/"+flNms[i]+suffix)).reset_index()
        df1.columns = list(map(lambda x : str(x + "_" + flNms[i]) if x != hdr else x,list(df1.columns)))
        df = pd.DataFrame.merge(df,df1,on=hdr)

    df.index = df[hdr]
    df.drop([hdr],1,inplace=True)
    

    # calculate errors in estimation
    # % error = (predicted - actual) * 100 / actual
    for col in df.columns:
        if 'all' in col:
            df[str(col+'_err')] = (df[col] - 3620) / 36.20
        elif 'zebras' in col:
            df[str(col+'_err')] = (df[col] - 3468) / 34.68
        elif 'giraffes' in col:
            df[str(col+'_err')] = (df[col] - 177) / 1.77

    figs=[]
    errorCols = [col for col in df.columns if 'err' in col]
    df = df[errorCols]

    for alg in algTypes:
        algCol = [col for col in df.columns if alg in col]
        algDf = df[algCol]
        titleAlg = "All %s %s" %(alg,titleSuffix)
        figs.append(algDf.iplot(kind='line',title=titleAlg))

    for attrib in attribTypes:
        attribCol = [col for col in df.columns if attrib in col]
        attribDf = df[attribCol]
        titleAttrib = "All %s %s" %(attrib,titleSuffix)
        figs.append(attribDf.iplot(kind='line',title=titleAttrib))

    figCodes = [fig.embed_code for fig in figs]
    return figCodes

def populationEstimatorFileGen():
    clfTypes = ['bayesian','logistic','svm','dtree','random_forests','ada_boost']
    attribTypes = ['sparse','non_sparse','non_zero','abv_mean']
    
    kwargsDict = {'dummy' : {'strategy' : 'most_frequent'},
            'bayesian' : {'fit_prior' : True},
            'logistic' : {'penalty' : 'l2'},
            'svm' : {'kernel' : 'rbf','probability' : True},
            'dtree' : {'criterion' : 'entropy'},
            'random_forests' : {'n_estimators' : 10 },
            'ada_boost' : {'n_estimators' : 50 }}

    if len(sys.argv) > 1:
        outFlNm = sys.argv[1]
    else:
        outFlNm = "/tmp/PopulationEstimatorFromClf.dump.json"

    estimates = []
    for clf in clfTypes:
        for attribType in attribTypes:
            print("Classifier : %s" %clf)
            print("Attribute selection method : %s" %attribType)
            clfObj,predResults = trainTestClf("../FinalResults/ImgShrRnkListWithTags.csv",
                     "../data/full_gid_aid_ftr_agg.csv",
                     clf,
                     attribType,
                     "../data/infoGainsExpt2.csv",
                     kwargsDict)
            thisObjhead = {'Classifier' : clf , 'Attribute' : attribType,'shared_images_count' : int(sum(clfObj.preds))}
            thisObj = estimatePopulation(predResults,
                            "../data/imgs_exif_data_full.json",
                            "../data/full_gid_aid_map.json",
                            "../data/full_aid_features.json")

            estimates.append({**thisObjhead,**thisObj})
            print("Complete")
            print()

    with open(outFlNm,"w") as jsonFl:
        json.dump(estimates,jsonFl,indent=4)

def __main__():
    parser = argparse.ArgumentParser()
    parser.add_argument('-cr','--clfOrRgr', help='Classifier method or regressor to be used', required=True)
    parser.add_argument('-t','--threshold', help='Threshold method, if set will be used', required=False)
    parser.add_argument('-r','--random', help='Share any random images, if set will be used', required=False)
    parser.add_argument('-b','--beauty', help='Use only the beauty features, if set will be used', required=False)

    args = vars(parser.parse_args())
    if args["beauty"] != None:
        beauty = True
    else:
        beauty = False

    if args["clfOrRgr"] == 'clf' and args["random"] != None:
        print("Running classifiers with random image selections")
        runSyntheticExptsClf(inExifFl,inGidAidMapFl,inAidFtrFl,range(2,75),randomShare=True,beautyFtrs=beauty)

    if args["clfOrRgr"] == 'clf' and args["random"] == None:
        print("Running classifiers with bottom k image selections")
        runSyntheticExptsClf(inExifFl,inGidAidMapFl,inAidFtrFl,range(2,75),beautyFtrs=beauty)

    if args["clfOrRgr"] == 'rgr' and args["random"] != None:
        print("Running regressors with random image selections")
        runSyntheticExptsRgr(inExifFl,inGidAidMapFl,inAidFtrFl,range(2,75),randomShare = True,beautyFtrs=beauty)
    
    if args["clfOrRgr"] == 'rgr' and args["random"] == None and args["threshold"] != None:
        print("Running regressors with top image above thresholds")
        runSyntheticExptsRgr(inExifFl,inGidAidMapFl,inAidFtrFl,range(40,90,5),thresholdMeth = True,beautyFtrs=beauty)
    
    if args["clfOrRgr"] == 'rgr' and args["random"] == None and args["threshold"] == None:
        print("Running regressors with bottom k image selections")
        runSyntheticExptsRgr(inExifFl,inGidAidMapFl,inAidFtrFl,range(2,75),thresholdMeth = False,beautyFtrs=beauty)



if __name__ == "__main__":
    __main__()