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

DTYPE = np.float64
ctypedef np.float_t DTYPE_t

cdef DTYPE_t weight_victim(int dist, int dt, DTYPE_t alpha,
                           DTYPE_t delta, DTYPE_t gamma):
    cdef DTYPE_t structural, temporal
    structural = delta ** dist
    temporal = (gamma - 1. / alpha) * 1. / (1. + dt / alpha) ** gamma
    return structural * temporal


cdef DTYPE_t weight_non_victim(int dist, int t, DTYPE_t alpha,
                               DTYPE_t delta, DTYPE_t gamma):
    cdef DTYPE_t structural, temporal
    structural = delta ** dist
    temporal = 1. - 1. / (1. + (3012. - t) / alpha) ** gamma
    return 1. - structural * temporal


def ml(dict root_victims, dict victims, dict non_victims,
       DTYPE_t alpha, DTYPE_t delta, DTYPE_t gamma=1.01):
    cdef:
        int n_roots, n_victims, n_nodes, roots, i, dist, dt, t
        DTYPE_t beta
        list parents, parents_weights
    n_roots, n_victims = len(root_victims), len(victims)
    n_nodes = n_victims + len(non_victims)
    cdef:
        np.ndarray[DTYPE_t] probs = 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()):
        parents_weights = [weight_victim(dist, dt, alpha, delta, gamma)
                           for (dist, dt) in parents]
        probs[i] = max(parents_weights)
    for i, parents in enumerate(non_victims.itervalues()):
        parents_weights = [weight_non_victim(dist, t, alpha, delta, gamma)
                           for (dist, t) in parents]
        probs_nv[i] = max(parents_weights)
    probs.sort()
    probs = probs[::-1]
    cdef:
        np.ndarray[DTYPE_t] betas = probs / (1. + probs)
        np.ndarray[DTYPE_t] cums = np.log(probs.cumsum())
    for i in xrange(n_victims - 1, 0, -1):
        roots = n_roots + n_victims - 1 - i
        if betas[i] > roots / float(n_nodes):
            break
    else:
        print "alpha: {0}, delta: {1}. Everyone is a root".format(alpha, delta)
        roots = n_roots + n_victims
    beta = roots / float(n_nodes)
    return (beta, roots,
            roots * log(beta) + (n_nodes - roots) * log(1 - beta) + cums[i]
            + np.log(probs_nv).sum())