#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# Author: Marcus Kasdorf
# Date:   July 28, 2021
"""
Functions to evolve nanowire networks over time.
"""

import numpy as np
import networkx as nx
from scipy.integrate import solve_ivp

from numbers import Number
from typing import Callable, Iterable, Optional
from scipy.integrate._ivp.ivp import OdeResult
import numpy.typing as npt

from .nanowire_network import NanowireNetwork
from .typing import *

from .nanowires import get_edge_indices
from .calculations import solve_drain_current, solve_network
from . import models


def solve_evolution(
    NWN: NanowireNetwork,
    t_eval: npt.NDArray,
    source_node: NWNNode | list[NWNNode],
    drain_node: NWNNode | list[NWNNode],
    voltage_func: Callable,
    window_func: Callable = None,
    tol: float = 1e-12,
    model: str = "default",
    solver: str = "spsolve",
    **kwargs,
) -> tuple[OdeResult, list[NWNEdge]]:
    """
    Deprecated. Please use NanowireNetwork.evolve() instead.

    Solve for the state variables `w` of the junctions of the given nanowire
    network at various points in time with an applied voltage.

    Parameters
    ----------
    NWN: Graph
        Nanowire network.

    t_eval : ndarray
        Time points to evaluate the nanowire network at. These should have
        units of `t0`.

    source_node : tuple, or list of tuples
        Voltage source nodes.

    drain_node : tuple, or list of tuples
        Grounded output nodes.

    voltage_func : Callable
        The applied voltage with the calling signature `func(t)`. The voltage
        should have units of `v0`.

    window_func : Callable, optional
        The window function used in the derivative of `w`. The calling
        signature is `f(w)` where w is the array of state variables.
        The default window function is `f(w) = 1`.

    tol : float, optional
        Tolerance of `scipy.integrate.solve_ivp`. Defaults to 1e-12.

    model : {"default", "decay", "chen"}, optional
        Evolutionary model type. Default: "default".

    solver : str, optional
        Name of sparse matrix solving algorithm to use. Default: "spsolve".

    **kwargs
        Keyword arguments passed to the solver.

    Returns
    -------
    sol : OdeResult
        Output from `scipy.integrate.solve_ivp`. See the SciPy documentation
        for information on this output's formatting.

    edge_list : list of tuples
        List of the edges corresponding with each `w`.

    """
    if model not in ("default", "HP", "decay", "chen", "SLT"):
        raise ValueError(f'Unsupported model type: model="{model}"')

    # Default window function
    if window_func is None:
        window_func = lambda x: 1

    # Get list of junction edges and state variable w
    edge_list, w0 = map(
        list,
        zip(*[((u, v), w) for u, v, w in NWN.edges.data("w") if w is not None]),
    )

    # Get edge indices
    start_nodes, end_nodes = get_edge_indices(NWN, edge_list)

    # Get list of tau
    tau0 = [tau for _, _, tau in NWN.edges.data("tau") if tau is not None]

    # Get list of epsilon
    epsilon0 = [ep for _, _, ep in NWN.edges.data("epsilon") if ep is not None]

    # Define initial state and associated derivative
    if model == "default" or model == "HP":
        _deriv = models.HP_model
        y0 = w0
    elif model == "decay":
        _deriv = models.decay_HP_model
        y0 = w0
    elif model == "chen" or model == "SLT":
        _deriv = models.SLT_HP_model
        y0 = np.hstack((w0, tau0, epsilon0))
        if "sigma" not in NWN.graph.keys():
            raise AttributeError(
                "sigma, theta, and a must be set before using the Chen model."
            )

    # Solve the system of ODEs
    t_span = (t_eval[0], t_eval[-1])
    sol = solve_ivp(
        _deriv,
        t_span,
        y0,
        "DOP853",
        t_eval,
        atol=tol,
        rtol=tol,
        args=(
            NWN,
            source_node,
            drain_node,
            voltage_func,
            window_func,
            solver,
            kwargs,
        ),
    )

    # Update NWN to final state variables.
    if model == "default":
        set_state_variables(NWN, sol.y[:, -1], edge_list)
    elif model == "decay":
        set_state_variables(NWN, sol.y[:, -1], edge_list)
    elif model == "chen":
        w_list, tau_list, eps_list = np.split(sol.y, 3)
        set_state_variables(
            NWN, w_list[:, -1], tau_list[:, -1], eps_list[:, -1], edge_list
        )

    return sol, edge_list


def set_state_variables(NWN: NanowireNetwork, *args):
    """
    Deprecated. Please use NanowireNetwork.set_state_var() instead.

    Sets the given nanowire network's state variable. Can be called in the
    following ways:

        set_state_variables(NWN, w)
            where `w` is a scalar value which is set for all edges.

        set_state_variables(NWN, w, edge_list)
            where `w` is an ndarray and edge_list is a list. The `w` array
            contains the state variable for the corresponding edge in
            `edge_list`.

        set_state_variables(NWN, w, tau)
            where `w` and `tau` are scalar values which is set for all edges.

        set_state_variables(NWN, w, tau, edge_list)
            where `w` and `tau` are ndarrays and edge_list is a list. The `w`
            and `tau` arrays contains the state variables for the
            corresponding edge in `edge_list`.

        set_state_variables(NWN, w, tau, epsilon)
            where `w`, `tau`, and `epsilon` are scalar values which is set for
            all edges.

        set_state_variables(NWN, w, tau, epsilon, edge_list)
            where `w`,`tau`, and `epsilon` are ndarrays and edge_list is a list.
            The `w`,`tau`, and `epsilon` arrays contains the state variables
            for the corresponding edge in `edge_list`.

    Parameters
    ----------
    NWN: Graph
        Nanowire network.

    *args
        See above.

    """
    # Only scalar w passed
    if len(args) == 1:
        if isinstance(args[0], Number):
            w = args[0]
            R = models.resist_func(NWN, w)

            attrs = {
                edge: {"w": w, "conductance": 1 / R}
                for edge in NWN.edges
                if NWN.edges[edge]["type"] == "junction"
            }
            nx.set_edge_attributes(NWN, attrs)

        else:
            raise ValueError("Invalid arguments.")

    elif len(args) == 2:
        # vector w and edge_list passed
        if isinstance(args[0], np.ndarray) and isinstance(args[1], Iterable):
            w, edge_list = args
            R = models.resist_func(NWN, w)

            attrs = {
                edge: {"w": w[i], "conductance": 1 / R[i]}
                for i, edge in enumerate(edge_list)
            }
            nx.set_edge_attributes(NWN, attrs)

        # scalar w and scalar tau passed
        elif isinstance(args[0], Number) and isinstance(args[1], Number):
            w, tau = args
            R = models.resist_func(NWN, w)

            attrs = {
                edge: {"w": w, "conductance": 1 / R, "tau": tau}
                for edge in NWN.edges
                if NWN.edges[edge]["type"] == "junction"
            }
            nx.set_edge_attributes(NWN, attrs)
            NWN.graph["tau"] = tau

        else:
            raise ValueError("Invalid arguments.")

    elif len(args) == 3:
        # vector w, vector tau, and edge_list passed
        if (
            isinstance(args[0], np.ndarray)
            and isinstance(args[1], np.ndarray)
            and isinstance(args[2], Iterable)
        ):
            w, tau, edge_list = args
            R = models.resist_func(NWN, w)

            attrs = {
                edge: {"w": w[i], "conductance": 1 / R[i], "tau": tau[i]}
                for i, edge in enumerate(edge_list)
            }
            nx.set_edge_attributes(NWN, attrs)

        # scalar w, scalar tau, scalar epsilon passed
        if (
            isinstance(args[0], Number)
            and isinstance(args[1], Number)
            and isinstance(args[2], Number)
        ):
            w, tau, epsilon = args
            R = models.resist_func(NWN, w)

            attrs = {
                edge: {
                    "w": w,
                    "conductance": 1 / R,
                    "tau": tau,
                    "epsilon": epsilon,
                }
                for edge in NWN.edges
                if NWN.edges[edge]["type"] == "junction"
            }
            nx.set_edge_attributes(NWN, attrs)
            NWN.graph["tau"] = tau
            NWN.graph["epsilon"] = epsilon

        else:
            raise ValueError("Invalid arguments.")

    elif len(args) == 4:
        # vector w, vector tau, vector epsilon, and edge_list passed
        if (
            isinstance(args[0], np.ndarray)
            and isinstance(args[1], np.ndarray)
            and isinstance(args[2], np.ndarray)
            and isinstance(args[3], Iterable)
        ):
            w, tau, epsilon, edge_list = args
            R = models.resist_func(NWN, w)

            attrs = {
                edge: {
                    "w": w[i],
                    "conductance": 1 / R[i],
                    "tau": tau[i],
                    "epsilon": epsilon[i],
                }
                for i, edge in enumerate(edge_list)
            }
            nx.set_edge_attributes(NWN, attrs)

        else:
            raise ValueError("Invalid arguments.")

    else:
        raise ValueError("Invalid number of arguments.")


def get_evolution_current(
    NWN: NanowireNetwork,
    sol: OdeResult,
    source_node: NWNNode | list[NWNNode],
    drain_node: NWNNode | list[NWNNode],
    voltage_func: Callable,
    state_vars: Optional[list[str]] = None,
    scaled: bool = False,
    solver: str = "spsolve",
    **kwargs,
) -> npt.NDArray:
    """
    Calculates the current draw from each drain node as a function of time,
    given the output state variables from `NWN.evolve()`.

    The source, drain, and voltage are expected to be the same as in the
    evolution of the state variables.

    This assumes the resistance function only depends on the first state variable
    set.

    Parameters
    ----------
    NWN : NanowireNetwork
        Nanowire network.

    sol : OdeResult
        Output state variables from `NWN.evolve()`.

    source_node : NWNNode, or list of NWNNodes
        Voltage source nodes.

    drain_node : NWNNode, or list of NWNNodes
        Grounded output nodes.

    voltage_func : Callable
        Function which inputs the time as a scalar and returns the voltage of
        all the source nodes as a scalar. Should be dimensionless.

    state_vars : list of str, optional
        List of state variables to use. Should not need to be provided.
        Defaults to `NWN.state_vars`.

    scaled : bool, optional
        Scale the output by the characteristic values. Default: False.

    solver : str, optional
        Name of sparse matrix solving algorithm to use. Default: "spsolve".

    **kwargs
        Keyword arguments passed to the solver.

    Returns
    -------
    current_array: ndarray
        Array containing the current flow through each drain node. Each column
        corresponds to a drain node in the order passed.

    """
    # Get lists of source and drain nodes
    if isinstance(source_node, tuple):
        source_node = [source_node]
    if isinstance(drain_node, tuple):
        drain_node = [drain_node]

    if state_vars is None:
        state_vars = NWN.state_vars

    # Preallocate output
    current_array = np.zeros((len(sol.t), len(drain_node)))

    # Loop through each time step
    for i, state in enumerate(sol.y.T):

        # Set state variables
        split = np.split(state, len(state_vars))
        for j, var in enumerate(state_vars):
            NWN.set_state_var(var, split[j])

        # Update resistance
        NWN.update_resistance(split[0])

        # Find current draw through drain nodes
        V = voltage_func(sol.t[i])
        current_array[i] = solve_drain_current(
            NWN, source_node, drain_node, V, scaled, solver, **kwargs
        )

    return current_array.squeeze()


def get_evolution_node_voltages(
    NWN: nx.Graph,
    sol: OdeResult,
    edge_list: list[NWNEdge],
    source_node: NWNNode | list[NWNNode],
    drain_node: NWNNode | list[NWNNode],
    voltage_func: Callable,
    solver: str = "spsolve",
    **kwargs,
) -> npt.NDArray:
    """
    To be used in conjunction with `solve_evolution`. Takes the output from
    `solve_evolution` and finds the voltage of all nodes in the network at each
    time step.

    The appropriate parameters passed should be the same as `solve_evolution`.

    Parameters
    ----------
    NWN : Graph
        Nanowire network.

    sol : OdeResult
        Output from `solve_evolution`.

    edge_list : list of tuples
        Output from `solve_evolution`.

    source_node : tuple, or list of tuples
        Voltage source nodes.

    drain_node : tuple, or list of tuples
        Grounded output nodes.

    voltage_func : Callable
        The applied voltage with the calling signature `func(t)`. The voltage
        should have units of `v0`.

    solver : str, optional
        Name of sparse matrix solving algorithm to use. Default: "spsolve".

    **kwargs
        Keyword arguments passed to the solver.

    Returns
    -------
    node_voltage_array: ndarray
        An `n` x `m` array containing the voltage of each node in the network,
        corresponding to the `n` time steps and `m` nodes in the network.

    """
    # Get lists of source and drain nodes
    if isinstance(source_node, tuple):
        source_node = [source_node]
    if isinstance(drain_node, tuple):
        drain_node = [drain_node]

    # Preallocate output
    node_voltage_array = np.zeros((len(sol.t), len(NWN.nodes)))

    # Loop through each time step
    for i in range(len(sol.t)):
        # Set state variables and get drain currents
        input_V = voltage_func(sol.t[i])
        set_state_variables(NWN, sol.y.T[i], edge_list)
        *node_voltage_array[i], I = solve_network(
            NWN, source_node, drain_node, input_V, "voltage", solver, **kwargs
        )

    return node_voltage_array.squeeze()
