#!/usr/bin/env python
# coding: utf-8

# In[25]:


get_ipython().run_line_magic('load_ext', 'autoreload')
get_ipython().run_line_magic('autoreload', '2')
get_ipython().run_line_magic('matplotlib', 'inline')
import matplotlib.pyplot as plt
import MESS
import numpy as np
import pandas as pd
import random

from IPython.display import display

pd.set_option('display.max_rows', 500)
pd.set_option('display.max_columns', 500)
pd.set_option('display.width', 1000)


# In[2]:


r = MESS.Region("watdo")

r.set_param("J", 250)
r.set_param("speciation_model", "none")
r.set_param("generations", 0.25)
r.run(sims=1)


# ## Competition Dev Code

# In[57]:


loc1 = r.islands.values()[0]
mean_local_trait = loc1.region.get_trait_mean(loc1.local_community)
print(mean_local_trait)
local_traits = map(loc1.region.get_trait, loc1.local_community)

death_probs = map(lambda x: np.exp(-((x - mean_local_trait) ** 2)\
                                        /loc1.region.metacommunity.paramsdict["ecological_strength"]),\
                                        local_traits)
df = pd.DataFrame(local_traits)
df = pd.concat([df, pd.DataFrame(death_probs)], axis=1)

death_probs = np.array(death_probs)/np.sum(death_probs)
df = pd.concat([df, pd.DataFrame(death_probs)], axis=1)
df.columns = ["trait", "death_probs", "proba"]
display(df)
vic_idx = list(np.random.multinomial(1, death_probs)).index(1)
victim = loc1.local_community[vic_idx]
print(victim)


# ## Filtering Dev Code

# In[55]:


n = 20
loc1 = r.islands.values()[0]
print(loc1.region.metacommunity.paramsdict["ecological_strength"],\
      loc1.region.metacommunity._hackersonly["filtering_optimum"])
local_traits = map(loc1.region.get_trait, loc1.local_community)
print(local_traits[:n])
fitness = map(lambda x: 1 - (np.exp(-((x - loc1.region.metacommunity._hackersonly["filtering_optimum"]) ** 2)\
                        /loc1.region.metacommunity.paramsdict["ecological_strength"])),\
                        local_traits)
## Scale all fitness values to proportions
df = pd.DataFrame(local_traits)
df = pd.concat([df, pd.DataFrame(fitness)], axis=1)
fitness = np.array(fitness)/np.sum(fitness)
df = pd.concat([df, pd.DataFrame(fitness)], axis=1)
df.columns = ["trait", "fitness", "proba"]
display(df)
## Get the victim conditioning on unequal death probability
vic_idx = list(np.random.multinomial(1, fitness)).index(1)
victim = loc1.local_community[vic_idx]


# In[7]:


r.set_param("community_assembly_model", "filtering")
r.set_param("generations", 0.25)
r.run(sims=1)


# ## Assembly Model Horserace

# In[51]:


r = MESS.Region("watdo")
r.set_param("J", 250)
r.set_param("community_assembly_model", "neutral")
r.set_param("generations", 10)
r.run(sims=1)


# In[31]:


r.set_param("community_assembly_model", "neutral")
get_ipython().run_line_magic('timeit', 'r.simulate(nsteps=10)')


# In[32]:


r.set_param("community_assembly_model", "filtering")
get_ipython().run_line_magic('timeit', 'r.simulate(nsteps=10)')


# In[33]:


r.set_param("community_assembly_model", "competition")
get_ipython().run_line_magic('timeit', 'r.simulate(nsteps=10)')


# ## Trash below here

# In[52]:


isl1 = r.islands.values()[0]
isl1.get_stats().T
for i in range(10):
    isl1.step()
    display(isl1.get_stats().T)


# In[153]:


r.paramsdict
r.set_param("speciation_rate", 5)


# In[28]:


r = MESS.Region("watdo")
r.set_param("J", 1000)
r.set_param("community_assembly_model", "filtering")
r.set_param("speciation_prob", 0.001)
r.set_param("generations", 0.2)
r.set_param("ecological_strength", 0.01)
r.run(sims=1)
isl1 = r.islands.values()[0]
print(isl1.current_time * 2/isl1.paramsdict["J"])


# In[23]:


get_ipython().run_line_magic('load_ext', 'line_profiler')
r.set_param("community_assembly_model", "filtering")

get_ipython().run_line_magic('lprun', '-T /tmp/rpt.txt -f isl1._filtering_death_step isl1._filtering_death_step()')
get_ipython().system('cat /tmp/rpt.txt')


# ## Can we improve the call to multinomial?
# Not with np.random.choice

# In[317]:


loc_inds = [x for x in isl1.local_community if x != None]
local_traits = map(isl1.region.get_trait, loc_inds)
fo = isl1.region.metacommunity._hackersonly["filtering_optimum"]
es = isl1.region.metacommunity.paramsdict["ecological_strength"]
death_probs = [1 - (np.exp(-((x - fo) ** 2)/es)) for x in local_traits]
death_probs = np.array(death_probs)/np.sum(death_probs)
get_ipython().run_line_magic('timeit', 'vic_idx = list(np.random.multinomial(1, death_probs)).index(1)')

get_ipython().run_line_magic('timeit', 'np.random.choice(loc_inds, size=1, p=death_probs)')


# ## Can we improve calculating the death_probs?
# Not much

# In[180]:


loc_inds = [x for x in isl1.local_community if x != None]
local_traits = map(isl1.region.get_trait, loc_inds)
fo = isl1.region.metacommunity._hackersonly["filtering_optimum"]
es = isl1.region.metacommunity.paramsdict["ecological_strength"]
get_ipython().run_line_magic('timeit', 'death_probs = map(lambda x: 1 - (np.exp(-((x - fo) ** 2)/es)), local_traits)')
get_ipython().run_line_magic('timeit', '[1 - (np.exp(-((x - fo) ** 2)/es)) for x in local_traits]')


# ## Improve getting trait values

# In[171]:


import random
loc_inds = [x for x in isl1.local_community if x != None]
# 1000000 loops, best of 3: 1.37 µs per loop
get_ipython().run_line_magic('timeit', 'victim = random.choice(isl1.local_community)')
# 100 loops, best of 3: 2.68 ms per loop
get_ipython().run_line_magic('timeit', 'local_traits = map(isl1.region.get_trait, loc_inds)')
dat = isl1.region.metacommunity.community

## DataFrame
df = pd.DataFrame(dat, index=dat["ids"])
# 10 loops, best of 3: 136 ms per loop
#%timeit map(lambda x: df.loc[x]["trait_values"], loc_inds)
# 10 loops, best of 3: 113 ms per loop
#%timeit [df.loc[x]["trait_values"] for x in loc_inds]

## Dictionary
trait_dict = {x["ids"]:x["trait_values"] for x in dat}
get_ipython().run_line_magic('timeit', '[isl1.region.get_trait(x) for x in loc_inds]')
get_ipython().run_line_magic('timeit', '[trait_dict[x] for x in loc_inds]')
get_ipython().run_line_magic('timeit', 'map(lambda x:trait_dict[x], loc_inds)')


# ## Actually really trash below here

# In[230]:


isl1.paramsdict


# In[45]:


r = MESS.Region("watdo")
r.set_param("J", 500)
r.set_param("community_assembly_model", "filtering")
r.set_param("speciation_prob", 0.005)
r.set_param("m", 0.01)
r.set_param("generations", 0.5)
r.set_param("ecological_strength", 1)
r.run(sims=1)
isl1 = r.islands.values()[0]
print(isl1.current_time * 2/isl1.paramsdict["J"])


# In[46]:


print(isl1.current_time * 2/isl1.paramsdict["J"])
ff_df = pd.DataFrame([isl1.founder_flags, isl1.local_community])
ff_df.T.groupby([0, 1]).size()


# In[73]:


es = isl1.region.metacommunity.paramsdict["ecological_strength"]
es = 1
def _filtering_death_step(self, es, comp=False):
    victim = random.choice(self.local_community)
    if victim == None:
        pass
    else:
        ## Get local traits for all individuals in the community (remove None first)
        loc_inds = [x for x in self.local_community if x != None]
        local_traits = map(self.region.get_trait, loc_inds)
        if comp:
            mean_local_trait = self.region.get_trait_mean(self.local_community)
            print(mean_local_trait)
            death_probs = map(lambda x: np.exp(-((x - mean_local_trait) ** 2)/(1./(es))),\
                                        local_traits)
        else:
            print(isl1.region.metacommunity._hackersonly["filtering_optimum"])
            death_probs = map(lambda x: 1 - (np.exp(-((x - self.region.metacommunity._hackersonly["filtering_optimum"]) ** 2)\
                                /es)),\
                                local_traits)
        ## Scale all fitness values to proportions
        unscaled_dp = death_probs
        death_probs = np.array(death_probs)/np.sum(death_probs)
        df = pd.DataFrame([loc_inds, local_traits, unscaled_dp, list(death_probs)])
        ## Get the victim conditioning on unequal death probability
        vic_idx = list(np.random.multinomial(1, death_probs)).index(1)
        victim = loc_inds[vic_idx]
    return df
for es in [0.001, 0.01, 0.1, 1, 10, 100, 1000]:
    dp_df = _filtering_death_step(isl1, es, comp=False)
    dp_df = pd.DataFrame(dp_df.T.groupby([0, 1, 2, 3]).size())
    #dp_df.columns=["species", "trait value", "raw death prob", "scaled death prob"]
    display(dp_df)#[0][[0, 1, 2, 3]])


# In[72]:


dp_df[0]


# In[357]:


from collections import Counter

def neut_samp(isl):
    self=isl
    victim = random.choice(self.local_community)
    return victim

def filt_samp(isl, es=0.01):
    self = isl
    loc_inds = [x for x in self.local_community if x != None]
    local_traits = map(self.region.get_trait, loc_inds)
    death_probs = map(lambda x: 1 - (np.exp(-((x - self.region.metacommunity._hackersonly["filtering_optimum"]) ** 2)\
                            /es)),\
                            local_traits)
    ## Scale all fitness values to proportions
    #unscaled_dp = death_probs
    death_probs = np.array(death_probs)/np.sum(death_probs)
    #df = pd.DataFrame([loc_inds, local_traits, unscaled_dp, list(death_probs)])
    ## Get the victim conditioning on unequal death probability
    vic_idx = list(np.random.multinomial(1, death_probs)).index(1)
    victim = loc_inds[vic_idx]
    return victim
print(Counter([neut_samp(isl1) for _ in range(500)]))
print(Counter([filt_samp(isl1, es=0.01) for _ in range(500)]))