import numpy as np
import scipy as sp
from scipy.stats import poisson, expon
import matplotlib.pyplot as plt
import math
import random

random.seed(42)


def hawkes_process_simulation(
    horizon = 100,
    decay_rate = .995,
    intensity_bump_per_event = 1,
    max_intensity = 100,
):
    """
    Simulates event arrival times for a self-exciting point-process,
    i.e. simulating traffic from "viral" attention.

    This is basically a fancy poisson process. What makes it special is that the rate of the process is not fixed.
    Instead, the rate is modeled by an "intensity function". Whenever a new event is observed, the intensity function
    experiences a bump. While no events are happening, the intensity function experiences exponential decay. This manifests
    occasional random spikes of accelerating event rates.

    Algorithm modified from: https://scse.d.umn.edu/sites/scse.d.umn.edu/files/obral_master.pdf

    Parameters:
        horizon: duration of simulation. event times will be generated in the range [0,horizon] with a presumptive event at t=0.
        decay_rate: rate at which the intensity function decays in absence of events
        intensity_bump_per_event: amount by which intensity function increases when a new event is observed
        max_intensity: this can be finicky to parameterize, so I added the ability to threshold the intensity function

    Returns:
        event_history: list of times at which simulated events occur
        excitation_history: simulated intensity function associated with the generated events
    """   
    
    def exp_decay(t, decay_rate=0.5, v0=intensity_bump_per_event):
        return v0 * math.exp(-t * decay_rate)
    
    excitation_history = [1]
    event_history = [0]
    
    while event_history[-1] < horizon:
    
        # sample new event interval
        excitation = excitation_history[-1]
        poisson_rate = 1/min(excitation, max_intensity)
        time_to_candidate_event = expon.rvs(loc=0, scale=poisson_rate, size=1)[0]
    
        #excitation_at_candidate = exp_decay(t=time_to_candidate_event, decay_rate=decay_rate, v0=excitation)
        decayed_excitation = exp_decay(t=time_to_candidate_event, decay_rate=decay_rate, v0=excitation)

        # "thinning" via rejection sampling
        if random.random() < decayed_excitation:
            # accept event
            t_prev = event_history[-1]
            t = t_prev + time_to_candidate_event
            event_history.append(t)
            
            # exp decay self-excitation
            excitation = decayed_excitation + intensity_bump_per_event
            excitation_history.append(excitation)

    return event_history, excitation_history

#######

import seaborn as sns

for i in range(4):
    event_history, excitation_history = hawkes_process_simulation(horizon=100)
    
    #plt.plot(event_history, excitation_history, label=f"n={len(event_history)}")
    sns.kdeplot(event_history, bw_adjust=.2, label=f"n={len(event_history)}")
    #sns.rugplot(event_history)
plt.title("Event Density for Hawkes Process Simulations.\n`n`= # events simulated")
plt.legend()
plt.show()