# cython: boundscheck=False, cdivision=True
import numpy as np
cimport numpy as np
from libc.math cimport log, exp

DTYPE = np.float64
ctypedef np.float_t DTYPE_t


cdef DTYPE_t weight_success(int dist, int dt, DTYPE_t alpha,
                           DTYPE_t delta, DTYPE_t gamma):
    """weight for successful infection, exponential time model"""
    cdef DTYPE_t structural, temporal, result
    structural = delta ** (dist)
    temporal = exp(-alpha * dt) * (1 - exp(-alpha))
    result = log(structural * temporal)
    return result


cdef DTYPE_t weight_success_power(int dist, int dt, DTYPE_t alpha,
                           DTYPE_t delta, DTYPE_t gamma):
    """weight for successful infection, power-law time model"""
    cdef DTYPE_t structural, temporal, result
    structural = delta ** (dist)
    temporal = 1. / (1. + (dt - 1.)/alpha)**0.01 - 1. / (1. + dt/alpha)**0.01
    result = log(structural * temporal)
    return result


cdef DTYPE_t weight_failure(int dist, int dt, DTYPE_t alpha,
                               DTYPE_t delta, DTYPE_t gamma):
    """weight for failed infection, exponential time model"""
    cdef DTYPE_t structural, temporal, result
    structural = delta ** (dist)
    temporal = 1. - exp(-alpha * dt)
    #result = log(1. - structural)
    result = log(1. - structural * temporal)
    return result


cdef DTYPE_t weight_failure_power(int dist, int dt, DTYPE_t alpha,
                               DTYPE_t delta, DTYPE_t gamma):
    """weight for failed infection, power-law time model"""
    cdef DTYPE_t structural, temporal, result
    structural = delta ** (dist)
    temporal = 1.  - 1. / (1. + dt/alpha)**0.01
    result = log(1. - structural * temporal)
    return result

def ml(dict root_victims, dict victims, dict non_victims, DTYPE_t age,
       DTYPE_t alpha, DTYPE_t delta, DTYPE_t gamma=10):
    cdef:
        int n_roots, n_victims, n_nodes, roots, i, dist, dt, t, l
        DTYPE_t beta, all_failures, ll, beta2
        list parents, failures, successes
    n_roots, n_victims = len(root_victims), len(victims)
    n_nodes = n_victims + len(non_victims) + n_roots
    cdef:
        np.ndarray[DTYPE_t] probs = np.zeros(n_victims, dtype=DTYPE)
        np.ndarray[DTYPE_t] probs_fail = np.zeros(n_victims, dtype=DTYPE)
        np.ndarray[DTYPE_t] probs_nv = np.zeros(len(non_victims), dtype=DTYPE)
    for i, parents in enumerate(victims.itervalues()):
        # for each victim node i, compute the probability that all its parents
        # fail to infect it, also computes the probability that its most
        # likely parent infects it
        failures = [weight_failure(dist, dt, alpha, delta, gamma)
                    for (dist, dt) in parents]
        all_failures = sum(failures)
        successes = [weight_success(dist, dt, alpha, delta, gamma)
                     for (dist, dt) in parents]
        probs[i] = max(s - failures[l] for l, s in enumerate(successes))
        probs_fail[i] = all_failures

    for i, parents in enumerate(non_victims.itervalues()):
        # for each non victim node, compute the probability that all its
        # parents fail to infect it
        failures = [weight_failure(dist, dt, alpha, delta, gamma)
                    for (dist, dt) in parents]
        probs_nv[i] = sum(failures)
    probs.sort()
    probs = probs[::-1]
    cdef:
        np.ndarray[DTYPE_t] cums = probs.cumsum()
    ll =  probs_fail.sum()
    ll += probs_nv.sum()

    for i in xrange(n_victims - 1, 0, -1):
        # iterate over all victim nodes to find the optimal threshold
        roots = n_roots + n_victims - 1 - i
        beta = 1. / (1. + exp(-probs[i]))
        if beta > float(roots) / age:
            break
    else:
        print "alpha: {0}, delta: {1}. Everyone is a root".format(alpha, delta)
        roots = n_victims + n_roots
        i = -1
    beta = float(roots) / age
    for i in xrange(n_victims - 1, 0, -1):
        if probs[i] >= log(beta/(1.- beta)):
            break
    ll += age * log(1 - beta)
    if i >= 0:
        ll += cums[i]
    if roots > 0:
        ll += roots * log(beta) - roots * log(1 - beta)
    return (beta, roots, ll)