# 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 lmbda,
                            DTYPE_t w1, DTYPE_t w2, DTYPE_t w3):
    """weight for successful infection, exponential time model"""
    cdef DTYPE_t structural, temporal, result
    structural = delta ** dist
    # structural = delta/(1. + 1./(w1*lmbda) + 1./(w2*lmbda) + 1./(w3*lmbda))
    temporal = log(exp(alpha)-1.) - alpha*dt
    # 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 lmbda,
                            DTYPE_t w1, DTYPE_t w2, DTYPE_t w3):
    """weight for failed infection, exponential time model"""
    cdef DTYPE_t structural, temporal, result
    structural = delta ** dist
    # structural = delta/(1. + 1./(w1*lmbda) + 1./(w2*lmbda) + 1./(w3*lmbda))
    temporal = exp(-alpha * dt)
    # temporal = 1.  - 1. / (1. + dt/alpha)**0.01
    result = log(1. - structural + structural * temporal)
    return result

def ml2(dict root_victims, dict victims, dict non_victims,
       DTYPE_t alpha, DTYPE_t delta, DTYPE_t lmbda):
    cdef:
        int n_roots, n_victims, roots, i, dist, dt, t, l
        DTYPE_t ll
        list parents, failures, successes
    n_roots, n_victims = len(root_victims), len(victims)
    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)
        np.ndarray[DTYPE_t] parent_dists = np.zeros(n_victims, dtype=DTYPE)
        np.ndarray[DTYPE_t] parent_dts = np.zeros(n_victims, dtype=DTYPE)
        np.ndarray[DTYPE_t] infectors = np.zeros(n_victims, dtype=DTYPE)

    # loop through victims
    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, lmbda, w1, w2, w3)
                    for (prnt, dist, dt, w1, w2, w3) in parents]
        probs_fail[i] = sum(failures)
        successes = [weight_success(dist, dt, alpha, delta, lmbda, w1, w2, w3)
                     for (prnt, dist, dt, w1, w2, w3) in parents]
        dists = [dist for (prnt, dist, dt, w1, w2, w3) in parents]
        dts = [dt for (prnt, dist, dt, w1, w2, w3) in parents]
        prnts = [prnt for (prnt, dist, dt, w1, w2, w3) in parents]
        # find parent that maximizes log(p) - log(\tilde{p})
        # probs[i] = max(s - failures[l] for l, s in enumerate(successes)) 
        probs[i] = float("-inf")
        for l, s in enumerate(successes):
            prob = s - failures[l]
            if prob > probs[i]:
                probs[i] = prob
                infectors[i] = prnts[l]
                parent_dists[i] = dists[l]
                parent_dts[i] = dts[l]

    # loop through non-victims
    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, lmbda, w1, w2, w3)
                    for (prnt, dist, dt, w1, w2, w3) in parents]
        probs_nv[i] = sum(failures)

    # calculate log likelihood
    ll =  probs_fail.sum() # add probability that all edges to victims fail
    ll += probs_nv.sum() # add probability that all edges to non_victims fail
    ll += probs.sum() # add probability for realized edges and subtract probability these edges fail

    roots = n_roots
    # print n_nodes, n_roots, n_victims, max_i, roots
    # print parent_dists[1:100]
    # print parent_dts[1:100]
    # print victims.keys()
    # print infectors
    # print np.mean(parent_dists)
    # print np.mean(parent_dts)
    
    with open('../../Results/infectors.csv', 'w') as infectors_file:
        for i, infector in enumerate(infectors):
            infectors_file.write("%s, %s\n" % ((victims.keys())[i], infector))
    return (lmbda, roots, ll)