# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03_utils.ipynb.

# %% auto 0
__all__ = ['trace', 'draw_dot', 'derivatives_matrix', 'fisher_info_matrix', 'cramer_rao_bound', 'CRLB_pipeline']

# %% ../nbs/03_utils.ipynb 5
from graphviz import Digraph

# %% ../nbs/03_utils.ipynb 8
def trace(root):
    nodes, edges = set(), set()
    def build(v):
        if v not in nodes:
            nodes.add(v)
            for child in v._prev:
                edges.add((child, v))
                build(child)
    build(root)
    return nodes, edges

def draw_dot(root, format='svg', rankdir='LR'):
    """
    format: png | svg | ...
    rankdir: TB (top to bottom graph) | LR (left to right)
    """
    assert rankdir in ['LR', 'TB']
    nodes, edges = trace(root)
    dot = Digraph(format=format, graph_attr={'rankdir': rankdir}) #, node_attr={'rankdir': 'TB'})
    
    for n in nodes:
        dot.node(name=str(id(n)), label = "{ data %.4f | grad %.4f }" % (n.data, n.grad), shape='record')
        if n._op:
            dot.node(name=str(id(n)) + n._op, label=n._op)
            dot.edge(str(id(n)) + n._op, str(id(n)))
    
    for n1, n2 in edges:
        dot.edge(str(id(n1)), str(id(n2)) + n2._op)
    
    return dot

# %% ../nbs/03_utils.ipynb 14
import numpy as np
import inspect

# %% ../nbs/03_utils.ipynb 15
def derivatives_matrix(func, ts: np.ndarray, params: dict, param_labels: list) -> np.ndarray:
    """calculate the derivatives matrix of a function w.r.t. its parameters at given measurements at ts"""
    # Initialize a dict to store grads for each parameter
    grads = {label: np.zeros(ts.size) for label in param_labels}
    
    for idx, t in enumerate(ts):
        # auto-diff
        y = func(t, *params.values())
        y.backward()
        
        # store grad for each parameter in the grads dictionary
        for label in param_labels:
            grads[label][idx] = params[label].grad
            params[label].zero_grad()
    
    # derivatives matrix
    Dij = np.zeros((ts.size, len(param_labels)))
    for i, label in enumerate(param_labels):
        Dij[:, i] = grads[label]
    
    return Dij

# %% ../nbs/03_utils.ipynb 16
def fisher_info_matrix(Dij: np.ndarray, sigma: float) -> np.ndarray:
    """calculate the Fisher Information Matrix (FIM) from the derivatives matrix Dij"""
    Iij = 1/(sigma**2) * Dij.T @ Dij
    return Iij

# %% ../nbs/03_utils.ipynb 17
def cramer_rao_bound(Iij: np.ndarray) -> np.ndarray:
    """calculate the Cramer-Rao Lower Bound (CRLB) from the Fisher Information Matrix (FIM)"""
    CRLB = np.linalg.inv(Iij)
    return CRLB

# %% ../nbs/03_utils.ipynb 18
def CRLB_pipeline(func, ts: np.ndarray, params: dict, param_labels: list, sigma:float, cond_num_cutoff=1e15) -> np.ndarray:
    """calculate the Cramer-Rao Lower Bound (CRLB) from a given function, measurements, parameters, and noise std"""
    Dij = derivatives_matrix(func, ts, params, param_labels)
    Iij = fisher_info_matrix(Dij, sigma)
    # condition number check
    if np.linalg.cond(Iij) < cond_num_cutoff:
        CRLB = cramer_rao_bound(Iij)
    else:
        print(f'Condition number of the Fisher Information Matrix is too high: {np.linalg.cond(Iij)}')
        CRLB = np.full(Iij.shape, np.nan)
    return CRLB
