#import dependencies
import glob
import itertools
import json
import multiprocessing as mp
import os
import pickle
import re
import sys
import time
import warnings

import numpy as np
import pandas as pd
from scipy import stats
from TempFolder.TempFolder import Temp
from tqdm import tqdm

from collections import namedtuple

warnings.filterwarnings('ignore')

with open('nogit\\path','r') as f:
    path_to_csv = f.read()

#set up path
Temp.set_path(path_to_csv +'SC-temp\\')

#load input files
exact_match = ['Make','Model_short','Model_Age','Prim_Damage']

ModelType = namedtuple('ModelType',['model_name','smalb','big_b','match_within','pct_match','truncated'])

models_pop = [
    ModelType('Pop_10p_500k1m_with_odm_no_trunc',5e5,1e6,.1,['Actual_Cash_Value_adj','Odometer_Replace'],False)
    ,ModelType('Pop_05p_500k1m_with_odm_no_trunc',5e5,1e6,.05,['Actual_Cash_Value_adj','Odometer_Replace'],False)
    ,ModelType('Pop_20p_500k1m_with_odm_no_trunc',5e5,1e6,.20,['Actual_Cash_Value_adj','Odometer_Replace'],False)

    ,ModelType('Pop_10p_1m1m_with_odm_trunc',1e6,1e6,.1,['Actual_Cash_Value_adj','Odometer_Replace'],True)
    ,ModelType('Pop_05p_1m1m_with_odm_trunc',1e6,1e6,.05,['Actual_Cash_Value_adj','Odometer_Replace'],True)
    ,ModelType('Pop_20p_1m1m_with_odm_trunc',1e6,1e6,.2,['Actual_Cash_Value_adj','Odometer_Replace'],True)

    ,ModelType('Pop_10p_500k1m_with_odm_trunc',5e5,1e6,.1,['Actual_Cash_Value_adj','Odometer_Replace'],True)
    ,ModelType('Pop_05p_500k1m_with_odm_trunc',5e5,1e6,.05,['Actual_Cash_Value_adj','Odometer_Replace'],True)
    ,ModelType('Pop_20p_500k1m_with_odm_trunc',5e5,1e6,.2,['Actual_Cash_Value_adj','Odometer_Replace'],True)

    ,ModelType('Pop_10p_1m1m_no_odm_trunc',1e6,1e6,.1,['Actual_Cash_Value_adj'],True)
    ,ModelType('Pop_05p_1m1m_no_odm_trunc',1e6,1e6,.05,['Actual_Cash_Value_adj'],True)
    ,ModelType('Pop_20p_1m1m_no_odm_trunc',1e6,1e6,.2,['Actual_Cash_Value_adj'],True)
    ,ModelType('Pop_05p_1m1m_no_odm_no_trunc',1e6,1e6,.05,['Actual_Cash_Value_adj'],False)

    ,ModelType('Pop_10p_5e51m_no_odm_trunc',5e5,1e6,.1,['Actual_Cash_Value_adj'],True)
    ,ModelType('Pop_05p_5e51m_no_odm_trunc',5e5,1e6,.05,['Actual_Cash_Value_adj'],True)
    ,ModelType('Pop_20p_5e51m_no_odm_trunc',5e5,1e6,.2,['Actual_Cash_Value_adj'],True)
    ,ModelType('Pop_05p_5e51m_no_odm_no_trunc',5e5,1e6,.05,['Actual_Cash_Value_adj'],False)

]

models_crime = [
    ModelType('Crime_10p_5k_with_odm_no_trunc',5e3,5e3,.1,['Actual_Cash_Value_adj','Odometer_Replace'],False)
    ,ModelType('Crime_05p_5k_with_odm_no_trunc',5e3,5e3,.05,['Actual_Cash_Value_adj','Odometer_Replace'],False)
    ,ModelType('Crime_20p_5k_with_odm_no_trunc',5e3,5e3,.20,['Actual_Cash_Value_adj','Odometer_Replace'],False)

    ,ModelType('Crime_10p_5k_with_odm_trunc',5e3,5e3,.1,['Actual_Cash_Value_adj','Odometer_Replace'],True)
    ,ModelType('Crime_05p_5k_with_odm_trunc',5e3,5e3,.05,['Actual_Cash_Value_adj','Odometer_Replace'],True)
    ,ModelType('Crime_20p_5k_with_odm_trunc',5e3,5e3,.2,['Actual_Cash_Value_adj','Odometer_Replace'],True)

    ,ModelType('Crime_10p_5k_no_odm_trunc',5e3,5e3,.1,['Actual_Cash_Value_adj'],True)
    ,ModelType('Crime_05p_5k_no_odm_trunc',5e3,5e3,.05,['Actual_Cash_Value_adj'],True)
    ,ModelType('Crime_20p_5k_no_odm_trunc',5e3,5e3,.2,['Actual_Cash_Value_adj'],True)
    ,ModelType('Crime_05p_5k_no_odm_no_trunc',5e3,5e3,.05,['Actual_Cash_Value_adj'],False)

]

models_unemp = [
    ModelType('Unemp_10p_5k_with_odm_no_trunc',5.,5.,.1,['Actual_Cash_Value_adj','Odometer_Replace'],False)
    ,ModelType('Unemp_05p_5k_with_odm_no_trunc',5.,5.,.05,['Actual_Cash_Value_adj','Odometer_Replace'],False)
    ,ModelType('Unemp_20p_5k_with_odm_no_trunc',5.,5.,.20,['Actual_Cash_Value_adj','Odometer_Replace'],False)

    ,ModelType('Unemp_10p_5k_with_odm_trunc',5.,5.,.1,['Actual_Cash_Value_adj','Odometer_Replace'],True)
    ,ModelType('Unemp_05p_5k_with_odm_trunc',5.,5.,.05,['Actual_Cash_Value_adj','Odometer_Replace'],True)
    ,ModelType('Unemp_20p_5k_with_odm_trunc',5.,5.,.2,['Actual_Cash_Value_adj','Odometer_Replace'],True)

    ,ModelType('Unemp_10p_5.1m_no_odm_trunc',5.,5.,.1,['Actual_Cash_Value_adj'],True)
    ,ModelType('Unemp_05p_5.1m_no_odm_trunc',5.,5.,.05,['Actual_Cash_Value_adj'],True)
    ,ModelType('Unemp_20p_5.1m_no_odm_trunc',5.,5.,.2,['Actual_Cash_Value_adj'],True)
    ,ModelType('Unemp_05p_5.1m_no_odm_no_trunc',5.,5.,.05,['Actual_Cash_Value_adj'],False)
]

models_income = [
    ModelType('Income_10p_5k_with_odm_no_trunc',5e4,5e4,.1,['Actual_Cash_Value_adj','Odometer_Replace'],False)
    ,ModelType('Income_05p_5k_with_odm_no_trunc',5e4,5e4,.05,['Actual_Cash_Value_adj','Odometer_Replace'],False)
    ,ModelType('Income_20p_5k_with_odm_no_trunc',5e4,5e4,.20,['Actual_Cash_Value_adj','Odometer_Replace'],False)

    ,ModelType('Income_10p_5k_with_odm_trunc',5e4,5e4,.1,['Actual_Cash_Value_adj','Odometer_Replace'],True)
    ,ModelType('Income_05p_5k_with_odm_trunc',5e4,5e4,.05,['Actual_Cash_Value_adj','Odometer_Replace'],True)
    ,ModelType('Income_20p_5k_with_odm_trunc',5e4,5e4,.2,['Actual_Cash_Value_adj','Odometer_Replace'],True)

    ,ModelType('Income_10p_5e41m_no_odm_trunc',5e4,5e4,.1,['Actual_Cash_Value_adj'],True)
    ,ModelType('Income_05p_5e41m_no_odm_trunc',5e4,5e4,.05,['Actual_Cash_Value_adj'],True)
    ,ModelType('Income_20p_5e41m_no_odm_trunc',5e4,5e4,.2,['Actual_Cash_Value_adj'],True)
    ,ModelType('Income_05p_5e41m_no_odm_no_trunc',5e4,5e4,.05,['Actual_Cash_Value_adj'],False)
]

run_on = 'CENSUS_closest'
# run_on = 'GRNDTOT_closest'
# run_on = 'Unemp_pct_closest'
# run_on = 'Income_adj_closest'
models_dict = dict(CENSUS_closest = models_pop
                    ,GRNDTOT_closest = models_crime
                    ,Unemp_pct_closest=models_unemp
                    ,Income_adj_closest = models_income
                    )
def main():
    for id, model in enumerate(models_dict[run_on]):

        car_demo_joined = Temp.load_obj('car_demo_joined')

        car_demo_joined.loc[car_demo_joined[run_on]<=model.smalb , 'group'] = 0
        car_demo_joined.loc[car_demo_joined[run_on]>model.big_b, 'group'] = 1
        car_demo_joined = car_demo_joined.dropna(subset = ['group'])

        compare_columns = model.pct_match + ['Price_Sold_or_Highest_Bid_adj']
        #truncate quantiles
        original_size =car_demo_joined.shape[0]
        if model.truncated:
            for col_ in compare_columns:
                car_demo_joined = car_demo_joined.loc[car_demo_joined[col_] < car_demo_joined[col_].quantile(0.99) ]

        print('Model name: ',model.model_name,' #',id,' of ',len(models_dict[run_on]),' \n'
            ,'                                               small county boundary:' , model.smalb, ' bigg county boundary', model.big_b,
            '\nRows excluded from the analysis:',original_size - car_demo_joined.shape[0]
        ,'\nLooking for a match within ', model.match_within,' for ', model.pct_match
        ,'\nExact match: ', exact_match
        )

        def find_match_expected(small_c, big_c):
            '''
            Find expected match - for each row in the `small_c` (small county) dataset find ID of a matching values from the `big_c` dataset

            Params:
            small_c: pd.DataFrame
            big_c: pd.DataFrame

            Returns:
            pd.DataFrame
            '''
            def find_match_for_row(x):
                # try:
                bs = big_c.copy()
                # import pdb; pdb.set_trace()
                bs[model.pct_match[0]] = np.abs((x[model.pct_match[0]]/bs[model.pct_match[0]])-1)

                if len(model.pct_match)>1:
                    bs[model.pct_match[1]]  = np.abs((x[model.pct_match[1]]/bs[model.pct_match[1]])-1)
                    bs = bs.fillna(9999)
                    bs = bs.loc[
                                (bs[model.pct_match[0]]  <=model.match_within )
                                & (bs[model.pct_match[1]]  <=model.match_within )
                                ].index.tolist()
                else:
                    bs = bs.fillna(9999)
                    bs = bs.loc[
                                (bs[model.pct_match[0]]  <=model.match_within )
                                ].index.tolist()
                return bs

            return small_c.apply(find_match_for_row,axis=1)

        #Iterate over all dataframe and find matches
        match_l = []
        i = 0
        for grp , df in tqdm(car_demo_joined.groupby(exact_match)):
            if len(df) <2:
                continue
            group0_df = df.loc[df.group == 0]
            group1_df = df.loc[df.group == 1]

            match_l.append(find_match_expected(group0_df,group1_df))
            i+=1

            # if i>550:
            #     break

        #concat to the dataframe
        match_null = pd.concat(match_l).dropna(axis=1).iloc[:,0].rename('match')

        #filter out rows without match
        match_e = match_null.loc[match_null.map(lambda x: len(x)>0)]

        #save calculation results
        Temp.save_obj(match_null,'matched_candidates_'+model.model_name)

        print('% of records matched: ', 100*len(match_e) / len(match_null)
        ,'# of records matched:', len(match_e)

        ,'%\n',match_e.head())

        #expected match

        def expected_mean_f(idx):
            '''
            The method select the number of observations and calculates the average value across them
            '''
            return car_demo_joined.loc[idx,compare_columns].mean()

        expected_mean = (match_e.apply(expected_mean_f)
                        .join(car_demo_joined[compare_columns],rsuffix='_matched')
                                )

        # expected_match = (match_e.apply(expected_mean).rename('expected_value_big_city').to_frame()
        #                 .join(car_demo_joined.Price_Sold_or_Highest_Bid_adj.rename('actual_value_small_city')))

        #save output
        # Temp.save_obj(expected_match,'expected_match_'+model.model_name)

        Temp.save_obj(expected_mean,'expected_match_'+model.model_name)

        # expected_match = Temp.load_obj('expected_match_'+model.model_name)

        #calculate difference between small and the large cities

        for ccol in compare_columns:
            expected_match_diff = expected_mean[[ccol,ccol+'_matched']].diff(axis=1).dropna(axis=1,how='all').dropna()
            print('\n-'+ccol,'\n'
            ,'mean value: ',round(expected_mean[ccol].mean(),2),'\n'
            ,'mean value (matched): ',round(expected_mean[ccol+'_matched'].mean(),2),'\n'
            ,expected_match_diff.agg(['mean','std']).rename({ccol+'_matched':'difference'},axis=1)
            ,'\nt-test p value:', round(stats.ttest_1samp(expected_match_diff.values.reshape(-1,),0)[1],3)#,'\n',expected_mean.head()
            )

        del car_demo_joined , expected_match_diff , expected_mean , match_null , match_e , match_l

if __name__ == "__main__":
    main()


#!/usr/bin/env python

import os
import time as t
import subprocess
import argparse

import numpy as np

from processing_components.image.operations import import_image_from_fits, export_image_to_fits, qa_image
from processing_components.visibility.operations import append_visibility

from ska_sip.metamorphosis.filter import uv_cut, uv_advice
from ska_sip.telescopetools.initinst import init_inst
from ska_sip.accretion.ms import load
from ska_sip.pipelines.dprepb import dprepb_imaging, arl_data_future

import dask
import dask.array as da
from dask.distributed import Client
from distributed.diagnostics import progress
from distributed import wait

__author__ = "Jamie Farnes"
__email__ = "jamie.farnes@oerc.ox.ac.uk"

"""
This code is something of a "Workflow Script Wrapper".

Execution control would load this code. At the moment, Execution Control is essentially 'python -i pipe.py'!
"""

def main(args):
    """
    Initialising launch sequence.
    """
    # ------------------------------------------------------
    # Print some stuff to show that the code is running:
    print("")
    os.system("printf 'A demonstration of a \033[5mDPrepB/DPrepC\033[m SDP pipeline\n'")
    print("")
    # Set the directory for the moment images:
    MOMENTS_DIR = args.outputs + '/MOMENTS'
    # Check that the output directories exist, if not then create:
    os.makedirs(args.outputs, exist_ok=True)
    os.makedirs(MOMENTS_DIR, exist_ok=True)
    # Set the polarisation definition of the instrument:
    POLDEF = init_inst(args.inst)

    # Setup Variables for SIP services
    # ------------------------------------------------------
    # Define the Queue Producer settings:
    if args.queues:
        queue_settings = {'bootstrap.servers': 'scheduler:9092', 'message.max.bytes': 100000000}  #10.60.253.31:9092

    # Setup the Confluent Kafka Queue
    # ------------------------------------------------------
    if args.queues:
        from confluent_kafka import Producer
        import pickle
        # Create an SDP queue:
        sip_queue = Producer(queue_settings)

    # Define a Data Array Format
    # ------------------------------------------------------
    def gen_data(channel):
        return np.array([vis1[channel], vis2[channel], channel, None, None, False, False, args.plots, float(args.uvcut), float(args.pixels), POLDEF, args.outputs, float(args.angres), None, None, None, None, None, None, args.twod, npixel_advice, cell_advice])

    # Setup the Dask Cluster
    # ------------------------------------------------------
    starttime = t.time()

    dask.config.set(get=dask.distributed.Client.get)
    client = Client(args.daskaddress)  # scheduler for Docker container, localhost for P3.

    print("Dask Client details:")
    print(client)
    print("")

    # Define channel range for 1 subband, each containing 40 channels:
    channel_range = np.array(range(int(args.channels)))

    # Load the data into memory:
    """
    The input data should be interfaced with Buffer Management.
    """
    print("Loading data:")
    print("")
    vis1 = [load('%s/%s' % (args.inputs, args.ms1), range(channel, channel+1), POLDEF) for channel in range(0, int(args.channels))]
    vis2 = [load('%s/%s' % (args.inputs, args.ms2), range(channel, channel+1), POLDEF) for channel in range(0, int(args.channels))]

    # Prepare Measurement Set
    # ------------------------------------------------------
    # Combine MSSS snapshots:
    vis_advice = append_visibility(vis1[0], vis2[0])

    # Apply a uv-distance cut to the data:
    vis_advice = uv_cut(vis_advice, float(args.uvcut))
    npixel_advice, cell_advice = uv_advice(vis_advice, float(args.uvcut), float(args.pixels))

    # Begin imaging via the Dask cluster
    # ------------------------------------------------------
    # Submit data for each channel to the client, and return an image:

    # Scatter all the data in advance to all the workers:
    """
    The data here could be passed via Data Queues.
    Queues may not be ideal. Data throughput challenges.
    Need to think more about the optimum approach.
    """
    print("Scatter data to workers:")
    print("")
    big_job = [client.scatter(gen_data(channel)) for channel in channel_range]

    # Submit jobs to the cluster and create a list of futures:
    futures = [client.submit(dprepb_imaging, big_job[channel], pure=False, retries=3) for channel in channel_range]
    """
    The dprepb_imaging function could generate QA, logging, and pass this information via Data Queues.
    Queues work well for this.
    Python logging calls are preferable. Send them to a text file on the node.
    Run another service that watches that file. Or just read from standard out.
    The Dockerisation will assist with logs.
    """

    print("Imaging on workers:")
    # Watch progress:
    progress(futures)

    # Wait until all futures are complete:
    wait(futures)

    # Check that no futures have errors, if so resubmit:
    for future in futures:
        if future.status == 'error':
            print("ERROR: Future", future, "has 'error' status, as:")
            print(client.recreate_error_locally(future))
            print("Rerunning...")
            print("")
            index = futures.index(future)
            futures[index].cancel()
            futures[index] = client.submit(dprepb_imaging, big_job[index], pure=False, retries=3)

    # Wait until all futures are complete:
    wait(futures)

    # Gather results from the futures:
    results = client.gather(futures, errors='raise')

    # Run QA on ARL objects and produce to queue:
    if args.queues:
        print("Adding QA to queue:")
        for result in results:
            sip_queue.produce('qa', pickle.dumps(qa_image(result), protocol=2))

        sip_queue.flush()

    # Return the data element of each ARL object, as a Dask future:
    futures = [client.submit(arl_data_future, result, pure=False, retries=3) for result in results]

    progress(futures)

    wait(futures)

    # Calculate the Moment images
    # ------------------------------------------------------
    # Now use 'distributed Dask arrays' in order to parallelise the Moment image calculation:
    # Construct a small Dask array for every future:
    print("")
    print("Calculating Moment images:")
    print("")
    arrays = [da.from_delayed(future, dtype=np.dtype('float64'), shape=(1, 4, 512, 512)) for future in futures]

    # Stack all small Dask arrays into one:
    stack = da.stack(arrays, axis=0)

    # Combine chunks to reduce overhead - is initially (40, 1, 4, 512, 512):
    stack = stack.rechunk((1, 1, 4, 64, 64))

    # Spread the data around on the cluster:
    stack = client.persist(stack)
    # Data is now coordinated by the single logical Dask array, 'stack'.

    # Save the Moment images:
    """
    The output moment images should be interfaced with Buffer Management.

    Need to know more about the Buffer specification.
    Related to initial data distribution also/staging.
    """
    print("Saving Moment images to disk:")
    print("")
    # First generate a template:
    image_template = import_image_from_fits('%s/imaging_dirty_WStack-%s.fits' % (args.outputs, 0))

    # Output mean images:
    # I:
    image_template.data = stack[:, :, 0, :, :].mean(axis=0).compute()
    # Run QA on ARL objects and produce to queue:
    if args.queues:
        sip_queue.produce('qa', pickle.dumps(qa_image(image_template), protocol=2))
    # Export the data to disk:
    export_image_to_fits(image_template, '%s/Mean-%s.fits' % (MOMENTS_DIR, 'I'))

    # Q:
    image_template.data = stack[:, :, 1, :, :].mean(axis=0).compute()
    # Run QA on ARL objects and produce to queue:
    if args.queues:
        sip_queue.produce('qa', pickle.dumps(qa_image(image_template), protocol=2))
    # Export the data to disk:
    export_image_to_fits(image_template, '%s/Mean-%s.fits' % (MOMENTS_DIR, 'Q'))

    # U:
    image_template.data = stack[:, :, 2, :, :].mean(axis=0).compute()
    # Run QA on ARL objects and produce to queue:
    if args.queues:
        sip_queue.produce('qa', pickle.dumps(qa_image(image_template), protocol=2))
    # Export the data to disk:
    export_image_to_fits(image_template, '%s/Mean-%s.fits' % (MOMENTS_DIR, 'U'))

    # P:
    image_template.data = da.sqrt((da.square(stack[:, :, 1, :, :]) + da.square(stack[:, :, 2, :, :]))).mean(axis=0).compute()
    # Run QA on ARL objects and produce to queue:
    if args.queues:
        sip_queue.produce('qa', pickle.dumps(qa_image(image_template), protocol=2))
    # Export the data to disk:
    export_image_to_fits(image_template, '%s/Mean-%s.fits' % (MOMENTS_DIR, 'P'))

    # Output standard deviation images:
    # I:
    image_template.data = stack[:, :, 0, :, :].std(axis=0).compute()
    # Run QA on ARL objects and produce to queue:
    if args.queues:
        sip_queue.produce('qa', pickle.dumps(qa_image(image_template), protocol=2))
    # Export the data to disk:
    export_image_to_fits(image_template, '%s/Std-%s.fits' % (MOMENTS_DIR, 'I'))

    # Q:
    image_template.data = stack[:, :, 1, :, :].std(axis=0).compute()
    # Run QA on ARL objects and produce to queue:
    if args.queues:
        sip_queue.produce('qa', pickle.dumps(qa_image(image_template), protocol=2))
    # Export the data to disk:
    export_image_to_fits(image_template, '%s/Std-%s.fits' % (MOMENTS_DIR, 'Q'))

    # U:
    image_template.data = stack[:, :, 2, :, :].std(axis=0).compute()
    # Run QA on ARL objects and produce to queue:
    if args.queues:
        sip_queue.produce('qa', pickle.dumps(qa_image(image_template), protocol=2))
    # Export the data to disk:
    export_image_to_fits(image_template, '%s/Std-%s.fits' % (MOMENTS_DIR, 'U'))

    # P:
    image_template.data = da.sqrt((da.square(stack[:, :, 1, :, :]) + da.square(stack[:, :, 2, :, :]))).std(axis=0).compute()
    # Run QA on ARL objects and produce to queue:
    if args.queues:
        sip_queue.produce('qa', pickle.dumps(qa_image(image_template), protocol=2))
    # Export the data to disk:
    export_image_to_fits(image_template, '%s/Std-%s.fits' % (MOMENTS_DIR, 'P'))

    # Flush queue:
    if args.queues:
        sip_queue.flush()

    # Make a tarball of moment images:
    subprocess.call(['tar', '-cvf', '%s/moment.tar' % (MOMENTS_DIR), '%s/' % (MOMENTS_DIR)])
    subprocess.call(['gzip', '-9f', '%s/moment.tar' % (MOMENTS_DIR)])

    endtime = t.time()
    print(endtime-starttime)

# Define the arguments for the pipeline:
ap = argparse.ArgumentParser()
ap.add_argument('-d', '--daskaddress', help='Address of the Dask scheduler [default scheduler:8786]', default='scheduler:8786')
ap.add_argument('-c', '--channels', help='Number of channels to process [default 40]', default=40)

ap.add_argument('-inp', '--inputs', help='Input data directory [default /data/inputs]', default='/data/inputs')
ap.add_argument('-out', '--outputs', help='Output data directory [default /data/outputs]', default='/data/outputs')
ap.add_argument('-ms1', '--ms1', help='Measurement Set 1 [default sim-1.ms]', default='sim-1.ms')
ap.add_argument('-ms2', '--ms2', help='Measurement Set 2 [default sim-2.ms]', default='sim-2.ms')
ap.add_argument('-q', '--queues', help='Enable Queues? [default False]', default=True)
ap.add_argument('-p', '--plots', help='Output diagnostic plots? [default False]', default=False)
ap.add_argument('-2d', '--twod', help='2D imaging [True] or wstack imaging [False]? [default False]', default=False)

ap.add_argument('-uv', '--uvcut', help='Cut-off for the uv-data [default 450]', default=450.0)
ap.add_argument('-a', '--angres', help='Force the angular resolution to be consistent across the band, in arcmin FWHM [default 8.0]', default=8.0)
ap.add_argument('-pix', '--pixels', help='The number of pixels/sampling across the observing beam [default 5.0]', default=5.0)
ap.add_argument('-ins', '--inst', help='Instrument name (for future use) [default LOFAR]', default='LOFAR')
args = ap.parse_args()
main(args)


import numpy as np
from annoy import AnnoyIndex
import os
import re

class _Annoy:

    def __init__(self, annoy_path):
        self.annoy_path = annoy_path
        self.annoy = 0

    def save_to_annoy(self, n_trees=1000, embs_path='CelebsEmbeddings/', metrica='euclidean'):
        """
        Метод переводит сохраненные эмбендинги в .ann
        :param n_trees:
        :param embs_path:
        :param metrica:
        :return:
        """
        dict_of_files = {}
        new_filelist = []
        for name in os.listdir(embs_path):
            dict_of_files[int(re.match('[0-9]+', name).group(0))] = name
        for el in sorted(dict_of_files.items()):
            new_filelist.append(el[1])

        annoy = AnnoyIndex(128, metrica)

        for file in new_filelist:
            data = np.load(embs_path+file)
            embedding, index = data['arr_0'], data['arr_1']
            for idx in range(index.shape[0]):
                annoy.add_item(index[idx], embedding[idx])

        annoy.build(n_trees)
        annoy.save(self.annoy_path)
        print('saved correctly at {}'.format(self.annoy_path))

    def load_annoy(self, metrica='euclidean'):
        """
        Загрузка модели annoy
        :param metrica:
        :return:
        """
        self.annoy = AnnoyIndex(128, metrica)
        print('Loading annoy...')
        self.annoy.load(self.annoy_path)

    def get_nns(self, embedding, n_neighbors=3):
        """
        Поиск ближайших эмбендингов (фото звезд)
        :param embedding:
        :param n_neighbors:
        :return: indexes of neighbors celebrities
        """
        indexes = self.annoy.get_nns_by_vector(embedding, n_neighbors)
        return indexes


import logging
import utils.gpu as gpu
from model.build_model import Build_Model
from model.loss.yolo_loss import YoloV4Loss
import torch
import torch.optim as optim
from torch.utils.data import DataLoader
import utils.datasets as data
import time
import random
import argparse
from eval.evaluator import *
from utils.tools import *
from tensorboardX import SummaryWriter
import config.yolov4_config as cfg
from utils import cosine_lr_scheduler
from utils.log import Logger
# from apex import amp
import torchvision.models as models
from eval_coco import *
from eval.cocoapi_evaluator import COCOAPIEvaluator
import pdb
import ap
import random
import pruning
import numpy as np

def detection_collate(batch):
    targets = []
    imgs = []
    for sample in batch:
        imgs.append(sample[0])
        targets.append(sample[1])
    return torch.stack(imgs, 0), targets

def seed_setup(args):

    print('INFO: Seeds : [{}]'.format(args.seed))
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)

class Trainer(object):
    def __init__(self, weight_path, resume, gpu_id, accumulate, fp_16, args):
        init_seeds(0)
        seed_setup(args)
        self.args = args
        self.fp_16 = fp_16
        self.device = gpu.select_device(gpu_id)
        self.start_epoch = 0
        self.best_mAP = 0.0
        self.accumulate = accumulate
        self.weight_path = weight_path
        self.multi_scale_train = cfg.TRAIN["MULTI_SCALE_TRAIN"]
        if self.multi_scale_train:
            print("Using multi scales training")
        else:
            print("train img size is {}".format(cfg.TRAIN["TRAIN_IMG_SIZE"]))
        self.train_dataset = data.Build_Dataset(
            anno_file_type="train", img_size=cfg.TRAIN["TRAIN_IMG_SIZE"]
        )
        self.epochs = (
            cfg.TRAIN["YOLO_EPOCHS"]
            if cfg.MODEL_TYPE["TYPE"] == "YOLOv4"
            else cfg.TRAIN["Mobilenet_YOLO_EPOCHS"]
        )
        self.train_dataloader = DataLoader(
            self.train_dataset,
            batch_size=cfg.TRAIN["BATCH_SIZE"],
            num_workers=cfg.TRAIN["NUMBER_WORKERS"],
            shuffle=True,
            pin_memory=True,
        )

        self.yolov4 = Build_Model(weight_path=weight_path, resume=resume).to(self.device)
        self.optimizer = optim.SGD(
            self.yolov4.parameters(),
            lr=cfg.TRAIN["LR_INIT"],
            momentum=cfg.TRAIN["MOMENTUM"],
            weight_decay=cfg.TRAIN["WEIGHT_DECAY"],
        )

        self.criterion = YoloV4Loss(
            anchors=cfg.MODEL["ANCHORS"],
            strides=cfg.MODEL["STRIDES"],
            iou_threshold_loss=cfg.TRAIN["IOU_THRESHOLD_LOSS"],
        )

        self.scheduler = cosine_lr_scheduler.CosineDecayLR(
            self.optimizer,
            T_max=self.epochs * len(self.train_dataloader),
            lr_init=cfg.TRAIN["LR_INIT"],
            lr_min=cfg.TRAIN["LR_END"],
            warmup=cfg.TRAIN["WARMUP_EPOCHS"] * len(self.train_dataloader),
        )
        if resume:
            self.__load_resume_weights(weight_path)

    def __load_resume_weights(self, weight_path):

        last_weight = os.path.join(os.path.split(weight_path)[0], "last.pt")
        chkpt = torch.load(last_weight, map_location=self.device)
        self.yolov4.load_state_dict(chkpt["model"])

        self.start_epoch = chkpt["epoch"] + 1
        if chkpt["optimizer"] is not None:
            self.optimizer.load_state_dict(chkpt["optimizer"])
            self.best_mAP = chkpt["best_mAP"]
        del chkpt

    def train(self):
        global writer
        logger.info(
            "Training start,img size is: {:d},batchsize is: {:d},work number is {:d}".format(
                cfg.TRAIN["TRAIN_IMG_SIZE"],
                cfg.TRAIN["BATCH_SIZE"],
                cfg.TRAIN["NUMBER_WORKERS"],
            )
        )
        logger.info("Train datasets number is : {}".format(len(self.train_dataset)))

        print('Loading base network...')
        if self.args.imp_type == 'imagenet':
            print("-" * 100)
            print("Use resnet50 [imagnet] weight")
            print("-" * 100)

        elif self.args.imp_type == 'moco':
            print("-" * 100)
            print("load resnet50 [moco] weight")
            print("-" * 100)
            moco_ckpt = torch.load('./weight/moco_v2_800ep_pretrain.pth.tar', map_location='cuda')['state_dict']
            resnet50_pytorch = models.resnet50(pretrained=False)
            moco_state_dict = {k[17:] : v for k , v in moco_ckpt.items() if k[17:] in resnet50_pytorch.state_dict().keys()}
            overlap_state_dict = {k : v for k , v in moco_state_dict.items() if k in self.yolov4._Build_Model__yolov4.backbone.state_dict().keys()}
            print("MOCO Overlap[{}/{}]".format(overlap_state_dict.keys().__len__(), self.yolov4._Build_Model__yolov4.backbone.state_dict().keys().__len__()))
            ori = self.yolov4._Build_Model__yolov4.backbone.state_dict()
            ori.update(overlap_state_dict)
            self.yolov4._Build_Model__yolov4.backbone.load_state_dict(ori)

        elif self.args.imp_type == 'simclr':
            print("-" * 100)
            print("load resnet50 [SimCLR] weight")
            print("-" * 100)
            base_weights = torch.load('./weight/simclr_weight.pt', map_location='cuda')['state_dict']
            overlap_state_dict = {k : v for k , v in base_weights.items() if k in self.yolov4._Build_Model__yolov4.backbone.state_dict().keys()}
            print("SimCLR Overlap[{}/{}]".format(overlap_state_dict.keys().__len__(), self.yolov4._Build_Model__yolov4.backbone.state_dict().keys().__len__()))
            self.yolov4._Build_Model__yolov4.backbone.load_state_dict(overlap_state_dict)

        else: assert False

        OUTDIR = 'seed{}_{}_imp_ckpt'.format(self.args.seed, self.args.imp_type)
        OUTDIRS = OUTDIR + '/model{}'.format(self.args.imp_num)

        if not os.path.exists(OUTDIRS): os.makedirs(OUTDIRS)

        if self.args.imp_num == 0:
            print("Begining [{}] IMP:[{}]".format(self.args.imp_type, self.args.imp_num))
            pruning.save_dicts_and_masks(imp_num=self.args.imp_num, model=self.yolov4, name='./'+ OUTDIR +'/backbone_ori.pth')
        else:

            print("Begining [{}] IMP:[{}]".format(self.args.imp_type, self.args.imp_num))
            load_name = './'+ OUTDIR +'/backbone_imp{}.pth'.format(self.args.imp_num)
            print("Load All CKPT Dir: {}".format(load_name))
            ###### load all others ########
            all_ckpt = torch.load(load_name, map_location="cuda")
            model_state_dict = self.yolov4.state_dict()
            overlap_all = {k : v for k , v in all_ckpt['all_state_dict'].items() if k in model_state_dict.keys()}
            model_state_dict.update(overlap_all)
            self.yolov4.load_state_dict(model_state_dict)
            ###### load conv1 ########
            self.yolov4._Build_Model__yolov4.backbone.conv1.load_state_dict(all_ckpt['first_conv1'])
            ###### pruning ######
            mask_dict = pruning.extract_mask(all_ckpt['backbone'])
            print("Load Mask Len:[{}]".format(len(mask_dict.keys())))
            pruning.imp_pruning_yolo_resnet50(self.yolov4._Build_Model__yolov4.backbone, mask_dict)
            pruning.see_zero_rate(self.yolov4._Build_Model__yolov4.backbone)
            print("-" * 100)
            print("Finish Process!")
            print("-" * 100)

        if self.args.resume_all:
            self.start_epoch = pruning.resume_begin(self.yolov4, self.optimizer, OUTDIRS)
            logger.info(" =======  Resume  Training at Epoch:[{}]  ======".format(self.start_epoch))
        else:
            logger.info(" =======  Start  Training   ======")
        start = time.time()

        print("=========EVAL=========")
        mAP = 0.0
        with torch.no_grad():
            OUTDIRS += "_eval_debug"
            if not os.path.exists(OUTDIRS): os.makedirs(OUTDIRS)
            APs, inference_time = Evaluator(self.yolov4, showatt=False, result_path=OUTDIRS).APs_voc()

            for i in APs:
                logger.info("{} --> mAP : {}".format(i, APs[i]))
                mAP += APs[i]
            mAP = mAP / self.train_dataset.num_classes
            logger.info("mAP : {}".format(mAP))

if __name__ == "__main__":
    global logger, writer
    parser = argparse.ArgumentParser()
    parser.add_argument("--resume_all",action="store_true",default=False,help="resume training flag")
    parser.add_argument('--imp_type', default='imagenet', type=str, help='imagenet, moco, simclr')
    parser.add_argument('--mask_dir', default='', type=str, help='imagenet, moco, simclr')
    parser.add_argument('--imp_num', default=-1, type=int, help='mask number')
    parser.add_argument('--seed', default=1, type=int, help='Seed')
    parser.add_argument(
        "--weight_path",
        type=str,
        default="weight/mobilenetv2.pth",
        help="weight file path",
    )  # weight/darknet53_448.weights
    parser.add_argument(
        "--resume",
        action="store_true",
        default=False,
        help="resume training flag",
    )
    parser.add_argument(
        "--gpu_id",
        type=int,
        default=0,
        help="whither use GPU(0) or CPU(-1)",
    )
    parser.add_argument("--log_path", type=str, default="log/", help="log path")
    parser.add_argument(
        "--accumulate",
        type=int,
        default=2,
        help="batches to accumulate before optimizing",
    )
    parser.add_argument(
        "--fp_16",
        type=bool,
        default=False,
        help="whither to use fp16 precision",
    )

    opt = parser.parse_args()
    writer = SummaryWriter(logdir=opt.log_path + "/event")
    logger = Logger(
        log_file_name=opt.log_path + "/log.txt",
        log_level=logging.DEBUG,
        logger_name="YOLOv4",
    ).get_log()

    Trainer(
        weight_path=opt.weight_path,
        resume=opt.resume,
        gpu_id=opt.gpu_id,
        accumulate=opt.accumulate,
        fp_16=opt.fp_16,
        args=opt

    ).train()


from db import db
import datetime
from enum import Enum
from scipy.interpolate import interp1d

author_coll = db['author']  # 获得collection的句柄
video_coll = db['video']  # 获得collection的句柄
fans_variation_coll = db['fans_variation']  # 获得collection的句柄

def dateRange(beginDate, endDate):
  dates = []
  date = beginDate
  while date <= endDate:
    dates.append(date.date())
    date += datetime.timedelta(1)
  return dates

class FansWatcher(object):
  def __insert_event(self, delta_rate, d_daily, author, info, date):
    print('变化率：{}% \n单日涨幅：{} \nUP主：{} \n信息：{}\n日期：{}\n\n'.format(
        delta_rate, d_daily, author['name'], info, date))
    out_data = {
        'variation': int(d_daily),
        'mid': author['mid'],
        'author': author['name'],
        'face': author['face'],
        'deltaRate': delta_rate,
        'datetime': date.strftime("%Y-%m-%d"),
        'info': info,
    }

    videos = video_coll.find({'mid': author['mid']})
    temp_video = {}
    cause = {'type': 'video'}
    for each_v in videos:
      # 相差一日之内
      if type(each_v['datetime']) == str:
        pass
      elif (date - each_v['datetime']).days >= -1 and (date - each_v['datetime']).days <= 7:
        temp_video['aid'] = each_v['aid']
        temp_video['title'] = each_v['title']
        temp_video['pic'] = each_v['pic']
        temp_video['cView'] = each_v['data'][0]['view']
        temp_video['channel'] = each_v['channel']
        temp_video['subChannel'] = each_v['subChannel']
        if 'cView' not in temp_video or 'aid' not in cause or temp_video[
                'cView'] > cause['cView']:
          cause['aid'] = temp_video['aid']
          cause['title'] = temp_video['title']
          cause['pic'] = temp_video['pic']
          cause['cView'] = temp_video['cView']
          cause['channel'] = temp_video['channel']
          cause['subChannel'] = temp_video['subChannel']

    if cause != {'type': 'video'}:
      out_data['cause'] = cause
    fans_variation_coll.replace_one(
        {'mid': out_data['mid'], 'datetime': out_data['datetime']}, out_data, upsert=True)

  def __judge(self, author):
    '''
        一共有这样几种可能：
            1、 大量涨粉        日涨粉数超过上周平均的25倍
            2、 史诗级涨粉      日涨粉数超过上周平均的50倍或单日涨粉超过10W
            3、 传说级涨粉      日涨粉数超过上周平均的100倍或单日涨粉超过20W
            4、 急转直下        上升轨道中的UP主突然掉粉
            5、 大量掉粉        每日掉粉数突破5K
            6、 雪崩级掉粉      每日掉粉数突破2W
            7、 末日级掉粉      每日掉粉数突破5W
            8、 新星爆发         日涨粉超过粉丝总数的20%
    '''

    data = sorted(author['data'], key=lambda x: x['datetime'])
    start_date = data[0]['datetime'].timestamp()
    end_date = data[-1]['datetime'].timestamp()
    x = []
    y = []
    for each in data:
      x.append(each['datetime'].timestamp())
      y.append(each['fans'])
    if len(x) <= 1:
      return
    # 线性插值
    interrupted_fans = interp1d(x, y, kind='linear')
    temp_date = datetime.datetime.fromtimestamp(start_date)
    c_date = datetime.datetime(
        temp_date.year, temp_date.month, temp_date.day).timestamp() + 86400 * 3
    if c_date - 86400 * 2 <= start_date:
      return
    while (c_date <= end_date):
      date = datetime.datetime.fromtimestamp(c_date)
      daily_array = interrupted_fans([c_date - 86400, c_date])
      p_daily_array = interrupted_fans(
          [c_date - 86400 * 2, c_date - 86400])

      # 24小时前涨粉数
      pd_daily = p_daily_array[1] - p_daily_array[0]

      # 每日涨粉数
      d_daily = daily_array[1] - daily_array[0]

      if (d_daily >= 5000 or d_daily <= -2000):

        delta_rate = round(d_daily / pd_daily * 100, 2)
        if (d_daily >= daily_array[1] * 0.50):
          self.__insert_event(round(d_daily/daily_array[1]*100, 2), d_daily,
                              author, '新星爆发', date)

        if (d_daily <= 0 and pd_daily >= 0):
          self.__insert_event('-', d_daily,
                              author, '急转直下', date)
          c_date += 86400
          continue

        if (d_daily <= -50000):
          # 每日掉粉数突破5K
          self.__insert_event(delta_rate, d_daily,
                              author, '末日级掉粉', date)
          pass
        elif (d_daily <= -20000):
          # 每日掉粉数突破2W
          self.__insert_event(delta_rate, d_daily,
                              author, '雪崩级掉粉', date)
          pass
        elif (d_daily <= -5000):
          # 每日掉粉数突破5W
          self.__insert_event(delta_rate, d_daily,
                              author, '大量掉粉', date)
          pass

        if (c_date >= start_date + 86400 * 8 and d_daily > 0):
          weekly_array = interrupted_fans([
              c_date - 86400 * 8, c_date - 86400])
          # 上月平均涨粉数
          weekly_mean = (weekly_array[1] - weekly_array[0]) / 7
          # 上周平均涨粉数
          delta_rate = round(d_daily / weekly_mean * 100, 2)
          if delta_rate >= 10000 or d_daily >= 200000:
            # 日涨粉数超过上日的100倍
            self.__insert_event(delta_rate, d_daily,
                                author, '传说级涨粉', date)
            pass
          elif delta_rate >= 5000 or d_daily >= 100000:
            # 日涨粉数超过上日的50倍
            self.__insert_event(delta_rate, d_daily,
                                author, '史诗级涨粉', date)
            pass
          elif delta_rate >= 2500:
            # 日涨粉数超过上日的25倍
            self.__insert_event(delta_rate, d_daily,
                                author, '大量涨粉', date)
            pass

      c_date += 86400
      pass

  def watchAllAuthor(self):
    author_filter = {'data': {'$exists': True}}
    self.__judge_author(author_filter)

  def watchBigAuthor(self):
    author_filter = {'data': {'$exists': True}, 'cFans': {'$gt': 10000}}
    self.__judge_author(author_filter)

  def __judge_author(self, author_filter):
    author_cursor = author_coll.find(author_filter)
    count = author_cursor.count()
    a = author_coll.aggregate([
        {
            '$match': author_filter
        }, {
            '$project': {
                "mid": 1,
                "face": 1,
                "name": 1,
                "data": {
                    "$filter": {
                        "input":
                        "$data",
                        "as": "data",
                        "cond": {"$gt": ["$$data.datetime", datetime.datetime.utcnow() + datetime.timedelta(hours=8)-datetime.timedelta(32)]}
                    }
                }
            }
        }, {
            "$match": {
                "data.0": {
                    "$exists": True
                }
            }
        }
    ])
    print("待爬取作者数量：{}".format(count))
    for each_author in a:
      print(each_author['mid'])
      self.__judge(each_author)
    pass

if __name__ == "__main__":
  w = FansWatcher()
  w.watchBigAuthor()
  pass


#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Tue May 16 16:29:55 2017

@author: haolin
"""

import numpy as np

import sys

import cv2
import json
from collections import defaultdict

from sklearn import svm
from sklearn.externals import joblib
from sklearn import preprocessing
from scipy.signal import medfilt

from haolin.HFEClass import HistFeatureExtractor
from haolin.FDClass import FaceDetector

import os

class DCU_EmotionService(object):
    def __init__(self, feature_type='LBP', face_det_type='cv'):

        # @vlaand
        self.local_path = os.path.dirname(os.path.abspath(__file__))
        self._file_paths = {
            'arousal_model': os.path.join(self.local_path,'arousal_model.pkl') ,
            'valence_model': os.path.join(self.local_path,'valence_model.pkl'),
            'fea_scaler': os.path.join(self.local_path,'feature_scaler.pkl')
            }
        # @vlaand

        self.hist_fea_ext = HistFeatureExtractor(method = feature_type)
        self.face_detector = FaceDetector(method = face_det_type)
        self.FACE_SIZE = (76,76)
        self.ki=4
        self.kj=4
        self.num_fea = 59
        # load arousal model
        self.arousal_model = joblib.load(self._file_paths['arousal_model'])
        # load valence model
        self.valence_model = joblib.load(self._file_paths['valence_model'])
        # load feature scaler
        self.fea_scaler = joblib.load(self._file_paths['fea_scaler'])

    def defDictOfDict(self):
        return defaultdict(dict)

    # To Do: add post processing method
    def post_processing_predication():
        return

    def analysis_video(self, path_to_video_file, vis=False):
        cap = cv2.VideoCapture(path_to_video_file)
        total_frame = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        face = np.zeros(self.FACE_SIZE).astype('uint8')
        allFrameData = defaultdict(self.defDictOfDict)
        frameID = 0

        # Extract features from given video
        print 'Extracting features .......'
        while(True):
            # Capture frame-by-frame
            ret, frame = cap.read()
            if not ret:
                break
            # Our operations on the frame come here
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            [x,y,w,h] = self.face_detector.detect_face(gray)

            if sum([x,y,w,h]) == 0:
                lbp_fea = np.zeros((1, self.num_fea * self.ki * self.kj))
            else:
                face = gray[y:y+h, x:x+w]
                face = cv2.resize(face, self.FACE_SIZE)

                lbp_fea = self.hist_fea_ext.ExtFea(face)
                lbp_fea = self.fea_scaler.transform(lbp_fea)

            # Predict arousal, valence value use pre-trained model
            arousal_res = self.arousal_model.predict(lbp_fea)
            valence_res = self.valence_model.predict(lbp_fea)
            # print frameID, arousal_res, valence_res
            allFrameData[frameID][0]['emotion'] = {"pad:arousal":arousal_res[0], "pad:pleasure":valence_res[0]}

            frameID = frameID + 1

            if vis:
                cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
                cv2.putText(frame, "arousal {0:.2f}, valence {0:.2f}".format(arousal_res[0], valence_res[0]),
                                              (x - 10, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
                cv2.imshow('frame', frame)
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break

        cap.release()

        if vis:
            cv2.destroyAllWindows()

        return json.dumps(allFrameData)


# coding: utf-8
from package.utils import KPIPoint
from package.utils import KPISet
from package.utils import KPITest
import pandas as pd
import numpy as np
from math import *
import random

class HotSpot():

    def __init__(self, KPIPoint, L, pt, m):
        # All the KPI values of elements when the total KPI is found anomalous
        self._KPIPoint = KPIPoint
        # maximum ID of Layer
        self._max_layer_id = L
        # potential score threshold
        self._ps_threshold = pt
        # maximum number of Iteration
        self._max_iteration = m
        # the actual leaf value vertor
        self._actaul_vector = []
        # the feature leaf value vertor
        self._feature_vector = []
        self.cal_vector()

    # find the root cause set
    def find_root_cause_set(self):
        RSets = None
        BSets = {}
        candidateSet = []
        for l in range(1, self._max_layer_id + 1):
            # Parallel Execution in each cuboid
            """
            # 暴力搜索
            elements_set = self._KPIPoint.get_elements_coms_by_layer(l)
            """
            # MCTS搜索
            # prune strategy
            elements_set = self._KPIPoint.get_elements_set_by_layer_with_prune(l, candidateSet)
            print('layer #%d all elements: ' % l, elements_set)

            for cuboid in elements_set:
                # Calculate Potential Scores ps(ek) of each element ek
                # Sort ek in a descending order of ps(ek)
                ps_set = {}
                for ele in elements_set[cuboid]:
                    ps = self.cal_potential_scores(ele, elements_set[cuboid])
                    ps_set[ele] = ps

                temp = sorted(ps_set.items(), key=lambda x: x[1], reverse=True)
                ps_set_sorted = {}
                for i in range(len(temp)):
                    ps_set_sorted[temp[i][0]] = temp[i][1]
                print('sorted ps set: ', ps_set_sorted)
                state = self.State(ps_set_sorted)
                bestSet, bestPS = self.MCTS(state, self._max_iteration, True)
                parentList = []
                for e in bestSet:
                    parentList.append(e)
                candidateSet = parentList
                BSets[bestSet] = bestPS
                print('best set', bestSet)
                print('\n')
                # break
                # Obtain BSetl,j
                # Prune e_c in layer l+1 whose father e_f are not in BSetl,j
                # if All the e_c in layer l+1 are pruned then break
        print('best set is : ', BSets)
        print('candidate set are :', candidateSet)
        # Choose RSet form BSetl,j with the largest ps
        # ps(RSet) = Max{ps(BSetl,j)}

        RSets = sorted(BSets.items(), key=lambda c : c[1], reverse=True)[:3]
        return RSets

    # calculate the potential score of element
    def cal_potential_scores(self, elements_set, value):
        print('elements_set: ', elements_set)
        print(value)
        deduced_leaf = {}
        for leaf in self._KPIPoint._leaf:
            deduced_leaf[leaf] = self._KPIPoint._leaf[leaf][1]
        visited = {}
        for element in elements_set:
            # print('element', element)
            deduced_value = 0
            # if element is in LEAF
            if len(element) == len(self._KPIPoint._attribute_names):
                # print('element is in LEAF')
                deduced_leaf[element] = self._KPIPoint._leaf[element][0]
            else:
                leaves, _ = self._KPIPoint.get_descendant_elements_ele(element)
                # print('leaves involved', leaves)
                for leaf in leaves:
                    if leaf not in visited:
                        visited[leaf] = 0
                    visited[leaf] = visited[leaf] + 1
                    # print('leaf: ', leaf)
                    inter = list(set(element).intersection(set(leaf)))
                    if len(inter) == 0:
                        # print('inter %d', len(inter))
                        deduced_value = self._KPIPoint._leaf[leaf][1]
                    else:
                        f_c = value[(element,)][1]
                        if f_c == 0:
                            print('error, f(c) == 0')
                            continue
                        f_l = self._KPIPoint._leaf[leaf][1]
                        a_c = value[(element,)][0]
                        deduced_value = f_l - (f_c - a_c) * (f_l / f_c)
                    # print('deduced_value %f' % deduced_value)
                    deduced_leaf[leaf] = deduced_value

        # update deduced vector
        deduced_vector = []
        for leaf in self._KPIPoint._leaf:
            deduced_value = deduced_leaf[leaf]
            deduced_vector.append(deduced_value)
        # print('visited: ', visited)
        dis_a_v = self.cal_euclidean_distance(self._actaul_vector, deduced_vector)
        dis_a_f = self.cal_euclidean_distance(self._actaul_vector, self._feature_vector)
        rate = dis_a_v / dis_a_f
        ps = np.max([1 - rate, 0])
        # print('ps of ', elements_set, ' is', ps, '\n')
        return ps

    # calculate potential scores Revised
    def cal_potential_scores_revised(self, ele, value):
        return

    # calculate the actual and predict vector
    def cal_vector(self):
        self._actaul_vector = []
        self._feature_vector = []
        for leaf in self._KPIPoint._leaf:
            self._actaul_vector.append(self._KPIPoint._leaf[leaf][0])
            self._feature_vector.append(self._KPIPoint._leaf[leaf][1])
        print(self._actaul_vector,
              self._feature_vector)

    # calculate the Euclidean distance
    def cal_euclidean_distance(self, v1, v2):
        vec1 = np.array(v1)
        vec2 = np.array(v2)
        return np.linalg.norm(vec1 - vec2)

    class State:
        """ A state of current tree
        """

        def __init__(self, elements=None):
            # all the elements to select
            self._elements = elements
            # record the visited elements
            self._visited_elements = []
            # record the visited times of edge(s,a)
            self._edge_visited = {}

        # create a deep clone of this state.
        def Clone(self):
            state = HotSpot.State(self._elements)
            state._edge_visited = self._edge_visited
            return state

        # visit an element
        def visit_element(self, ele):
            self._visited_elements.append(ele)

        # visit an to_ele from from_ele
        def update_edge_visits(self, from_ele, to_ele):
            if from_ele is None:
                from_ele = ()
            if (from_ele, to_ele) not in self._edge_visited:
                self._edge_visited[(from_ele, to_ele)] = 0
            self._edge_visited[(from_ele, to_ele)] += 1

        # get the unvisited elements
        def get_unvisited_elements(self):
            unvisited_elements = {}
            for element in self._elements:
                if element not in self._visited_elements:
                    unvisited_elements[element] = self._elements[element]
            return unvisited_elements

        # get the element with max potential score.
        def get_ele_with_max_ps(self):
            return sorted(self._elements.items(), key=lambda x: x[1])[-1]

    class Node:
        """ A node in the tree.
        """

        def __init__(self, elements=None, parent=None, state=None):
            self._elementsSet = elements  # the elements that got us to this node - "None" for the root node
            self._parent_node = parent  # "None" for the root node
            self._child_nodes = []
            self._score = state.get_ele_with_max_ps()[1]
            self._visits = 0
            self._state = state  # all nodes share one state
            self._unvisited_elements = state.get_unvisited_elements()  # future child nodes

        def select_child(self):
            """ Use the UCB formula to select a child node.
            """
            # if N(s,a) = 0, then assign a probability of taking unvisited actions to be
            # R = (1 − Q(s,amax)), where amax = argmaxa∈A(s)∩N(s,a)=0 Q(s,a).
            # print("parent: ", self._elementsSet, "child: ", self._child_nodes)
            # print(self._state._edge_visited)
            node = sorted(self._child_nodes,
                          key=lambda c: c._score +
                                        sqrt(2 * log(c._visits)
                                             / self._state._edge_visited[(self._elementsSet, c._elementsSet)]))[-1]
            if len(self._unvisited_elements) == 0:
                return node

            prob = random.random()
            element = sorted(self._unvisited_elements.items(), key=lambda x: x[1])[-1]
            if prob > 1 - element[1]:
                print('explore')
                self._state.visit_element(element[0])
                node = self.add_child(element[0], self._state)
            return node

        def add_child(self, e, s):
            """ Remove m from _unvisited_elements and add a new child node for this elements.
                Return the added child node
                * We choose e∗ to have the largest ps(S) value of the remaining elements rather than choosing e∗ randomly.
            """
            # print(self._elementsSet, ' add child: ', e)
            # print('state unvisited elements', s.get_unvisited_elements())
            if self._elementsSet is None:
                self._elementsSet = ()
            elementsSet = self._elementsSet + e
            n = HotSpot.Node(elements=elementsSet, parent=self, state=s)

            if e in self._unvisited_elements:
                del self._unvisited_elements[e]
            self._child_nodes.append(n)
            return n

        def update(self, result):
            """ update this node. We update the Q of a father only when the child’s Q is greater than the father’s.
            """
            self._visits += 1
            if self._score < result:
                self._score = result

            from_elementsSet = ()
            if self._parent_node != None:
                from_elementsSet = self._parent_node._elementsSet

            if len(self._elementsSet) == 0:
                return
            to_elementsSet = self._elementsSet
            self._state.update_edge_visits(from_elementsSet, to_elementsSet)

        # get the element with max potential score.
        def get_unvisited_ele_with_max_ps(self):
            return sorted(self._unvisited_elements.items(), key=lambda x: x[1])[-1]

        def __repr__(self):
            return "[elementsSet:" + str(self._elementsSet) + " Score/Visits:" + str(self._score) \
                   + "/" + str(self._visits) + " Unvisited:" + str(self._unvisited_elements) + "]"

        def TreeToString(self, indent):
            s = self.IndentString(indent) + str(self)
            for c in self._child_nodes:
                s += c.TreeToString(indent + 1)
            return s

        def IndentString(self, indent):
            s = "\n"
            for i in range(1, indent + 1):
                s += "| "
            return s

        def ChildrenToString(self):
            s = ""
            for c in self._child_nodes:
                s += str(c) + "\n"
            return s

    def MCTS(self, root_state, itermax, verbose=False):
        """ Conduct a MCTS search for itermax iterations starting from root_state.
            Return the best set from the root_state.
        """

        retSet = {}
        rootnode = self.Node(state=root_state)
        bestPS = 0
        bestSet = None
        for i in range(itermax):
            node = rootnode
            state = root_state.Clone()

            ps = node._score
            print('ps origin', ps)

            # Select
            while node._child_nodes != []:
                node = node.select_child()
                for element in node._elementsSet:
                    state.visit_element((element,))
            print('Select Node', node)

            # Expand
            if len(node._unvisited_elements) != 0:  # if we can expand (i.e. state/node is non-terminal)
                element_choose = node.get_unvisited_ele_with_max_ps()[0]
                state.visit_element(element_choose)
                node = node.add_child(element_choose, state)  # add child and descend tree
            print('Expand Node', node)

            # Evaluation
            elementsSet = node._elementsSet
            if elementsSet in state._elements:
                ps = state._elements[elementsSet]
            else:
                _, value = self._KPIPoint.get_descendant_elements_coms(elementsSet)
                print('eval ', _, value)
                ps = self.cal_potential_scores(elementsSet, value)
            print('Evaluation', ps)

            if ps > bestPS:
                bestPS = ps
                bestSet = elementsSet

            # Early stop
            if ps >= self._ps_threshold:
                print('Early stop')
                elementsSet = node._elementsSet
                return elementsSet, ps

            # Backpropagate
            while node != None:  # backpropagate from the expanded node and work back to the root node
                node.update(ps)  # state is terminal. update node with potential score
                node = node._parent_node
            print('Backpropagate')

            # Output some information about the tree - can be omitted
            if (verbose):
                print(rootnode.TreeToString(0))
            else:
                print(rootnode.ChildrenToString())

            if len(state.get_unvisited_elements()) == 0:
                print('all elements visited')
                break

            print('\n#%d edge visited' % i, state._edge_visited, node, '\n')

        # retNode = sorted(rootnode._child_nodes, key=lambda c: c._score)[-1]
        # ps = retNode._score
        return bestSet, bestPS

if __name__ == "__main__":
    attr_map = {'a': ['a1', 'a2'], 'b': ['b1', 'b2', 'b3']}
    kPoint = KPIPoint(attr_map,
                      1000, {('a1', 'b1'): [14, 20],
                             ('a1', 'b2'): [9, 15],
                             ('a1', 'b3'): [10, 10],
                             ('a2', 'b1'): [7, 10],
                             ('a2', 'b2'): [15, 25],
                             ('a2', 'b3'): [20, 20],
                             })
    hotSpot = HotSpot(kPoint, 2, 0.9, 100)
    rootCauseSet = hotSpot.find_root_cause_set()
    print(rootCauseSet)
    _, value = kPoint.get_descendant_elements_coms((('b2',), ('b1',),))
    print(value)
    score = hotSpot.cal_potential_scores((('b2',), ('b1',),), value)
    print(score)


"""
=============================
Plotting the Optic Radiations
=============================

pyAFQ is designed to be customizable. This example shows how
you can customize it to define a new bundle based
on both waypoint ROIs of your design, as well as endpoint
ROIs of your design.

For now, this is a hypothetical example, as we do not yet
provide these ROIs as part of the software.
"""
from dipy.data import get_fnames
import os.path as op
import matplotlib.pyplot as plt
import numpy as np
import nibabel as nib
import dipy.data as dpd
from dipy.data import fetcher
from dipy.io.streamline import save_tractogram, load_tractogram
from dipy.stats.analysis import afq_profile, gaussian_weights
from dipy.io.stateful_tractogram import StatefulTractogram
from dipy.io.stateful_tractogram import Space
from dipy.reconst import shm

import AFQ.data as afd
import AFQ.tractography as aft
import AFQ.registration as reg
import AFQ.models.dti as dti
import AFQ.models.csd as csd
import AFQ.segmentation as seg
from AFQ.utils.volume import transform_inverse_roi

import logging
logging.basicConfig(level=logging.INFO)

# Target directory for this example's output files
working_dir = "./optic_radiations"

##########################################################################
# Get example data:
# -------------------------

dpd.fetch_stanford_hardi()
hardi_dir = op.join(fetcher.dipy_home, "stanford_hardi")
hardi_fdata = op.join(hardi_dir, "HARDI150.nii.gz")
hardi_fbval = op.join(hardi_dir, "HARDI150.bval")
hardi_fbvec = op.join(hardi_dir, "HARDI150.bvec")
img = nib.load(hardi_fdata)

##########################################################################
# Calculate DTI:
# -------------------------

print("Calculating DTI...")
if not op.exists(op.join(working_dir, 'dti_FA.nii.gz')):
    dti_params = dti.fit_dti(hardi_fdata, hardi_fbval, hardi_fbvec,
                             out_dir=working_dir)
else:
    dti_params = {'FA': op.join(working_dir, 'dti_FA.nii.gz'),
                  'params': op.join(working_dir, 'dti_params.nii.gz')}

FA_img = nib.load(dti_params['FA'])
FA_data = FA_img.get_fdata()

##########################################################################
# Calculate CSD:
# -------------------------
print("Calculating CSD...")
if not op.exists(op.join(working_dir, 'csd_sh_coeff.nii.gz')):
    sh_coeff = csd.fit_csd(hardi_fdata, hardi_fbval, hardi_fbvec,
                           sh_order=4, out_dir=working_dir)
else:
    sh_coeff = op.join(working_dir, "csd_sh_coeff.nii.gz")

apm = shm.anisotropic_power(nib.load(sh_coeff).get_fdata())

##########################################################################
# Register the individual data to a template:
# -------------------------------------------
# For the purpose of bundle segmentation, the individual brain is registered to
# the MNI T1 template. The waypoint ROIs used in segmentation are then each
# brought into each subject's native space to test streamlines for whether they
# fulfill the segmentation criteria.
#
# .. note::
#
#     To find the right place for the waypoint ROIs, we calculate a non-linear
#     transformation between the individual's brain DWI measurement (the b0
#     measurements) and the MNI T1 template.
#     Before calculating this non-linear warping, we perform a pre-alignment
#     using an affine transformation.

print("Registering to template...")

MNI_T1w_img = afd.read_mni_template(weight="T1w")

if not op.exists(op.join(working_dir, 'mapping.nii.gz')):
    import dipy.core.gradients as dpg
    gtab = dpg.gradient_table(hardi_fbval, hardi_fbvec)
    # Prealign using affine registration
    _, prealign = reg.affine_registration(
        apm,
        MNI_T1w_img.get_fdata(),
        img.affine,
        MNI_T1w_img.affine)

    # Then register using a non-linear registration using the affine for
    # prealignment
    warped_hardi, mapping = reg.syn_register_dwi(hardi_fdata, gtab,
                                                 prealign=prealign)
    reg.write_mapping(mapping, op.join(working_dir, 'mapping.nii.gz'))
else:
    mapping = reg.read_mapping(op.join(working_dir, 'mapping.nii.gz'),
                               img, MNI_T1w_img)

##########################################################################
# Bundle specification
# -------------------------------------------
#
# Here, a bundle specification is defined as a series of waypoint ROIs.
# For each hemisphere, two ROIs are inclusion ROIs and three ROIs are
# exclusion ROIs.

roi_folder = op.join(op.expanduser('~'), "AFQ_Data", "visual")
waypoint_folder = op.join(roi_folder, "waypoint")

waypoint_roi_fnames = [
    "left_OR_1.nii.gz",
    "left_OR_2.nii.gz",
    "left_OP_MNI.nii.gz",
    "left_TP_MNI.nii.gz",
    "left_pos_thal_MNI.nii.gz",
    "right_OR_1.nii.gz",
    "right_OR_2.nii.gz",
    "right_pos_thal_MNI.nii.gz",
    "right_OP_MNI.nii.gz",
    "right_TP_MNI.nii.gz"]

waypoint_rois = {}

for fname in waypoint_roi_fnames:
    waypoint_rois[fname.split('.')[0]] = afd.read_resample_roi(
        op.join(waypoint_folder, fname))

bundles = {
    "L_OR": {
        "ROIs": [waypoint_rois["left_OR_1"],
                 waypoint_rois["left_OR_2"],
                 waypoint_rois["left_OP_MNI"],
                 waypoint_rois["left_TP_MNI"],
                 waypoint_rois["left_pos_thal_MNI"]],
        "rules": [True, True, False, False, False],
        "cross_midline": False,
        "uid": 1
        },
    "R_OR": {
        "ROIs": [waypoint_rois["right_OR_1"],
                 waypoint_rois["right_OR_2"],
                 waypoint_rois["right_OP_MNI"],
                 waypoint_rois["right_TP_MNI"],
                 waypoint_rois["right_pos_thal_MNI"]],
        "rules": [True, True, False, False, False],
        "cross_midline": False,
        "uid": 2
        }
    }

##########################################################################
# Endpoints
# ----------
# In addition to the waypoint ROIs, we will customize the endpoint ROIs
# used for filtering the streamlines that are selected based on the
# waypoint ROIs defined above.

endpoint_folder = op.join(roi_folder, "endpoint")

endpoint_spec = {
    "L_OR": {
        "startpoint": nib.load(op.join(endpoint_folder,
                                       'left_thal_MNI.nii.gz')),
        "endpoint": nib.load(op.join(endpoint_folder,
                                     'left_V1_MNI.nii.gz'))},
    "R_OR": {
        "startpoint": nib.load(op.join(endpoint_folder,
                                       'right_thal_MNI.nii.gz')),
        "endpoint": nib.load(op.join(endpoint_folder,
                                     'right_V1_MNI.nii.gz'))}}

##########################################################################
# Tracking
# --------
# We will use PFT and generate a large number of streamlines from seeds
# placed only within the inclusion ROIs. This allows us to oversample that
# part of the brain, without having to deal with very large tractograms.

f_pve_csf, f_pve_gm, f_pve_wm = get_fnames('stanford_pve_maps')

pve_csf = nib.load(f_pve_csf)
pve_gm = nib.load(f_pve_gm)
pve_wm = nib.load(f_pve_wm)

print("Tracking...")
if not op.exists(op.join(working_dir, 'pft_streamlines.trk')):
    seed_roi = np.zeros(img.shape[:-1])
    for bundle in bundles:
        for idx, roi in enumerate(bundles[bundle]['ROIs']):
            warped_roi = transform_inverse_roi(
                roi,
                mapping,
                bundle_name=bundle)
            print(roi)
            nib.save(nib.Nifti1Image(warped_roi.astype(float), img.affine),
                     op.join(working_dir, f"{bundle}_{idx+1}.nii.gz"))

            # Add voxels that aren't there yet:
            if bundles[bundle]['rules'][idx]:
                seed_roi = np.logical_or(seed_roi, warped_roi)

        for ii, pp in enumerate(endpoint_spec[bundle].keys()):
            roi = endpoint_spec[bundle][pp]
            roi = reg.resample(roi.get_fdata(),
                               MNI_T1w_img,
                               roi.affine,
                               MNI_T1w_img.affine)

            warped_roi = transform_inverse_roi(
                roi,
                mapping,
                bundle_name=bundle)

            nib.save(nib.Nifti1Image(warped_roi.astype(float), img.affine),
                     op.join(working_dir, f"{bundle}_{pp}.nii.gz"))

    nib.save(nib.Nifti1Image(seed_roi.astype(float), img.affine),
             op.join(working_dir, 'seed_roi.nii.gz'))

    sft = aft.track(sh_coeff,
                    seed_mask=seed_roi,
                    n_seeds=5,
                    tracker="pft",
                    stop_mask=(pve_wm, pve_gm, pve_csf),
                    stop_threshold="ACT",
                    directions="prob",
                    odf_model="CSD")

    save_tractogram(sft, op.join(working_dir, 'pft_streamlines.trk'),
                    bbox_valid_check=False)
else:
    sft = load_tractogram(op.join(working_dir, 'pft_streamlines.trk'), img)

sft.to_vox()

##########################################################################
# Segmentation
# ------------
# We run the segmentation using both the ``bundles`` and the
# ``endpoint_spec`` we defined above. In this particular case, we set a
# rather lenient criterion for endpoint filtering, by changing from the
# default value of ``dist_to_atlas`` (4 mm) to 5 mm.
#
print("Segmenting fiber groups...")
segmentation = seg.Segmentation(return_idx=True,
                                dist_to_atlas=5)
segmentation.segment(bundles,
                     sft,
                     fdata=hardi_fdata,
                     fbval=hardi_fbval,
                     fbvec=hardi_fbvec,
                     mapping=mapping,
                     reg_template=MNI_T1w_img,
                     endpoint_dict=endpoint_spec)

fiber_groups = segmentation.fiber_groups

##########################################################################
# Cleaning
# --------
# We proceed to clean outliers and save out trk files with the bundles.

print("Cleaning fiber groups...")
for bundle in bundles:
    print(f"Cleaning {bundle}")
    print(f"Before cleaning: {len(fiber_groups[bundle]['sl'])} streamlines")
    new_fibers, idx_in_bundle = seg.clean_bundle(
        fiber_groups[bundle]['sl'],
        return_idx=True)
    print(f"Afer cleaning: {len(new_fibers)} streamlines")
    new_fibers = fiber_groups[bundle]['sl']
    idx_in_global = fiber_groups[bundle]['idx'][idx_in_bundle]
    np.save(op.join(working_dir, f'{bundle}_idx.npy'), idx_in_global)
    sft = StatefulTractogram(new_fibers.streamlines,
                             img,
                             Space.VOX)
    sft.to_rasmm()
    save_tractogram(sft, op.join(working_dir, f'{bundle}_afq.trk'),
                    bbox_valid_check=False)

##########################################################################
# Bundle profiles
# ---------------
# Finally, we can extract and plot bundle profiles.

print("Extracting tract profiles...")
for bundle in bundles:
    sft = load_tractogram(op.join(working_dir, f'{bundle}_afq.trk'),
                          img, to_space=Space.VOX)
    fig, ax = plt.subplots(1)
    weights = gaussian_weights(sft.streamlines)
    profile = afq_profile(FA_data, sft.streamlines,
                          np.eye(4), weights=weights)
    ax.plot(profile)
    ax.set_title(bundle)

plt.show()

##########################################################################
# References:
# -------------------------
# .. [Yeatman2012] Jason D Yeatman, Robert F Dougherty, Nathaniel J Myall,
#                  Brian A Wandell, Heidi M Feldman, "Tract profiles of
#                  white matter properties: automating fiber-tract
#                  quantification", PloS One, 7: e49790
#
# .. [Yeatman2014] Jason D Yeatman, Brian A Wandell, Aviv Mezer Feldman,
#                  "Lifespan maturation and degeneration of human brain white
#                  matter", Nature Communications 5: 4932


import numpy as np
import copy
from tqdm import tqdm

from gluoncv.torch.data.gluoncv_motion_dataset.dataset import DataSample

from .eval_utils import greedy_matching, compute_AP, bbs_iou

def eval_det_ap(gt: list, pred: dict, class_table=None, data_filter_fn=None, iou_threshold=[0.5]):
    """
    Evaluate the detection performance (COCO-style ap) on PoseTrack dataset
    :param gt: ground truth annotations for all videos
    :type gt: dict(vid_id: DataSample)
    :param pred: predictions for all videos
    :type pred: dict(vid_id: DataSample)
    :param data_filter_fn: a callable function that filters out detections that are not considered during evaluation
    :param class_table: class table specify the class order
    :param iou_threshold:
    :return: Average Precision (AP) over different thresholds
    """
    if class_table is None:
        class_table = ["person"]
    num_classes = len(class_table)

    all_scores = [[[] for _ in range(len(iou_threshold))] for _ in range(num_classes)]
    all_pr_ious = [[[] for _ in range(len(iou_threshold))] for _ in range(num_classes)]
    all_gt_ious = [[[] for _ in range(len(iou_threshold))] for _ in range(num_classes)]

    for (vid_id, vid_gt) in tqdm(gt):
        vid_pred = pred[vid_id]

        eval_frame_idxs = vid_gt.get_non_empty_frames()

        # Loop over all classes
        for class_id in range(0, num_classes):
            gt_class_entities = vid_gt.entities
            # gt_class_entities = vid_gt.get_entities_with_label(class_table[class_id])
            pred_class_entities = vid_pred.get_entities_with_label(class_table[class_id])

            # Wrap entities to a DataSample
            vid_class_gt = DataSample(vid_id, metadata=vid_gt.metadata)
            vid_class_pred = DataSample(vid_id, metadata=vid_pred.metadata)
            for _entity in gt_class_entities:
                vid_class_gt.add_entity(_entity)
            for _entity in pred_class_entities:
                vid_class_pred.add_entity(_entity)

            # Get AP for this class and video
            vid_class_scores, vid_class_pr_ious, vid_class_gt_ious = \
                get_ap(vid_class_gt, vid_class_pred, data_filter_fn, eval_frame_idxs, iou_threshold)

            for iou_id in range(len(iou_threshold)):
                all_scores[class_id][iou_id] += vid_class_scores[iou_id]
                all_pr_ious[class_id][iou_id] += vid_class_pr_ious[iou_id]
                all_gt_ious[class_id][iou_id] += vid_class_gt_ious[iou_id]

    class_ap_matrix = np.zeros((num_classes, len(iou_threshold)))
    for class_id in range(num_classes):
        class_ap_matrix[class_id, :] = compute_AP(all_scores[class_id],
                                                  all_pr_ious[class_id],
                                                  all_gt_ious[class_id])

    return class_ap_matrix

def get_ap(vid_class_gt: DataSample, vid_class_pred: DataSample, filter_fn, eval_frame_idxs, iou_thresh=[0.5]):
    """
    :param vid_class_gt: the ground truths for a specific class, in DataSample format
    :param vid_class_pred: the predictions for a specific class, in DataSample format
    :param filter_fn: a callable function to filter out detections
    :param eval_frame_idxs: the frame indexs where evaluation happens
    :param iou_thresh: the list of iou threshod that determines whether a detection is TP
    :returns
           vid_scores: the confidence for every predicted entity (a Python list)
           vid_pr_ious: the iou between the predicted entity and its matching gt entity (a Python list)
           vid_gt_ious: the iou between the gt entity and its matching predicted entity (a Python list)
    """
    if not isinstance(iou_thresh, list):
        iou_thresh = [iou_thresh]
    vid_scores = [[] for _ in iou_thresh]
    vid_pr_ious = [[] for _ in iou_thresh]
    vid_gt_ious = [[] for _ in iou_thresh]
    for frame_idx in eval_frame_idxs:

        gt_entities = vid_class_gt.get_entities_for_frame_num(frame_idx)
        pred_entities = vid_class_pred.get_entities_for_frame_num(frame_idx)

        # Remove detections for evaluation that are within ignore regions
        if filter_fn is not None:
            # Filter out ignored gt entities
            gt_entities, ignore_gt_entities = filter_fn(gt_entities, meta_data=vid_class_gt.metadata)
            # Filter out predicted entities that overlaps with ignored gt entities
            pred_entities, ignore_pred_entities = filter_fn(pred_entities, ignore_gt_entities)

        # sort the entity based on confidence scores
        pred_entities = sorted(pred_entities, key=lambda x: x.confidence, reverse=True)
        iou_matrix = bbs_iou(pred_entities, gt_entities)
        scores = [entity.confidence for entity in pred_entities]
        for i, _iou in enumerate(iou_thresh):
            # pred_ious, gt_ious = target_matching(pred_entities, gt_entities)
            pred_ious, gt_ious = greedy_matching(copy.deepcopy(iou_matrix), _iou)
            vid_scores[i] += scores
            vid_pr_ious[i] += pred_ious
            vid_gt_ious[i] += gt_ious

    return vid_scores, vid_pr_ious, vid_gt_ious


from sklearn.externals import joblib
from cassiopeia import Match
from cassiopeia.core.match import Timeline
import numpy as np
import matplotlib.path as mpl_path
import pandas as pd
import os

# Initializing all the roles accepted
role_composition = {"JUNGLE_NONE", "TOP_SOLO", "MIDDLE_SOLO", "BOTTOM_DUO_CARRY", "BOTTOM_DUO_SUPPORT"}
clean_roles = {"JUNGLE_NONE": 'jungle', "TOP_SOLO": 'top', "MIDDLE_SOLO": 'mid',
               "BOTTOM_DUO_CARRY": 'bot', "BOTTOM_DUO_SUPPORT": 'supp'}

# Init spells
spells = {'spell-21': 0, 'spell-1': 0, 'spell-14': 0, 'spell-3': 0, 'spell-4': 0, 'spell-6': 0, 'spell-7': 0,
          'spell-11': 0, 'spell-12': 0}

# Static data for lane position
midlane = mpl_path.Path(
    np.array([[4200, 3500], [11300, 10500], [13200, 13200], [10500, 11300], [3300, 4400], [1600, 1600]]))
toplane = mpl_path.Path(
    np.array([[-120, 1600], [-120, 14980], [13200, 14980], [13200, 13200], [4000, 13200], [1600, 11000], [1600, 1600]]))
botlane = mpl_path.Path(
    np.array([[1600, -120], [14870, -120], [14870, 13200], [13200, 13200], [13270, 4000], [10500, 1700], [1600, 1600]]))
jungle1 = mpl_path.Path(
    np.array([[1600, 5000], [1600, 11000], [4000, 13200], [9800, 13200], [10500, 11300], [3300, 4400]]))
jungle2 = mpl_path.Path(
    np.array([[5000, 1700], [4200, 3500], [11300, 10500], [13270, 9900], [13270, 4000], [10500, 1700]]))

# Static data for item frequency
overUsedItems = {
    'BOTTOM_DUO_SUPPORT': ['3050', '3069', '3092', '3096', '3097', '3098', '3105', '3107', '3114', '3222', '3382', '3401', '3504'],
    'JUNGLE_NONE': ['1039', '1041', '1400', '1401', '1402', '1412', '1413', '1414', '1416', '1419', '2057', '3706', '3715'],
    'TOP_SOLO': ['3068', '3161', '3196', '3373', '3379'],
    'BOTTOM_DUO_CARRY': ['2319', '3004', '3042', '3095', '3389'],
    'MIDDLE_SOLO': ['3197']
}
mostlyUsedItems = {
    'BOTTOM_DUO_SUPPORT': ['1004', '2065', '2403', '3024', '3028', '3109', '3117', '3174', '3190', '3383', '3801'],
    'JUNGLE_NONE': ['3083', '3380', '3513'],
    'TOP_SOLO': ['2053', '3022', '3056', '3074', '3387', '3512', '3748', '3751', '3812'],
    'BOTTOM_DUO_CARRY': ['1042', '1043', '1051', '1055', '1083', '2011', '2015', '2060', '2061', '3006', '3025', '3031', '3046', '3072', '3085', '3086', '3087', '3094', '3124', '3139', '3140', '3144', '3153', '3363', '3371', '3508'],
    'MIDDLE_SOLO': ['1056', '3003', '3027', '3030', '3040', '3198', '3285', '3390']
}

underUsedItems = {
    'BOTTOM_DUO_SUPPORT': ['1011', '1018', '1036', '1037', '1038', '1039', '1041', '1042', '1043', '1051', '1053', '1054', '1055', '1056', '1057', '1058', '1083', '1400', '1401', '1402', '1412', '1413', '1414', '1416', '1419', '2011', '2015', '2033', '2057', '2058', '2059', '2061', '2062', '2140', '2319', '2420', '2421', '3001', '3004', '3006', '3020', '3022', '3025', '3026', '3027', '3031', '3033', '3035', '3036', '3040', '3042', '3044', '3046', '3047', '3052', '3053', '3057', '3065', '3068', '3071', '3072', '3074', '3075', '3076', '3077', '3078', '3083', '3085', '3086', '3087', '3089', '3091', '3094', '3095', '3100', '3101', '3102', '3115', '3123', '3124', '3135', '3139', '3140', '3142', '3143', '3144', '3146', '3152', '3153', '3155', '3156', '3161', '3165', '3193', '3194', '3196', '3197', '3198', '3211', '3285', '3340', '3363', '3371', '3373', '3374', '3379', '3380', '3384', '3386', '3387', '3389', '3508', '3513', '3706', '3715', '3742', '3748', '3751', '3812', '3814', '3907', '3916'],
    'JUNGLE_NONE': ['1004', '1006', '1018', '1051', '1053', '1054', '1055', '1056', '1083', '2003', '2004', '2010', '2011', '2013', '2015', '2033', '2053', '2056', '2058', '2059', '2060', '2061', '2062', '2065', '2319', '2403', '2422', '2423', '2424', '3003', '3004', '3006', '3009', '3024', '3025', '3027', '3028', '3030', '3031', '3040', '3042', '3046', '3050', '3056', '3068', '3069', '3070', '3072', '3086', '3087', '3092', '3094', '3095', '3096', '3097', '3098', '3105', '3107', '3114', '3115', '3116', '3123', '3124', '3139', '3140', '3146', '3152', '3153', '3158', '3161', '3190', '3196', '3197', '3198', '3222', '3285', '3301', '3302', '3363', '3371', '3373', '3379', '3382', '3389', '3390', '3401', '3504', '3508', '3751', '3801', '3802', '3812', '3905', '3907'],
    'TOP_SOLO': ['1004', '1039', '1041', '1042', '1051', '1400', '1401', '1402', '1412', '1413', '1414', '1416', '1419', '2003', '2015', '2055', '2057', '2065', '2319', '2423', '3004', '3006', '3009', '3028', '3031', '3041', '3042', '3046', '3050', '3069', '3072', '3085', '3086', '3092', '3094', '3095', '3096', '3097', '3098', '3100', '3105', '3107', '3109', '3113', '3114', '3117', '3124', '3139', '3145', '3153', '3158', '3174', '3190', '3197', '3222', '3285', '3303', '3363', '3364', '3371', '3374', '3380', '3382', '3386', '3388', '3389', '3390', '3401', '3504', '3513', '3706', '3715', '3801', '3905'],
    'BOTTOM_DUO_CARRY': ['1001', '1004', '1006', '1011', '1026', '1027', '1028', '1029', '1031', '1039', '1041', '1052', '1056', '1058', '1082', '1400', '1401', '1402', '1412', '1413', '1414', '1416', '1419', '2033', '2053', '2055', '2057', '2065', '2138', '2139', '2403', '3001', '3003', '3010', '3020', '3024', '3027', '3028', '3030', '3040', '3041', '3044', '3047', '3050', '3052', '3053', '3056', '3057', '3065', '3067', '3068', '3069', '3070', '3071', '3074', '3075', '3076', '3077', '3078', '3082', '3083', '3089', '3092', '3096', '3097', '3098', '3100', '3102', '3105', '3107', '3108', '3109', '3110', '3111', '3113', '3114', '3116', '3117', '3134', '3135', '3136', '3142', '3143', '3147', '3151', '3152', '3157', '3161', '3165', '3174', '3190', '3191', '3193', '3194', '3196', '3197', '3198', '3211', '3222', '3285', '3303', '3340', '3364', '3373', '3374', '3379', '3380', '3382', '3383', '3387', '3390', '3401', '3504', '3512', '3513', '3706', '3715', '3742', '3748', '3751', '3800', '3801', '3812', '3814', '3905', '3907', '3916'],
    'MIDDLE_SOLO': ['1004', '1006', '1011', '1029', '1031', '1033', '1039', '1041', '1042', '1043', '1051', '1055', '1057', '1083', '1400', '1401', '1402', '1412', '1413', '1414', '1416', '1419', '2010', '2011', '2013', '2015', '2053', '2056', '2057', '2058', '2060', '2061', '2062', '2065', '2138', '2140', '2319', '3004', '3006', '3009', '3010', '3022', '3024', '3025', '3028', '3042', '3046', '3047', '3050', '3056', '3065', '3067', '3068', '3069', '3071', '3072', '3074', '3075', '3076', '3082', '3083', '3085', '3086', '3092', '3094', '3095', '3096', '3097', '3098', '3101', '3105', '3107', '3109', '3110', '3114', '3117', '3123', '3124', '3139', '3143', '3144', '3153', '3161', '3174', '3190', '3193', '3194', '3196', '3211', '3222', '3301', '3364', '3373', '3379', '3382', '3383', '3387', '3389', '3401', '3504', '3508', '3512', '3513', '3706', '3715', '3742', '3751', '3800', '3801']
}

# TODO : these should not be static
#       Need to set up an external resource to call for item frequencies for the current patch.
#       Even better would be for each patch.

# Loading the model
roleml_model = joblib.load(os.path.join(os.path.dirname(__file__), "role_identification_model.sav"))

def get_positions(timeline):
    frames = timeline['frames']

    participants_positions = {1: [], 2: [], 3: [], 4: [], 5: [], 6: [], 7: [], 8: [], 9: [], 10: []}
    # 10 first frames except the really first one, when everybody spawn
    for i in range(1, 11):
        # For each participant
        for k in frames[i]['participantFrames']:
            participantId = frames[i]['participantFrames'][k]["participantId"]
            position = None

            # Position on the map
            coord = [frames[i]['participantFrames'][k]['position']['x'],
                     frames[i]['participantFrames'][k]['position']['y']]
            # Check where is the position
            if jungle1.contains_point(coord) or jungle2.contains_point(coord):
                position = "jungle"
            elif midlane.contains_point(coord):
                position = "mid"
            elif toplane.contains_point(coord):
                position = "top"
            elif botlane.contains_point(coord):
                position = "bot"
            # Save the position for the participant
            participants_positions[int(participantId)].append(position)

    return participants_positions

def get_most_frequent_lane(participants_positions):
    most_frequent_lane = {}
    for participant_id in participants_positions:
        lane_frequency = {"mid": 0, "top": 0, "bot": 0, "jungle": 0}

        for lane in participants_positions[participant_id]:
            if lane is not None:
                lane_frequency[lane] += 1

        most_frequent_lane[participant_id] = max(lane_frequency, key=lane_frequency.get)

    return most_frequent_lane

def get_lane_frequencies(participants_positions):
    lane_frequencies = {}
    for participant_id in participants_positions:
        lane_frequency = {"mid": 0, "top": 0, "bot": 0, "jungle": 0}

        for lane in participants_positions[participant_id]:
            if lane is not None:
                lane_frequency[lane] += 1

        lane_frequencies[participant_id] = lane_frequency

    return lane_frequencies

def get_stats_at_10(timeline):
    participant_frame = timeline['frames'][10]['participantFrames']
    stats_at_10 = {}
    for k in participant_frame:
        participantId = participant_frame[k]["participantId"]
        row = {
            "minionsKilled": participant_frame[k]["minionsKilled"],
            "jungleMinionsKilled": participant_frame[k]["jungleMinionsKilled"],
            "jungleMinionRatio": participant_frame[k]["jungleMinionsKilled"] /
                                 (participant_frame[k]["minionsKilled"] + participant_frame[k]["jungleMinionsKilled"])
            if participant_frame[k]["minionsKilled"] > 0 else 0
        }
        stats_at_10[participantId] = row

    return stats_at_10

def get_features(match, timeline):

    participants_features_list = []

    # Get the players positions form the timeline
    players_positions = get_positions(timeline)

    player_most_frequent_lane = get_most_frequent_lane(players_positions)
    player_lane_frequencies = get_lane_frequencies(players_positions)

    player_stats = get_stats_at_10(timeline)

    for participant in match['participants']:

        participant_features = {}

        participant_id = participant["participantId"]

        lanes = ["jungle", "top", "mid", "bot"]

        # one hot encode positions
        for lane in lanes:
            participant_features['most-frequent-' + lane] = 0

        participant_features['most-frequent-' + player_most_frequent_lane[participant_id]] = 1

        # Lane frequency
        participant_features.update({"lane-frequency-" + k: v for k, v in
                                     player_lane_frequencies[participant_id].items()})

        # Init items lists
        for role in role_composition:
            participant_features["has-item-overUsed-" + role] = 0
            participant_features["has-item-mostlyUsed-" + role] = 0
            participant_features["has-item-underUsed-" + role] = 0

        # Get the item ID for each of the 7 slots and check if they are in one of the three items lists
        for i in range(0, 7):
            # Check if there is an item for the slot
            if participant['stats']['item' + str(i)] > 0:
                for role in role_composition:

                    if str(participant['stats']['item' + str(i)]) in overUsedItems[role]:
                        participant_features["has-item-overUsed-" + role] = 1

                    if str(participant['stats']['item' + str(i)]) in mostlyUsedItems[role]:
                        participant_features["has-item-mostlyUsed-" + role] = 1

                    if str(participant['stats']['item' + str(i)]) in underUsedItems[role]:
                        participant_features["has-item-underUsed-" + role] = 1

        # Summoner spells
        participant_features.update(spells)
        participant_features["spell-" + str(participant["spell1Id"])] = 1
        participant_features["spell-" + str(participant["spell2Id"])] = 1

        # Player stats
        participant_features.update(player_stats[participant_id])

        participant_features["participantId"] = participant_id

        participants_features_list.append(participant_features)

    return pd.DataFrame(participants_features_list)

def __prepare_cass_match(match):
    if not match._Ghost__all_loaded:
        match.load()

    match_dict = match.to_dict()
    for participant in match_dict['participants']:
        participant['participantId'] = participant['id']
        participant["spell1Id"] = participant["summonerSpellDId"]
        participant["spell2Id"] = participant["summonerSpellFId"]
        del participant['id']
        del participant['summonerSpellDId']
        del participant['summonerSpellFId']
    return match_dict

def __prepare_cass_timeline(timeline):
    if not timeline._Ghost__all_loaded:
        timeline.load()

    timeline_dict = timeline.to_dict()
    minute_10_participant_frames = timeline_dict['frames'][10]['participantFrames']
    for k in minute_10_participant_frames:
        minute_10_participant_frames[k]['minionsKilled'] = minute_10_participant_frames[k]['creepScore']
        minute_10_participant_frames[k]["jungleMinionsKilled"] = minute_10_participant_frames[k]["neutralMinionsKilled"]
        minute_10_participant_frames[k]["jungleMinionRatio"] = minute_10_participant_frames[k]["neutralMinionsKilled"] / (minute_10_participant_frames[k]["creepScore"] + minute_10_participant_frames[k]["neutralMinionsKilled"])
        del minute_10_participant_frames[k]['creepScore']
        del minute_10_participant_frames[k]['neutralMinionsKilled']
    return timeline_dict

def predict(match, timeline):

    if isinstance(match, Match):
        match = __prepare_cass_match(match)

    if isinstance(timeline, Timeline):
        timeline = __prepare_cass_timeline(timeline)

    if match["gameDuration"] < 720:
        raise Exception("Match too short")

    df = get_features(match, timeline)
    df["role"] = roleml_model.predict(df.drop(["participantId"], axis=1))

    df = df.set_index(df["participantId"])

    participant_roles = df["role"].to_dict()

    return {k: clean_roles[participant_roles[k]] for k in participant_roles}

def fix_and_augment_game_and_timeline(game, timeline, upgrade_participant=False, upgrade_timeline=False):
    true_roles = predict(game, timeline)

    # First, set all roles properly so we know player's opponents
    for participant in game['participants']:
        participant['role'] = true_roles[participant['participantId']]

    # Second, add the diffs in the frames. We're not using the "deltas" since it's redundant info
    for participant in game['participants']:
        participant_id = participant['participantId']

        opponent_found = None
        for opponent in game['participants']:
            opponent_id = opponent['participantId']
            if opponent_id != participant_id and opponent['role'] == participant['role']:
                opponent_found = True
                break

        if opponent_found is None:
            raise Exception('Player without opponent found')

        if not game['participants'][opponent_id-1]['participantId'] == opponent_id:
            raise Exception('Opponent array detection screwed')

        participant_timeline = participant['timeline']
        opponent_timeline = game['participants'][opponent_id-1]['timeline']

        # Reseting or creating stats per min fields
        participant_timeline['csDiffPerMinDeltas'] = {}
        for i in participant_timeline['creepsPerMinDeltas']:
            participant_timeline['csDiffPerMinDeltas'][i] = round( participant_timeline['creepsPerMinDeltas'][i] - opponent_timeline['creepsPerMinDeltas'][i], 2)

        participant_timeline['xpDiffPerMinDeltas'] = {}
        for i in participant_timeline['xpPerMinDeltas']:
            participant_timeline['xpDiffPerMinDeltas'][i] = round( participant_timeline['xpPerMinDeltas'][i] - opponent_timeline['xpPerMinDeltas'][i], 2)

        participant_timeline['damageTakenDiffPerMinDeltas'] = {}
        for i in participant_timeline['damageTakenPerMinDeltas']:
            participant_timeline['damageTakenDiffPerMinDeltas'][i] = round( participant_timeline['damageTakenPerMinDeltas'][i] - opponent_timeline['damageTakenPerMinDeltas'][i], 2)

        if upgrade_participant:
            participant_timeline['goldDiffPerMinDeltas'] = {}
            for i in participant_timeline['goldPerMinDeltas']:
                participant_timeline['goldDiffPerMinDeltas'][i] = round( participant_timeline['goldPerMinDeltas'][i] - opponent_timeline['goldPerMinDeltas'][i], 2)

        if upgrade_timeline:
            for frame in timeline['frames']:
                participant_frame = frame['participantFrames'][str(participant_id)]
                opponent_frame = frame['participantFrames'][str(opponent_id)]

                participant_frame['totalGoldDiff'] = participant_frame['totalGold'] - opponent_frame['totalGold']
                participant_frame['xpDiff'] = participant_frame['xp'] - opponent_frame['xp']
                participant_frame['minionsKilledDiff'] = \
                    participant_frame['minionsKilled'] - opponent_frame['minionsKilled']
                participant_frame['jungleMinionsKilledDiff'] = \
                    participant_frame['jungleMinionsKilled'] - opponent_frame['jungleMinionsKilled']

    if not upgrade_participant:
        for participant in game['participants']:
            del(participant['role'])

    return game, timeline